diff --git a/roop/core.py b/roop/core.py index 30507e2..307e761 100755 --- a/roop/core.py +++ b/roop/core.py @@ -36,25 +36,27 @@ parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_ parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=4) parser.add_argument('--gpu-vendor', help='choice your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia']) -args = {} - -for name, value in vars(parser.parse_args()).items(): - args[name] = value +args = parser.parse_known_args()[0] if 'all_faces' in args: roop.globals.all_faces = True -if args['cpu_cores']: - roop.globals.cpu_cores = args['cpu_cores'] +if args.cpu_cores: + roop.globals.cpu_cores = int(args.cpu_cores) +# cpu thread fix for mac if sys.platform == 'darwin': roop.globals.cpu_cores = 1 -if args['gpu_threads']: - roop.globals.gpu_threads = args['gpu_threads'] +if args.gpu_threads: + roop.globals.gpu_threads = int(args.gpu_threads) -if args['gpu_vendor']: - roop.globals.gpu_vendor = args['gpu_vendor'] +# gpu thread fix for amd +if args.gpu_vendor == 'amd': + roop.globals.gpu_threads = 1 + +if args.gpu_vendor: + roop.globals.gpu_vendor = args.gpu_vendor else: roop.globals.providers = ['CPUExecutionProvider'] @@ -62,12 +64,10 @@ sep = "/" if os.name == "nt": sep = "\\" -POOL = None - def limit_resources(): - if args['max_memory']: - memory = args['max_memory'] * 1024 * 1024 * 1024 + if args.max_memory: + memory = args.max_memory * 1024 * 1024 * 1024 if str(platform.system()).lower() == 'windows': import ctypes kernel32 = ctypes.windll.kernel32 @@ -156,27 +156,27 @@ def process_video_multi_cores(source_img, frame_paths): def start(preview_callback = None): - if not args['source_img'] or not os.path.isfile(args['source_img']): + if not args.source_img or not os.path.isfile(args.source_img): print("\n[WARNING] Please select an image containing a face.") 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.") return - if not args['output_file']: - target_path = args['target_path'] - args['output_file'] = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path - target_path = args['target_path'] - test_face = get_face_single(cv2.imread(args['source_img'])) + if not args.output_file: + target_path = args.target_path + args.output_file = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path + target_path = args.target_path + test_face = get_face_single(cv2.imread(args.source_img)) if not test_face: print("\n[WARNING] No face detected in source image. Please try with another one.\n") return if is_img(target_path): if predict_image(target_path) > 0.85: quit() - process_img(args['source_img'], target_path, args['output_file']) + process_img(args.source_img, target_path, args.output_file) status("swap successful!") 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) if any(probability > 0.85 for probability in probabilities): quit() video_name_full = target_path.split("/")[-1] @@ -185,7 +185,7 @@ def start(preview_callback = None): Path(output_dir).mkdir(exist_ok=True) status("detecting video's FPS...") fps, exact_fps = detect_fps(target_path) - if not args['keep_fps'] and fps > 30: + if not args.keep_fps and fps > 30: this_path = output_dir + "/" + video_name + ".mp4" set_fps(target_path, this_path, 30) target_path, exact_fps = this_path, 30 @@ -193,33 +193,33 @@ def start(preview_callback = None): shutil.copy(target_path, output_dir) status("extracting frames...") extract_frames(target_path, output_dir) - args['frame_paths'] = tuple(sorted( + args.frame_paths = tuple(sorted( glob.glob(output_dir + "/*.png"), key=lambda x: int(x.split(sep)[-1].replace(".png", "")) )) status("swapping in progress...") - if roop.globals.cpu_cores > 0 and roop.globals.gpu_vendor is None: + if roop.globals.gpu_vendor is None and roop.globals.cpu_cores > 0: global POOL POOL = mp.Pool(roop.globals.cpu_cores) - process_video_multi_cores(args['source_img'], args['frame_paths']) + process_video_multi_cores(args.source_img, args.frame_paths) else: - process_video(args['source_img'], args["frame_paths"]) + process_video(args.source_img, args.frame_paths) status("creating video...") create_video(video_name, exact_fps, output_dir) status("adding audio...") - add_audio(output_dir, target_path, video_name_full, args['keep_frames'], args['output_file']) - save_path = args['output_file'] if args['output_file'] else output_dir + "/" + video_name + ".mp4" + add_audio(output_dir, target_path, video_name_full, args.keep_frames, args.output_file) + save_path = args.output_file if args.output_file else output_dir + "/" + video_name + ".mp4" print("\n\nVideo saved as:", save_path, "\n\n") status("swap successful!") def select_face_handler(path: str): - args['source_img'] = path + args.source_img = path def select_target_handler(path: str): - args['target_path'] = path - return preview_video(args['target_path']) + args.target_path = path + return preview_video(args.target_path) def toggle_all_faces_handler(value: int): @@ -227,21 +227,21 @@ def toggle_all_faces_handler(value: int): def toggle_fps_limit_handler(value: int): - args['keep_fps'] = int(value != 1) + args.keep_fps = int(value != 1) def toggle_keep_frames_handler(value: int): - args['keep_frames'] = value + args.keep_frames = value def save_file_handler(path: str): - args['output_file'] = path + args.output_file = path def create_test_preview(frame_number): return process_faces( - get_face_single(cv2.imread(args['source_img'])), - get_video_frame(args['target_path'], frame_number) + get_face_single(cv2.imread(args.source_img)), + get_video_frame(args.target_path, frame_number) ) @@ -250,16 +250,16 @@ def run(): pre_check() limit_resources() - if args['source_img']: - args['cli_mode'] = True + if args.source_img: + args.cli_mode = True start() quit() window = ui.init( { 'all_faces': roop.globals.all_faces, - 'keep_fps': args['keep_fps'], - 'keep_frames': args['keep_frames'] + 'keep_fps': args.keep_fps, + 'keep_frames': args.keep_frames }, select_face_handler, select_target_handler,