upagge/roop
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Normalize ARGS

Browse Source
This commit is contained in:
henryruhs 2023-06-05 14:45:33 +02:00
parent b116f2001a
commit 4420ba5bdd
2 changed files with 50 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
roop/core.py
View File

@ -24,11 +24,14 @@ from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, ext
from roop.analyser import get_face_single from roop.analyser import get_face_single
import roop.ui as ui import roop.ui as ui
def handle_parse():
global args
signal.signal(signal.SIGINT, lambda signal_number, frame: quit()) signal.signal(signal.SIGINT, lambda signal_number, frame: quit())
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('-f', '--face', help='use this face', dest='source_img') parser.add_argument('-f', '--face', help='use this face', dest='source_target')
parser.add_argument('-t', '--target', help='replace this face', dest='target_path') parser.add_argument('-t', '--target', help='replace this face', dest='target_path')
parser.add_argument('-o', '--output', help='save output to this file', dest='output_file') parser.add_argument('-o', '--output', help='save output to this file', dest='output_path')
parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False)
parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False)
parser.add_argument('--all-faces', help='swap all faces in frame', dest='all_faces', action='store_true', default=False) parser.add_argument('--all-faces', help='swap all faces in frame', dest='all_faces', action='store_true', default=False)
@ -39,8 +42,8 @@ parser.add_argument('--gpu-vendor', help='choice your GPU vendor', dest='gpu_ven
args = parser.parse_known_args()[0] args = parser.parse_known_args()[0]
if 'all_faces' in args: roop.globals.headless = args.source_target or args.target_path or args.output_path
roop.globals.all_faces = True roop.globals.all_faces = args.all_faces
if args.cpu_cores: if args.cpu_cores:
roop.globals.cpu_cores = int(args.cpu_cores) roop.globals.cpu_cores = int(args.cpu_cores)
@ -61,10 +64,6 @@ if args.gpu_vendor:
else: else:
roop.globals.providers = ['CPUExecutionProvider'] roop.globals.providers = ['CPUExecutionProvider']
sep = "/"
if os.name == "nt":
sep = "\\"
def limit_resources(): def limit_resources():
# prevent tensorflow memory leak # prevent tensorflow memory leak
@ -141,18 +140,18 @@ def preview_video(video_path):
def status(string): def status(string):
value = "Status: " + string value = "Status: " + string
if 'cli_mode' in args: if roop.globals.headless:
print(value) print(value)
else: else:
ui.update_status_label(value) ui.update_status_label(value)
def process_video_multi_cores(source_img, frame_paths): def process_video_multi_cores(source_target, frame_paths):
n = len(frame_paths) // roop.globals.cpu_cores n = len(frame_paths) // roop.globals.cpu_cores
if n > 2: if n > 2:
processes = [] processes = []
for i in range(0, len(frame_paths), n): for i in range(0, len(frame_paths), n):
p = POOL.apply_async(process_video, args=(source_img, frame_paths[i:i + n],)) p = POOL.apply_async(process_video, args=(source_target, frame_paths[i:i + n],))
processes.append(p) processes.append(p)
for p in processes: for p in processes:
p.get() p.get()
@ -161,24 +160,24 @@ def process_video_multi_cores(source_img, frame_paths):
def start(preview_callback = None): def start(preview_callback = None):
if not args.source_img or not os.path.isfile(args.source_img): if not args.source_target or not os.path.isfile(args.source_target):
print("\n[WARNING] Please select an image containing a face.") print("\n[WARNING] Please select an image containing a face.")
return return
elif not args.target_path or not os.path.isfile(args.target_path): elif not args.target_path or not os.path.isfile(args.target_path):
print("\n[WARNING] Please select a video/image to swap face in.") print("\n[WARNING] Please select a video/image to swap face in.")
return return
if not args.output_file: if not args.output_path:
target_path = args.target_path target_path = args.target_path
args.output_file = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path args.output_path = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path
target_path = args.target_path target_path = args.target_path
test_face = get_face_single(cv2.imread(args.source_img)) test_face = get_face_single(cv2.imread(args.source_target))
if not test_face: if not test_face:
print("\n[WARNING] No face detected in source image. Please try with another one.\n") print("\n[WARNING] No face detected in source image. Please try with another one.\n")
return return
if is_img(target_path): if is_img(target_path):
if predict_image(target_path) > 0.85: if predict_image(target_path) > 0.85:
quit() quit()
process_img(args.source_img, target_path, args.output_file) process_img(args.source_target, target_path, args.output_path)
status("swap successful!") status("swap successful!")
return return
seconds, probabilities = predict_video_frames(video_path=args.target_path, frame_interval=100) seconds, probabilities = predict_video_frames(video_path=args.target_path, frame_interval=100)
@ -200,29 +199,29 @@ def start(preview_callback = None):
extract_frames(target_path, output_dir) extract_frames(target_path, output_dir)
args.frame_paths = tuple(sorted( args.frame_paths = tuple(sorted(
glob.glob(output_dir + "/*.png"), glob.glob(output_dir + "/*.png"),
key=lambda x: int(x.split(sep)[-1].replace(".png", "")) key=lambda x: int(x.split(os.sep)[-1].replace(".png", ""))
)) ))
status("swapping in progress...") status("swapping in progress...")
if roop.globals.gpu_vendor is None and roop.globals.cpu_cores > 1: if roop.globals.gpu_vendor is None and roop.globals.cpu_cores > 1:
global POOL global POOL
POOL = mp.Pool(roop.globals.cpu_cores) POOL = mp.Pool(roop.globals.cpu_cores)
process_video_multi_cores(args.source_img, args.frame_paths) process_video_multi_cores(args.source_target, args.frame_paths)
else: else:
process_video(args.source_img, args.frame_paths) process_video(args.source_target, args.frame_paths)
# prevent out of memory while using ffmpeg with cuda # prevent out of memory while using ffmpeg with cuda
if args.gpu_vendor == 'nvidia': if args.gpu_vendor == 'nvidia':
torch.cuda.empty_cache() torch.cuda.empty_cache()
status("creating video...") status("creating video...")
create_video(video_name, exact_fps, output_dir) create_video(video_name, exact_fps, output_dir)
status("adding audio...") status("adding audio...")
add_audio(output_dir, target_path, video_name_full, args.keep_frames, args.output_file) add_audio(output_dir, target_path, video_name_full, args.keep_frames, args.output_path)
save_path = args.output_file if args.output_file else output_dir + "/" + video_name + ".mp4" save_path = args.output_path if args.output_path else output_dir + "/" + video_name + ".mp4"
print("\n\nVideo saved as:", save_path, "\n\n") print("\n\nVideo saved as:", save_path, "\n\n")
status("swap successful!") status("swap successful!")
def select_face_handler(path: str): def select_face_handler(path: str):
args.source_img = path args.source_target = path
def select_target_handler(path: str): def select_target_handler(path: str):
@ -243,26 +242,24 @@ def toggle_keep_frames_handler(value: int):
def save_file_handler(path: str): def save_file_handler(path: str):
args.output_file = path args.output_path = path
def create_test_preview(frame_number): def create_test_preview(frame_number):
return process_faces( return process_faces(
get_face_single(cv2.imread(args.source_img)), get_face_single(cv2.imread(args.source_target)),
get_video_frame(args.target_path, frame_number) get_video_frame(args.target_path, frame_number)
) )
def run(): def run():
global all_faces, keep_frames, limit_fps global all_faces, keep_frames, limit_fps
handle_parse()
pre_check() pre_check()
limit_resources() limit_resources()
if args.source_img: if roop.globals.headless:
args.cli_mode = True
start() start()
quit() quit()
window = ui.init( window = ui.init(
{ {
'all_faces': roop.globals.all_faces, 'all_faces': roop.globals.all_faces,
@ -279,5 +276,4 @@ def run():
get_video_frame, get_video_frame,
create_test_preview create_test_preview
) )
window.mainloop() window.mainloop()

1
roop/globals.py
View File

@ -5,6 +5,7 @@ log_level = 'error'
cpu_cores = None cpu_cores = None
gpu_threads = None gpu_threads = None
gpu_vendor = None gpu_vendor = None
headless = None
providers = onnxruntime.get_available_providers() providers = onnxruntime.get_available_providers()
if 'TensorrtExecutionProvider' in providers: if 'TensorrtExecutionProvider' in providers: