From 6daee960292dfc55ab17db932b11923c09847ce2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Phan=20Tu=E1=BA=A5n=20Anh?= Date: Mon, 5 Jun 2023 11:57:11 +0200 Subject: [PATCH 01/40] ffmpeg platform-agnostic hardware-acceleration --- roop/utils.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/roop/utils.py b/roop/utils.py index 3ec6872..450a7ed 100644 --- a/roop/utils.py +++ b/roop/utils.py @@ -31,8 +31,7 @@ def detect_fps(input_path): def run_ffmpeg(args): - log_level = f'-loglevel {roop.globals.log_level}' - run_command(f'ffmpeg {log_level} {args}') + run_command(f'ffmpeg -hide_banner -hwaccel auto -loglevel {roop.globals.log_level} {args}') def set_fps(input_path, output_path, fps): @@ -41,15 +40,13 @@ def set_fps(input_path, output_path, fps): def create_video(video_name, fps, output_dir): - hwaccel_option = '-hwaccel cuda' if roop.globals.gpu_vendor == 'nvidia' else '' output_dir = path(output_dir) - run_ffmpeg(f'{hwaccel_option} -framerate "{fps}" -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4"') + run_ffmpeg(f'-framerate "{fps}" -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4"') def extract_frames(input_path, output_dir): - hwaccel_option = '-hwaccel cuda' if roop.globals.gpu_vendor == 'nvidia' else '' input_path, output_dir = path(input_path), path(output_dir) - run_ffmpeg(f' {hwaccel_option} -i "{input_path}" "{output_dir}{sep}%04d.png"') + run_ffmpeg(f'-i "{input_path}" "{output_dir}{sep}%04d.png"') def add_audio(output_dir, target_path, video, keep_frames, output_file): From c59400f6a6aa5a21836f8245b5e148cd9421da2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Phan=20Tu=E1=BA=A5n=20Anh?= Date: Mon, 5 Jun 2023 11:58:24 +0200 Subject: [PATCH 02/40] clear CUDA cache after swapping on low VRAM + ffmpeg cuda acceleration, clearing cache prevent cuda out-of-memory error --- roop/core.py | 1 + 1 file changed, 1 insertion(+) diff --git a/roop/core.py b/roop/core.py index c14bc48..b667a97 100755 --- a/roop/core.py +++ b/roop/core.py @@ -209,6 +209,7 @@ def start(preview_callback = None): process_video_multi_cores(args.source_img, args.frame_paths) else: process_video(args.source_img, args.frame_paths) + torch.cuda.empty_cache() status("creating video...") create_video(video_name, exact_fps, output_dir) status("adding audio...") From 36683605207bc7de28ee27c634dcb641aa4bc044 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Phan=20Tu=E1=BA=A5n=20Anh?= Date: Mon, 5 Jun 2023 12:21:37 +0200 Subject: [PATCH 03/40] check torch gpu before clearing cache --- roop/core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/roop/core.py b/roop/core.py index b667a97..22113c2 100755 --- a/roop/core.py +++ b/roop/core.py @@ -209,7 +209,8 @@ def start(preview_callback = None): process_video_multi_cores(args.source_img, args.frame_paths) else: process_video(args.source_img, args.frame_paths) - torch.cuda.empty_cache() + if torch.cuda.is_available() and args.gpu_vendor in ['amd', 'nvidia']: + torch.cuda.empty_cache() status("creating video...") create_video(video_name, exact_fps, output_dir) status("adding audio...") From b9476034416850ca69258e7acaa7c257af944ac6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Phan=20Tu=E1=BA=A5n=20Anh?= Date: Mon, 5 Jun 2023 12:25:35 +0200 Subject: [PATCH 04/40] torch check nvidia only --- roop/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/roop/core.py b/roop/core.py index 22113c2..f392a2a 100755 --- a/roop/core.py +++ b/roop/core.py @@ -209,7 +209,7 @@ def start(preview_callback = None): process_video_multi_cores(args.source_img, args.frame_paths) else: process_video(args.source_img, args.frame_paths) - if torch.cuda.is_available() and args.gpu_vendor in ['amd', 'nvidia']: + if args.gpu_vendor = 'nvidia': torch.cuda.empty_cache() status("creating video...") create_video(video_name, exact_fps, output_dir) From 71af0fac69deeec3612ceb9cabaa04cbd1e381d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Phan=20Tu=E1=BA=A5n=20Anh?= Date: Mon, 5 Jun 2023 12:29:30 +0200 Subject: [PATCH 05/40] syntax error --- roop/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/roop/core.py b/roop/core.py index f392a2a..7a48f96 100755 --- a/roop/core.py +++ b/roop/core.py @@ -209,8 +209,8 @@ def start(preview_callback = None): process_video_multi_cores(args.source_img, args.frame_paths) else: process_video(args.source_img, args.frame_paths) - if args.gpu_vendor = 'nvidia': - torch.cuda.empty_cache() + if args.gpu_vendor == 'nvidia': + torch.cuda.empty_cache() # prevent CUDA OOM when using ffmpeg cuda accel status("creating video...") create_video(video_name, exact_fps, output_dir) status("adding audio...") From b116f2001a147342fd4ca7f8c26afda263c121eb Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 12:42:24 +0200 Subject: [PATCH 06/40] Adjust comment --- roop/core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/roop/core.py b/roop/core.py index db8d48d..b187fc7 100755 --- a/roop/core.py +++ b/roop/core.py @@ -209,8 +209,9 @@ def start(preview_callback = None): process_video_multi_cores(args.source_img, args.frame_paths) else: process_video(args.source_img, args.frame_paths) + # prevent out of memory while using ffmpeg with cuda if args.gpu_vendor == 'nvidia': - torch.cuda.empty_cache() # prevent CUDA OOM when using ffmpeg cuda accel + torch.cuda.empty_cache() status("creating video...") create_video(video_name, exact_fps, output_dir) status("adding audio...") From 4420ba5bdd4479fee29dccce87387b6057b2f393 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 14:45:33 +0200 Subject: [PATCH 07/40] Normalize ARGS --- roop/core.py | 102 +++++++++++++++++++++++------------------------- roop/globals.py | 1 + 2 files changed, 50 insertions(+), 53 deletions(-) diff --git a/roop/core.py b/roop/core.py index b187fc7..e548452 100755 --- a/roop/core.py +++ b/roop/core.py @@ -24,46 +24,45 @@ from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, ext from roop.analyser import get_face_single import roop.ui as ui -signal.signal(signal.SIGINT, lambda signal_number, frame: quit()) -parser = argparse.ArgumentParser() -parser.add_argument('-f', '--face', help='use this face', dest='source_img') -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('--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('--all-faces', help='swap all faces in frame', dest='all_faces', action='store_true', default=False) -parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) -parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) -parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) -parser.add_argument('--gpu-vendor', help='choice your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia']) -args = parser.parse_known_args()[0] +def handle_parse(): + global args + signal.signal(signal.SIGINT, lambda signal_number, frame: quit()) + parser = argparse.ArgumentParser() + 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('-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-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('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) + parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) + parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) + parser.add_argument('--gpu-vendor', help='choice your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia']) -if 'all_faces' in args: - roop.globals.all_faces = True + args = parser.parse_known_args()[0] -if args.cpu_cores: - roop.globals.cpu_cores = int(args.cpu_cores) + roop.globals.headless = args.source_target or args.target_path or args.output_path + roop.globals.all_faces = args.all_faces -# cpu thread fix for mac -if sys.platform == 'darwin': - roop.globals.cpu_cores = 1 + if args.cpu_cores: + roop.globals.cpu_cores = int(args.cpu_cores) -if args.gpu_threads: - roop.globals.gpu_threads = int(args.gpu_threads) + # cpu thread fix for mac + if sys.platform == 'darwin': + roop.globals.cpu_cores = 1 -# gpu thread fix for amd -if args.gpu_vendor == 'amd': - roop.globals.gpu_threads = 1 + if args.gpu_threads: + roop.globals.gpu_threads = int(args.gpu_threads) -if args.gpu_vendor: - roop.globals.gpu_vendor = args.gpu_vendor -else: - roop.globals.providers = ['CPUExecutionProvider'] + # gpu thread fix for amd + if args.gpu_vendor == 'amd': + roop.globals.gpu_threads = 1 -sep = "/" -if os.name == "nt": - sep = "\\" + if args.gpu_vendor: + roop.globals.gpu_vendor = args.gpu_vendor + else: + roop.globals.providers = ['CPUExecutionProvider'] def limit_resources(): @@ -141,18 +140,18 @@ def preview_video(video_path): def status(string): value = "Status: " + string - if 'cli_mode' in args: + if roop.globals.headless: print(value) else: 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 if n > 2: processes = [] 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) for p in processes: p.get() @@ -161,24 +160,24 @@ 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_target or not os.path.isfile(args.source_target): print("\n[WARNING] Please select an image containing a face.") return 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: + if not args.output_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 - test_face = get_face_single(cv2.imread(args.source_img)) + test_face = get_face_single(cv2.imread(args.source_target)) 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_target, target_path, args.output_path) status("swap successful!") return 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) args.frame_paths = tuple(sorted( 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...") if roop.globals.gpu_vendor is None and roop.globals.cpu_cores > 1: 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_target, args.frame_paths) 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 if args.gpu_vendor == 'nvidia': torch.cuda.empty_cache() 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_path) + save_path = args.output_path if args.output_path 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_target = path def select_target_handler(path: str): @@ -243,26 +242,24 @@ def toggle_keep_frames_handler(value: int): def save_file_handler(path: str): - args.output_file = path + args.output_path = path def create_test_preview(frame_number): 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) ) def run(): global all_faces, keep_frames, limit_fps - + handle_parse() pre_check() limit_resources() - if args.source_img: - args.cli_mode = True + if roop.globals.headless: start() quit() - window = ui.init( { 'all_faces': roop.globals.all_faces, @@ -279,5 +276,4 @@ def run(): get_video_frame, create_test_preview ) - window.mainloop() diff --git a/roop/globals.py b/roop/globals.py index 986bf91..9162a39 100644 --- a/roop/globals.py +++ b/roop/globals.py @@ -5,6 +5,7 @@ log_level = 'error' cpu_cores = None gpu_threads = None gpu_vendor = None +headless = None providers = onnxruntime.get_available_providers() if 'TensorrtExecutionProvider' in providers: From 249e6b3923995dcfd57a0130e4c896e088a10286 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:01:26 +0200 Subject: [PATCH 08/40] Remove path normalization --- roop/utils.py | 21 +++------------------ 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/roop/utils.py b/roop/utils.py index 450a7ed..eba456b 100644 --- a/roop/utils.py +++ b/roop/utils.py @@ -2,16 +2,6 @@ import os import shutil import roop.globals -sep = "/" -if os.name == "nt": - sep = "\\" - - -def path(string): - if sep == "\\": - return string.replace("/", "\\") - return string - def run_command(command, mode="silent"): if mode == "debug": @@ -20,7 +10,6 @@ def run_command(command, mode="silent"): def detect_fps(input_path): - input_path = path(input_path) output = os.popen(f'ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate "{input_path}"').read() if "/" in output: try: @@ -35,25 +24,21 @@ def run_ffmpeg(args): def set_fps(input_path, output_path, fps): - input_path, output_path = path(input_path), path(output_path) run_ffmpeg(f'-i "{input_path}" -filter:v fps=fps={fps} "{output_path}"') def create_video(video_name, fps, output_dir): - output_dir = path(output_dir) - run_ffmpeg(f'-framerate "{fps}" -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4"') + run_ffmpeg(f'-framerate "{fps}" -i "{output_dir}{os.sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{os.sep}output.mp4"') def extract_frames(input_path, output_dir): - input_path, output_dir = path(input_path), path(output_dir) - run_ffmpeg(f'-i "{input_path}" "{output_dir}{sep}%04d.png"') + run_ffmpeg(f'-i "{input_path}" "{output_dir}{os.sep}%04d.png"') def add_audio(output_dir, target_path, video, keep_frames, output_file): video_name = os.path.splitext(video)[0] save_to = output_file if output_file else output_dir + "/swapped-" + video_name + ".mp4" - save_to_ff, output_dir_ff = path(save_to), path(output_dir) - run_ffmpeg(f'-i "{output_dir_ff}{sep}output.mp4" -i "{output_dir_ff}{sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to_ff}"') + run_ffmpeg(f'-i "{output_dir}{os.sep}output.mp4" -i "{output_dir}{os.sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to}"') if not os.path.isfile(save_to): shutil.move(output_dir + "/output.mp4", save_to) if not keep_frames: From d8c6581900c3b7f7f69e5ca7fbb9ebe69f224a65 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:12:22 +0200 Subject: [PATCH 09/40] Remove args overrides --- roop/core.py | 13 +++++-------- roop/utils.py | 5 ----- 2 files changed, 5 insertions(+), 13 deletions(-) diff --git a/roop/core.py b/roop/core.py index e548452..f93785b 100755 --- a/roop/core.py +++ b/roop/core.py @@ -20,7 +20,7 @@ import cv2 import roop.globals from roop.swapper import process_video, process_img, process_faces, process_frames -from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames, rreplace +from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames from roop.analyser import get_face_single import roop.ui as ui @@ -166,9 +166,6 @@ def start(preview_callback = None): 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_path: - target_path = args.target_path - args.output_path = 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_target)) if not test_face: @@ -183,14 +180,14 @@ def start(preview_callback = None): 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] + video_name_full = target_path.split(os.sep)[-1] video_name = os.path.splitext(video_name_full)[0] - output_dir = os.path.dirname(target_path) + "/" + video_name if os.path.dirname(target_path) else video_name + output_dir = os.path.dirname(target_path) + os.sep + video_name if os.path.dirname(target_path) else video_name 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: - this_path = output_dir + "/" + video_name + ".mp4" + this_path = output_dir + os.sep + video_name + ".mp4" set_fps(target_path, this_path, 30) target_path, exact_fps = this_path, 30 else: @@ -215,7 +212,7 @@ def start(preview_callback = None): 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_path) - save_path = args.output_path if args.output_path else output_dir + "/" + video_name + ".mp4" + save_path = args.output_path if args.output_path else output_dir + os.sep + video_name + ".mp4" print("\n\nVideo saved as:", save_path, "\n\n") status("swap successful!") diff --git a/roop/utils.py b/roop/utils.py index eba456b..d4b8645 100644 --- a/roop/utils.py +++ b/roop/utils.py @@ -47,8 +47,3 @@ def add_audio(output_dir, target_path, video, keep_frames, output_file): def is_img(path): return path.lower().endswith(("png", "jpg", "jpeg", "bmp")) - - -def rreplace(s, old, new, occurrence): - li = s.rsplit(old, occurrence) - return new.join(li) From 49d8103a5339aad51e77b69a428b65ccfdb67116 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:22:22 +0200 Subject: [PATCH 10/40] Run test on Linux and Windows --- .github/workflows/ci.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ef5fa6c..1ea68c5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,7 +15,10 @@ jobs: - run: pip install flake8 - run: flake8 run.py core test: - runs-on: ubuntu-latest + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ ubuntu-latest, windows-latest ] steps: - name: Checkout uses: actions/checkout@v2 From 95096442e21ea2dca99246cc2ec0fba084903d90 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:43:49 +0200 Subject: [PATCH 11/40] Run test on Linux and Windows --- .github/workflows/ci.yml | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1ea68c5..edb6c44 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,14 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - os: [ ubuntu-latest, windows-latest ] + ubuntu: + os: ubuntu-latest + run-headless: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 + run-validate: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - + windows: + os: windows-latest + run-headless: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 + run-validate: ffmpeg -i .github\examples\snapshot.mp4 -i .github/examples\output.mp4 -filter_complex psnr -f null - steps: - name: Checkout uses: actions/checkout@v2 @@ -30,6 +37,6 @@ jobs: python-version: 3.9 - run: pip install -r requirements.txt gdown - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ - - run: ./run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 - - run: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex "psnr" -f null - + - run: ${{ matrix.run-headless }} + - run: ${{ matrix.run-validate }} From 7426295471ffc03a030255512b6d84939ee61b8a Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:47:02 +0200 Subject: [PATCH 12/40] Run test on Linux and Windows --- .github/workflows/ci.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index edb6c44..cf62d5f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,13 +19,13 @@ jobs: strategy: matrix: ubuntu: - os: ubuntu-latest - run-headless: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 - run-validate: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - + os: ubuntu-latest + run-headless: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 + run-validate: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - windows: - os: windows-latest - run-headless: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 - run-validate: ffmpeg -i .github\examples\snapshot.mp4 -i .github/examples\output.mp4 -filter_complex psnr -f null - + os: windows-latest + run-headless: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 + run-validate: ffmpeg -i .github\examples\snapshot.mp4 -i .github/examples\output.mp4 -filter_complex psnr -f null - steps: - name: Checkout uses: actions/checkout@v2 From 7f8406588ebfad5eb156af53be3fba113b375eac Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:49:32 +0200 Subject: [PATCH 13/40] Run test on Linux and Windows --- .github/workflows/ci.yml | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cf62d5f..c39b294 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,15 +17,15 @@ jobs: test: runs-on: ${{ matrix.os }} strategy: - matrix: - ubuntu: - os: ubuntu-latest - run-headless: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 - run-validate: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - - windows: - os: windows-latest - run-headless: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 - run-validate: ffmpeg -i .github\examples\snapshot.mp4 -i .github/examples\output.mp4 -filter_complex psnr -f null - + matrix: + ubuntu: + os: ubuntu-latest + run-headless: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 + run-validate: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - + windows: + os: windows-latest + run-headless: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 + run-validate: ffmpeg -i .github\examples\snapshot.mp4 -i .github/examples\output.mp4 -filter_complex psnr -f null - steps: - name: Checkout uses: actions/checkout@v2 From a58a376b8747dabf23cbe8fddeda3ece37c48abc Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:52:33 +0200 Subject: [PATCH 14/40] Run test on Linux and Windows --- .github/workflows/ci.yml | 73 ++++++++++++++++++++-------------------- 1 file changed, 36 insertions(+), 37 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c39b294..3576773 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,40 +3,39 @@ name: ci on: [ push, pull_request ] jobs: - lint: - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v2 - - name: Set up Python 3.9 - uses: actions/setup-python@v2 - with: - python-version: 3.9 - - run: pip install flake8 - - run: flake8 run.py core - test: - runs-on: ${{ matrix.os }} - strategy: - matrix: - ubuntu: - os: ubuntu-latest - run-headless: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 - run-validate: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - - windows: - os: windows-latest - run-headless: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 - run-validate: ffmpeg -i .github\examples\snapshot.mp4 -i .github/examples\output.mp4 -filter_complex psnr -f null - - steps: - - name: Checkout - uses: actions/checkout@v2 - - name: Set up ffmpeg - uses: FedericoCarboni/setup-ffmpeg@v2 - - name: Set up Python 3.9 - uses: actions/setup-python@v2 - with: - python-version: 3.9 - - run: pip install -r requirements.txt gdown - - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ - - run: ${{ matrix.run-headless }} - - run: ${{ matrix.run-validate }} - + lint: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Set up Python 3.9 + uses: actions/setup-python@v2 + with: + python-version: 3.9 + - run: pip install flake8 + - run: flake8 run.py core + test: + runs-on: ${{ matrix.os }} + strategy: + matrix: + ubuntu: + os: ubuntu-latest + run-headless: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 + run-validate: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - + windows: + os: windows-latest + run-headless: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 + run-validate: ffmpeg -i .github\examples\snapshot.mp4 -i .github/examples\output.mp4 -filter_complex psnr -f null - + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Set up ffmpeg + uses: FedericoCarboni/setup-ffmpeg@v2 + - name: Set up Python 3.9 + uses: actions/setup-python@v2 + with: + python-version: 3.9 + - run: pip install -r requirements.txt gdown + - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ + - run: ${{ matrix.run-headless }} + - run: ${{ matrix.run-validate }} From 8b774af1ac45517424cbae2204adfc8ef5b351f6 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:54:10 +0200 Subject: [PATCH 15/40] Run test on Linux and Windows --- .github/workflows/ci.yml | 70 +++++++++++++++++++--------------------- 1 file changed, 34 insertions(+), 36 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3576773..7d0eec8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,39 +3,37 @@ name: ci on: [ push, pull_request ] jobs: - lint: - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v2 - - name: Set up Python 3.9 - uses: actions/setup-python@v2 - with: - python-version: 3.9 - - run: pip install flake8 - - run: flake8 run.py core - test: - runs-on: ${{ matrix.os }} - strategy: - matrix: - ubuntu: - os: ubuntu-latest - run-headless: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 - run-validate: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - - windows: - os: windows-latest - run-headless: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 - run-validate: ffmpeg -i .github\examples\snapshot.mp4 -i .github/examples\output.mp4 -filter_complex psnr -f null - - steps: - - name: Checkout - uses: actions/checkout@v2 - - name: Set up ffmpeg - uses: FedericoCarboni/setup-ffmpeg@v2 - - name: Set up Python 3.9 - uses: actions/setup-python@v2 - with: - python-version: 3.9 - - run: pip install -r requirements.txt gdown - - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ - - run: ${{ matrix.run-headless }} - - run: ${{ matrix.run-validate }} + lint: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Set up Python 3.9 + uses: actions/setup-python@v2 + with: + python-version: 3.9 + - run: pip install flake8 + - run: flake8 run.py core + test: + runs-on: ${{ matrix.os }} + strategy: + matrix: + ubuntu: + os: ubuntu-latest + run: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 + windows: + os: windows-latest + run: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Set up ffmpeg + uses: FedericoCarboni/setup-ffmpeg@v2 + - name: Set up Python 3.9 + uses: actions/setup-python@v2 + with: + python-version: 3.9 + - run: pip install -r requirements.txt gdown + - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ + - run: ${{ matrix.run }} + From 4dc4436b2fb29e4227c86095a13b8af808304bc3 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 15:57:02 +0200 Subject: [PATCH 16/40] Run test on Linux and Windows --- .github/workflows/ci.yml | 29 ++++++++++++++++++----------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7d0eec8..78efb5c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -14,16 +14,8 @@ jobs: python-version: 3.9 - run: pip install flake8 - run: flake8 run.py core - test: - runs-on: ${{ matrix.os }} - strategy: - matrix: - ubuntu: - os: ubuntu-latest - run: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 - windows: - os: windows-latest - run: python run.py -f .github\examples\face.jpg -t .github\examples\target.mp4 -o .github\examples\output.mp4 + test-ubuntu: + runs-on: ubuntu-latest steps: - name: Checkout uses: actions/checkout@v2 @@ -35,5 +27,20 @@ jobs: python-version: 3.9 - run: pip install -r requirements.txt gdown - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ - - run: ${{ matrix.run }} + - run: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 + test-windows: + runs-on: windows-latest + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Set up ffmpeg + uses: FedericoCarboni/setup-ffmpeg@v2 + - name: Set up Python 3.9 + uses: actions/setup-python@v2 + with: + python-version: 3.9 + - run: pip install -r requirements.txt gdown + - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ + - run: python run.py -f=.github\examples\face.jpg -t=.github\examples\target.mp4 -o=.github\examples\output.mp4 + From 36787d739dbd5423c9a5e31b1d455a84416bb7ca Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 16:15:33 +0200 Subject: [PATCH 17/40] Revert to Ubuntu test only as Windows hangs --- .github/workflows/ci.yml | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 78efb5c..de2ea3e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -14,7 +14,7 @@ jobs: python-version: 3.9 - run: pip install flake8 - run: flake8 run.py core - test-ubuntu: + test: runs-on: ubuntu-latest steps: - name: Checkout @@ -28,19 +28,3 @@ jobs: - run: pip install -r requirements.txt gdown - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ - run: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 - test-windows: - runs-on: windows-latest - steps: - - name: Checkout - uses: actions/checkout@v2 - - name: Set up ffmpeg - uses: FedericoCarboni/setup-ffmpeg@v2 - - name: Set up Python 3.9 - uses: actions/setup-python@v2 - with: - python-version: 3.9 - - run: pip install -r requirements.txt gdown - - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ - - run: python run.py -f=.github\examples\face.jpg -t=.github\examples\target.mp4 -o=.github\examples\output.mp4 - - From 75b87f301901092d488105dec749c03877d81bf4 Mon Sep 17 00:00:00 2001 From: Antoine Buchser <10513467+AntwaneB@users.noreply.github.com> Date: Mon, 5 Jun 2023 19:34:38 +0100 Subject: [PATCH 18/40] Simplified the way to maintain aspect ratio of the preview, and maintaining aspect ratio of the miniatures --- roop/ui.py | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/roop/ui.py b/roop/ui.py index d8891fb..bedc4d4 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -1,6 +1,6 @@ import tkinter as tk from typing import Any, Callable, Tuple -from PIL import Image, ImageTk +from PIL import Image, ImageTk, ImageOps import webbrowser from tkinter import filedialog from tkinter.filedialog import asksaveasfilename @@ -69,19 +69,7 @@ def init_slider(frames_count, change_handler): def update_preview(frame): img = Image.fromarray(frame) - width, height = img.size - aspect_ratio = 1 - if width > height: - aspect_ratio = max_preview_size / width - else: - aspect_ratio = max_preview_size / height - img = img.resize( - ( - int(width * aspect_ratio), - int(height * aspect_ratio) - ), - Image.ANTIALIAS - ) + img = ImageOps.contain(img, (max_preview_size, max_preview_size), Image.LANCZOS) photo_img = ImageTk.PhotoImage(img) preview_image_frame.configure(image=photo_img) preview_image_frame.image = photo_img @@ -211,7 +199,7 @@ def open_preview_window(get_video_frame, target_path): def preview_face(path): img = Image.open(path) - img = img.resize((180, 180), Image.ANTIALIAS) + img = ImageOps.contain(img, (180, 180), Image.LANCZOS) photo_img = ImageTk.PhotoImage(img) face_label.configure(image=photo_img) face_label.image = photo_img @@ -219,7 +207,7 @@ def preview_face(path): def preview_target(frame): img = Image.fromarray(frame) - img = img.resize((180, 180), Image.ANTIALIAS) + img = ImageOps.contain(img, (180, 180), Image.LANCZOS) photo_img = ImageTk.PhotoImage(img) target_label.configure(image=photo_img) target_label.image = photo_img @@ -312,4 +300,4 @@ def init( status_label = tk.Label(window, width=580, justify="center", text="Status: waiting for input...", fg="#2ecc71", bg="#2d3436") status_label.place(x=10,y=640,width=580,height=30) - return window \ No newline at end of file + return window From 58da88e5d6f3ace18c5ea0e4b11f686179e30daf Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 5 Jun 2023 20:45:36 +0200 Subject: [PATCH 19/40] Change face and target images from contain to fit --- roop/ui.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/roop/ui.py b/roop/ui.py index bedc4d4..18863e2 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -199,7 +199,7 @@ def open_preview_window(get_video_frame, target_path): def preview_face(path): img = Image.open(path) - img = ImageOps.contain(img, (180, 180), Image.LANCZOS) + img = ImageOps.fit(img, (180, 180), Image.LANCZOS) photo_img = ImageTk.PhotoImage(img) face_label.configure(image=photo_img) face_label.image = photo_img @@ -207,7 +207,7 @@ def preview_face(path): def preview_target(frame): img = Image.fromarray(frame) - img = ImageOps.contain(img, (180, 180), Image.LANCZOS) + img = ImageOps.fit(img, (180, 180), Image.LANCZOS) photo_img = ImageTk.PhotoImage(img) target_label.configure(image=photo_img) target_label.image = photo_img From dcfd6cca4dc62a357fa4fc745e8ff8d46646acdd Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 6 Jun 2023 00:21:23 +0200 Subject: [PATCH 20/40] Improve status output --- roop/core.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/roop/core.py b/roop/core.py index f93785b..d633a79 100755 --- a/roop/core.py +++ b/roop/core.py @@ -140,9 +140,8 @@ def preview_video(video_path): def status(string): value = "Status: " + string - if roop.globals.headless: - print(value) - else: + print(value) + if not roop.globals.headless: ui.update_status_label(value) From 09ea59f66f986ca330792badb401e4d494b2ed54 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 6 Jun 2023 12:30:27 +0200 Subject: [PATCH 21/40] Massive utilities and core refactoring --- roop/core.py | 211 ++++++++++++++++++++++++---------------------- roop/globals.py | 5 +- roop/swapper.py | 10 +-- roop/ui.py | 4 +- roop/utilities.py | 92 ++++++++++++++++++++ roop/utils.py | 49 ----------- 6 files changed, 213 insertions(+), 158 deletions(-) create mode 100644 roop/utilities.py delete mode 100644 roop/utils.py diff --git a/roop/core.py b/roop/core.py index d633a79..92bc897 100755 --- a/roop/core.py +++ b/roop/core.py @@ -2,37 +2,39 @@ import os import sys +from typing import List +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # single thread doubles performance of gpu-mode - needs to be set before torch import if any(arg.startswith('--gpu-vendor') for arg in sys.argv): os.environ['OMP_NUM_THREADS'] = '1' import platform import signal import shutil -import glob import argparse import psutil import torch import tensorflow -from pathlib import Path -import multiprocessing as mp +import multiprocessing from opennsfw2 import predict_video_frames, predict_image import cv2 import roop.globals -from roop.swapper import process_video, process_img, process_faces, process_frames -from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames +from roop.swapper import process_video, process_img, process_faces +from roop.utilities import has_image_extention, is_image, detect_fps, create_video, extract_frames, \ + get_temp_frames_paths, restore_audio, create_temp, clean_temp, is_video from roop.analyser import get_face_single 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: destroy()) parser = argparse.ArgumentParser() - parser.add_argument('-f', '--face', help='use this face', dest='source_target') + parser.add_argument('-f', '--face', help='use this face', dest='source_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_path') parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) + parser.add_argument('--keep-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True) 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('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) @@ -42,7 +44,10 @@ def handle_parse(): args = parser.parse_known_args()[0] - roop.globals.headless = args.source_target or args.target_path or args.output_path + roop.globals.headless = args.source_path or args.target_path or args.output_path + roop.globals.keep_fps = args.keep_fps + roop.globals.keep_audio = args.keep_audio + roop.globals.keep_frames = args.keep_frames roop.globals.all_faces = args.all_faces if args.cpu_cores: @@ -83,7 +88,7 @@ def limit_resources(): def pre_check(): if sys.version_info < (3, 9): - quit('Python version is not supported - please upgrade to 3.9 or higher') + quit('Python version is not supported - please upgrade to 3.9 or higher.') if not shutil.which('ffmpeg'): quit('ffmpeg is not installed!') model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), '../inswapper_128.onnx') @@ -91,23 +96,23 @@ def pre_check(): quit('File "inswapper_128.onnx" does not exist!') if roop.globals.gpu_vendor == 'apple': if 'CoreMLExecutionProvider' not in roop.globals.providers: - quit("You are using --gpu=apple flag but CoreML isn't available or properly installed on your system.") + quit('You are using --gpu=apple flag but CoreML is not available or properly installed on your system.') if roop.globals.gpu_vendor == 'amd': if 'ROCMExecutionProvider' not in roop.globals.providers: - quit("You are using --gpu=amd flag but ROCM isn't available or properly installed on your system.") + quit('You are using --gpu=amd flag but ROCM is not available or properly installed on your system.') if roop.globals.gpu_vendor == 'nvidia': CUDA_VERSION = torch.version.cuda CUDNN_VERSION = torch.backends.cudnn.version() if not torch.cuda.is_available(): - quit("You are using --gpu=nvidia flag but CUDA isn't available or properly installed on your system.") + quit('You are using --gpu=nvidia flag but CUDA is not available or properly installed on your system.') if CUDA_VERSION > '11.8': - quit(f"CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8") + quit(f'CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8') if CUDA_VERSION < '11.4': - quit(f"CUDA version {CUDA_VERSION} is not supported - please upgrade to 11.8") + quit(f'CUDA version {CUDA_VERSION} is not supported - please upgrade to 11.8') if CUDNN_VERSION < 8220: - quit(f"CUDNN version {CUDNN_VERSION} is not supported - please upgrade to 8.9.1") + quit(f'CUDNN version {CUDNN_VERSION} is not supported - please upgrade to 8.9.1') if CUDNN_VERSION > 8910: - quit(f"CUDNN version {CUDNN_VERSION} is not supported - please downgrade to 8.9.1") + quit(f'CUDNN version {CUDNN_VERSION} is not supported - please downgrade to 8.9.1') def get_video_frame(video_path, frame_number = 1): @@ -115,19 +120,18 @@ def get_video_frame(video_path, frame_number = 1): amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1)) if not cap.isOpened(): - print("Error opening video file") + status('Error opening video file') return ret, frame = cap.read() if ret: return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - cap.release() def preview_video(video_path): cap = cv2.VideoCapture(video_path) if not cap.isOpened(): - print("Error opening video file") + status('Error opening video file') return 0 amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) ret, frame = cap.read() @@ -138,86 +142,89 @@ def preview_video(video_path): return (amount_of_frames, frame) -def status(string): - value = "Status: " + string +def status(message: str): + value = 'Status: ' + message print(value) if not roop.globals.headless: ui.update_status_label(value) -def process_video_multi_cores(source_target, frame_paths): - n = len(frame_paths) // roop.globals.cpu_cores - if n > 2: - processes = [] - for i in range(0, len(frame_paths), n): - p = POOL.apply_async(process_video, args=(source_target, frame_paths[i:i + n],)) - processes.append(p) - for p in processes: - p.get() - POOL.close() +def conditional_process_video(source_path: str, frame_paths: List[str]) -> None: + pool_amount = len(frame_paths) // roop.globals.cpu_cores + if pool_amount > 2 and roop.globals.cpu_cores > 1 and roop.globals.gpu_vendor is None: + status('Pool-Swapping in progress...') + global POOL + POOL = multiprocessing.Pool(roop.globals.cpu_cores, maxtasksperchild=1) + pools = [] + for i in range(0, len(frame_paths), pool_amount): + pool = POOL.apply_async(process_video, args=(source_path, frame_paths[i:i + pool_amount])) + pools.append(pool) + for pool in pools: + pool.get() POOL.join() + POOL.close() + else: + status('Swapping in progress...') + process_video(args.source_path, frame_paths) -def start(preview_callback = None): - if not args.source_target or not os.path.isfile(args.source_target): - print("\n[WARNING] Please select an image containing a face.") +def start(preview_callback = None) -> None: + if not args.source_path or not os.path.isfile(args.source_path): + status('Please select an image containing a face.') return 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.") + status('Please select a video/image target!') return - target_path = args.target_path - test_face = get_face_single(cv2.imread(args.source_target)) + test_face = get_face_single(cv2.imread(args.source_path)) if not test_face: - print("\n[WARNING] No face detected in source image. Please try with another one.\n") + status('No face detected in source image. Please try with another one!') return - if is_img(target_path): - if predict_image(target_path) > 0.85: - quit() - process_img(args.source_target, target_path, args.output_path) - status("swap successful!") + # process image to image + if has_image_extention(args.target_path): + if predict_image(args.target_path) > 0.85: + destroy() + process_img(args.source_path, args.target_path, args.output_path) + if is_image(args.target_path): + status('Swapping to image succeed!') + else: + status('Swapping to image failed!') return + # process image to videos 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(os.sep)[-1] - video_name = os.path.splitext(video_name_full)[0] - output_dir = os.path.dirname(target_path) + os.sep + video_name if os.path.dirname(target_path) else video_name - 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: - this_path = output_dir + os.sep + video_name + ".mp4" - set_fps(target_path, this_path, 30) - target_path, exact_fps = this_path, 30 - else: - shutil.copy(target_path, output_dir) - status("extracting frames...") - extract_frames(target_path, output_dir) - args.frame_paths = tuple(sorted( - glob.glob(output_dir + "/*.png"), - key=lambda x: int(x.split(os.sep)[-1].replace(".png", "")) - )) - status("swapping in progress...") - if roop.globals.gpu_vendor is None and roop.globals.cpu_cores > 1: - global POOL - POOL = mp.Pool(roop.globals.cpu_cores) - process_video_multi_cores(args.source_target, args.frame_paths) - else: - process_video(args.source_target, args.frame_paths) - # prevent out of memory while using ffmpeg with cuda + destroy() + status('Creating temp resources...') + create_temp(args.target_path) + status('Extracting frames...') + extract_frames(args.target_path) + frame_paths = get_temp_frames_paths(args.target_path) + conditional_process_video(args.source_path, frame_paths) + # prevent memory leak using ffmpeg with cuda if args.gpu_vendor == 'nvidia': torch.cuda.empty_cache() - 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_path) - save_path = args.output_path if args.output_path else output_dir + os.sep + video_name + ".mp4" - print("\n\nVideo saved as:", save_path, "\n\n") - status("swap successful!") + if roop.globals.keep_fps: + status('Detecting fps...') + fps = detect_fps(args.source_path) + status(f'Creating video with {fps} fps...') + create_video(args.target_path, fps) + else: + status('Creating video with 30 fps...') + create_video(args.target_path, 30) + if roop.globals.keep_audio: + if roop.globals.keep_fps: + status('Restoring audio...') + else: + status('Restoring audio might cause issues as fps are not kept...') + restore_audio(args.target_path) + clean_temp(args.target_path, args.output_path) + if is_video(args.target_path): + status('Swapping to video succeed!') + else: + status('Swapping to video failed!') def select_face_handler(path: str): - args.source_target = path + args.source_path = path def select_target_handler(path: str): @@ -243,33 +250,37 @@ def save_file_handler(path: str): def create_test_preview(frame_number): return process_faces( - get_face_single(cv2.imread(args.source_target)), + get_face_single(cv2.imread(args.source_path)), get_video_frame(args.target_path, frame_number) ) -def run(): +def destroy() -> None: + clean_temp(args.target_path) + quit() + + +def run() -> None: global all_faces, keep_frames, limit_fps handle_parse() pre_check() limit_resources() - if roop.globals.headless: - start() - quit() - window = ui.init( - { - 'all_faces': roop.globals.all_faces, - 'keep_fps': args.keep_fps, - 'keep_frames': args.keep_frames - }, - select_face_handler, - select_target_handler, - toggle_all_faces_handler, - toggle_fps_limit_handler, - toggle_keep_frames_handler, - save_file_handler, - start, - get_video_frame, - create_test_preview - ) - window.mainloop() + start() + if not roop.globals.headless: + window = ui.init( + { + 'all_faces': args.all_faces, + 'keep_fps': args.keep_fps, + 'keep_frames': args.keep_frames + }, + select_face_handler, + select_target_handler, + toggle_all_faces_handler, + toggle_fps_limit_handler, + toggle_keep_frames_handler, + save_file_handler, + start, + get_video_frame, + create_test_preview + ) + window.mainloop() diff --git a/roop/globals.py b/roop/globals.py index 9162a39..100c193 100644 --- a/roop/globals.py +++ b/roop/globals.py @@ -1,11 +1,14 @@ import onnxruntime +keep_fps = None +keep_audio = None +keep_frames = None all_faces = None -log_level = 'error' cpu_cores = None gpu_threads = None gpu_vendor = None headless = None +log_level = 'error' providers = onnxruntime.get_available_providers() if 'TensorrtExecutionProvider' in providers: diff --git a/roop/swapper.py b/roop/swapper.py index 81a6b1d..5f2eb1f 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -58,9 +58,8 @@ def multi_process_frame(source_img, frame_paths, progress): num_threads = roop.globals.gpu_threads num_frames_per_thread = len(frame_paths) // num_threads remaining_frames = len(frame_paths) % num_threads - - # create thread and launch start_index = 0 + # create threads by frames for _ in range(num_threads): end_index = start_index + num_frames_per_thread if remaining_frames > 0: @@ -71,8 +70,7 @@ def multi_process_frame(source_img, frame_paths, progress): threads.append(thread) thread.start() start_index = end_index - - # threading + # join threads for thread in threads: thread.join() @@ -83,13 +81,13 @@ def process_img(source_img, target_path, output_file): source_face = get_face_single(cv2.imread(source_img)) result = get_face_swapper().get(frame, face, source_face, paste_back=True) cv2.imwrite(output_file, result) - print("\n\nImage saved as:", output_file, "\n\n") def process_video(source_img, frame_paths): do_multi = roop.globals.gpu_vendor is not None and roop.globals.gpu_threads > 1 progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' - with tqdm(total=len(frame_paths), desc="Processing", unit="frame", dynamic_ncols=True, bar_format=progress_bar_format) as progress: + total = len(frame_paths) + with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress: if do_multi: multi_process_frame(source_img, frame_paths, progress) else: diff --git a/roop/ui.py b/roop/ui.py index 18863e2..bbca8bf 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -6,7 +6,7 @@ from tkinter import filedialog from tkinter.filedialog import asksaveasfilename import threading -from roop.utils import is_img +from roop.utilities import is_image max_preview_size = 800 @@ -114,7 +114,7 @@ def select_target(select_target_handler: Callable[[str], Tuple[int, Any]], targe def save_file(save_file_handler: Callable[[str], None], target_path: str): filename, ext = 'output.mp4', '.mp4' - if is_img(target_path): + if is_image(target_path): filename, ext = 'output.png', '.png' if save_file_handler: diff --git a/roop/utilities.py b/roop/utilities.py new file mode 100644 index 0000000..29789db --- /dev/null +++ b/roop/utilities.py @@ -0,0 +1,92 @@ +import glob +import os +import shutil +import subprocess +from pathlib import Path +from typing import List, Any + +import roop.globals +from PIL import Image + + +def run_ffmpeg(args: List) -> None: + commands = ['ffmpeg', '-hide_banner', '-hwaccel', 'auto', '-loglevel', roop.globals.log_level] + commands.extend(args) + try: + subprocess.check_output(commands, stderr=subprocess.STDOUT) + except Exception: + pass + + +def detect_fps(source_path: str) -> int: + command = ['ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries', 'stream=r_frame_rate', '-of', 'default=noprint_wrappers=1:nokey=1', source_path] + output = subprocess.check_output(command).decode().strip() + try: + return int(eval(output)) + except Exception: + pass + return 30 + + +def extract_frames(target_path: str) -> None: + temp_directory_path = get_temp_directory_path(target_path) + run_ffmpeg(['-i', target_path, temp_directory_path + os.sep + '%04d.png']) + + +def create_video(target_path: str, fps: int) -> None: + temp_directory_path = get_temp_directory_path(target_path) + temp_file_path = get_temp_file_path(target_path) + run_ffmpeg(['-i', temp_directory_path + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', '7', '-pix_fmt', 'yuv420p', '-y', temp_file_path]) + + +def restore_audio(target_path: str) -> None: + run_ffmpeg(['-i', target_path, '-i', get_temp_file_path(target_path), '-c:v', 'copy', '-map', '0:v:0', 'map', '1:a:0', '-y', get_temp_file_path(target_path)]) + + +def get_temp_frames_paths(target_path: str) -> List: + return glob.glob(get_temp_directory_path(target_path) + os.sep + '*.png') + + +def get_temp_directory_path(target_path: str) -> str: + target_directory_path = os.path.dirname(target_path) + return target_directory_path + os.sep + 'temp' + + +def get_temp_file_path(target_path: str) -> str: + return get_temp_directory_path(target_path) + os.sep + 'temp.mp4' + + +def create_temp(target_path: str) -> None: + Path(get_temp_directory_path(target_path)).mkdir(exist_ok=True) + + +def clean_temp(target_path: str, output_path: str) -> None: + temp_file_path = get_temp_file_path(target_path) + if os.path.isfile(temp_file_path): + shutil.move(temp_file_path, output_path) + if not roop.globals.keep_frames: + shutil.rmtree(get_temp_directory_path(target_path)) + + +def has_image_extention(image_path: str) -> bool: + return image_path.lower().endswith(('png', 'jpg', 'jpeg', 'bmp')) + + +def is_image(path: str) -> bool: + if os.path.isfile(path): + try: + image = Image.open(path) + image.verify() + return True + except Exception: + pass + return False + + +def is_video(path: str) -> bool: + try: + run_ffmpeg(['-v', 'error', '-i', path, '-f', 'null', '-']) + return True + except subprocess.CalledProcessError: + pass + return False diff --git a/roop/utils.py b/roop/utils.py deleted file mode 100644 index d4b8645..0000000 --- a/roop/utils.py +++ /dev/null @@ -1,49 +0,0 @@ -import os -import shutil -import roop.globals - - -def run_command(command, mode="silent"): - if mode == "debug": - return os.system(command) - return os.popen(command).read() - - -def detect_fps(input_path): - output = os.popen(f'ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate "{input_path}"').read() - if "/" in output: - try: - return int(output.split("/")[0]) // int(output.split("/")[1].strip()), output.strip() - except: - pass - return 30, 30 - - -def run_ffmpeg(args): - run_command(f'ffmpeg -hide_banner -hwaccel auto -loglevel {roop.globals.log_level} {args}') - - -def set_fps(input_path, output_path, fps): - run_ffmpeg(f'-i "{input_path}" -filter:v fps=fps={fps} "{output_path}"') - - -def create_video(video_name, fps, output_dir): - run_ffmpeg(f'-framerate "{fps}" -i "{output_dir}{os.sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{os.sep}output.mp4"') - - -def extract_frames(input_path, output_dir): - run_ffmpeg(f'-i "{input_path}" "{output_dir}{os.sep}%04d.png"') - - -def add_audio(output_dir, target_path, video, keep_frames, output_file): - video_name = os.path.splitext(video)[0] - save_to = output_file if output_file else output_dir + "/swapped-" + video_name + ".mp4" - run_ffmpeg(f'-i "{output_dir}{os.sep}output.mp4" -i "{output_dir}{os.sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to}"') - if not os.path.isfile(save_to): - shutil.move(output_dir + "/output.mp4", save_to) - if not keep_frames: - shutil.rmtree(output_dir) - - -def is_img(path): - return path.lower().endswith(("png", "jpg", "jpeg", "bmp")) From 7990896c9e7d2bad9b3962831b57322e7610051a Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 6 Jun 2023 14:10:04 +0200 Subject: [PATCH 22/40] Fix sound --- roop/core.py | 4 ++-- roop/utilities.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/roop/core.py b/roop/core.py index 92bc897..2bf3ebd 100755 --- a/roop/core.py +++ b/roop/core.py @@ -215,7 +215,7 @@ def start(preview_callback = None) -> None: status('Restoring audio...') else: status('Restoring audio might cause issues as fps are not kept...') - restore_audio(args.target_path) + restore_audio(args.target_path, args.output_path) clean_temp(args.target_path, args.output_path) if is_video(args.target_path): status('Swapping to video succeed!') @@ -256,7 +256,7 @@ def create_test_preview(frame_number): def destroy() -> None: - clean_temp(args.target_path) + clean_temp(args.target_path, args.output_path) quit() diff --git a/roop/utilities.py b/roop/utilities.py index 29789db..acf1632 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -39,8 +39,8 @@ def create_video(target_path: str, fps: int) -> None: run_ffmpeg(['-i', temp_directory_path + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', '7', '-pix_fmt', 'yuv420p', '-y', temp_file_path]) -def restore_audio(target_path: str) -> None: - run_ffmpeg(['-i', target_path, '-i', get_temp_file_path(target_path), '-c:v', 'copy', '-map', '0:v:0', 'map', '1:a:0', '-y', get_temp_file_path(target_path)]) +def restore_audio(target_path: str, output_path: str) -> None: + run_ffmpeg(['-i', get_temp_file_path(target_path), '-i', target_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path]) def get_temp_frames_paths(target_path: str) -> List: @@ -62,7 +62,7 @@ def create_temp(target_path: str) -> None: def clean_temp(target_path: str, output_path: str) -> None: temp_file_path = get_temp_file_path(target_path) - if os.path.isfile(temp_file_path): + if not roop.globals.keep_audio and os.path.isfile(temp_file_path): shutil.move(temp_file_path, output_path) if not roop.globals.keep_frames: shutil.rmtree(get_temp_directory_path(target_path)) From 957c5ec6ce01812d7e8a48ad4585952b59fa99bc Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 6 Jun 2023 15:01:27 +0200 Subject: [PATCH 23/40] Fix sound part2 --- .github/workflows/ci.yml | 2 ++ roop/core.py | 16 ++++++++++------ roop/utilities.py | 21 +++++++++++---------- 3 files changed, 23 insertions(+), 16 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index de2ea3e..c27c904 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,3 +28,5 @@ jobs: - run: pip install -r requirements.txt gdown - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ - run: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 + - run: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - + diff --git a/roop/core.py b/roop/core.py index 2bf3ebd..aaa9571 100755 --- a/roop/core.py +++ b/roop/core.py @@ -2,11 +2,13 @@ import os import sys -from typing import List os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # single thread doubles performance of gpu-mode - needs to be set before torch import if any(arg.startswith('--gpu-vendor') for arg in sys.argv): os.environ['OMP_NUM_THREADS'] = '1' +import warnings +warnings.simplefilter(action='ignore', category=FutureWarning) +from typing import List import platform import signal import shutil @@ -20,8 +22,7 @@ import cv2 import roop.globals from roop.swapper import process_video, process_img, process_faces -from roop.utilities import has_image_extention, is_image, detect_fps, create_video, extract_frames, \ - get_temp_frames_paths, restore_audio, create_temp, clean_temp, is_video +from roop.utilities import has_image_extention, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frames_paths, restore_audio, create_temp, move_temp, clean_temp from roop.analyser import get_face_single import roop.ui as ui @@ -216,7 +217,9 @@ def start(preview_callback = None) -> None: else: status('Restoring audio might cause issues as fps are not kept...') restore_audio(args.target_path, args.output_path) - clean_temp(args.target_path, args.output_path) + else: + move_temp(args.target_path, args.output_path) + clean_temp(args.target_path) if is_video(args.target_path): status('Swapping to video succeed!') else: @@ -265,8 +268,9 @@ def run() -> None: handle_parse() pre_check() limit_resources() - start() - if not roop.globals.headless: + if roop.globals.headless: + start() + else: window = ui.init( { 'all_faces': args.all_faces, diff --git a/roop/utilities.py b/roop/utilities.py index acf1632..487cbb7 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -14,7 +14,7 @@ def run_ffmpeg(args: List) -> None: commands.extend(args) try: subprocess.check_output(commands, stderr=subprocess.STDOUT) - except Exception: + except Exception as exception: pass @@ -29,18 +29,17 @@ def detect_fps(source_path: str) -> int: def extract_frames(target_path: str) -> None: - temp_directory_path = get_temp_directory_path(target_path) - run_ffmpeg(['-i', target_path, temp_directory_path + os.sep + '%04d.png']) + run_ffmpeg(['-i', target_path, get_temp_directory_path(target_path) + os.sep + '%04d.png']) def create_video(target_path: str, fps: int) -> None: - temp_directory_path = get_temp_directory_path(target_path) - temp_file_path = get_temp_file_path(target_path) - run_ffmpeg(['-i', temp_directory_path + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', '7', '-pix_fmt', 'yuv420p', '-y', temp_file_path]) + run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', '7', '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) def restore_audio(target_path: str, output_path: str) -> None: run_ffmpeg(['-i', get_temp_file_path(target_path), '-i', target_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path]) + if not os.path.isfile(output_path): + move_temp(target_path, output_path) def get_temp_frames_paths(target_path: str) -> List: @@ -48,8 +47,7 @@ def get_temp_frames_paths(target_path: str) -> List: def get_temp_directory_path(target_path: str) -> str: - target_directory_path = os.path.dirname(target_path) - return target_directory_path + os.sep + 'temp' + return os.path.dirname(target_path) + os.sep + 'temp' def get_temp_file_path(target_path: str) -> str: @@ -60,10 +58,13 @@ def create_temp(target_path: str) -> None: Path(get_temp_directory_path(target_path)).mkdir(exist_ok=True) -def clean_temp(target_path: str, output_path: str) -> None: +def move_temp(target_path: str, output_path: str) -> None: temp_file_path = get_temp_file_path(target_path) - if not roop.globals.keep_audio and os.path.isfile(temp_file_path): + if os.path.isfile(temp_file_path): shutil.move(temp_file_path, output_path) + + +def clean_temp(target_path: str) -> None: if not roop.globals.keep_frames: shutil.rmtree(get_temp_directory_path(target_path)) From a947a6e293ed403335546b0afabf59e0ca18f299 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 6 Jun 2023 15:13:52 +0200 Subject: [PATCH 24/40] Fix more --- roop/core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/roop/core.py b/roop/core.py index aaa9571..fb86d54 100755 --- a/roop/core.py +++ b/roop/core.py @@ -7,7 +7,6 @@ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' if any(arg.startswith('--gpu-vendor') for arg in sys.argv): os.environ['OMP_NUM_THREADS'] = '1' import warnings -warnings.simplefilter(action='ignore', category=FutureWarning) from typing import List import platform import signal @@ -26,6 +25,7 @@ from roop.utilities import has_image_extention, is_image, is_video, detect_fps, from roop.analyser import get_face_single import roop.ui as ui +warnings.simplefilter(action='ignore', category=FutureWarning) def handle_parse(): global args @@ -41,7 +41,7 @@ def handle_parse(): parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) - parser.add_argument('--gpu-vendor', help='choice your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia']) + parser.add_argument('--gpu-vendor', help='select your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia']) args = parser.parse_known_args()[0] @@ -259,7 +259,7 @@ def create_test_preview(frame_number): def destroy() -> None: - clean_temp(args.target_path, args.output_path) + clean_temp(args.target_path) quit() From 419458485431a899859bbff2b06fa63037495780 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 6 Jun 2023 19:29:00 +0200 Subject: [PATCH 25/40] Move every UI related thing to ui.py --- roop/analyser.py | 15 ++-- roop/core.py | 213 ++++++++++++++++------------------------------- roop/globals.py | 4 + roop/ui.py | 89 ++++++++++++++++---- 4 files changed, 156 insertions(+), 165 deletions(-) diff --git a/roop/analyser.py b/roop/analyser.py index 804f7a8..c2899e7 100644 --- a/roop/analyser.py +++ b/roop/analyser.py @@ -1,10 +1,11 @@ +from typing import Any import insightface import roop.globals FACE_ANALYSER = None -def get_face_analyser(): +def get_face_analyser() -> Any: global FACE_ANALYSER if FACE_ANALYSER is None: FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=roop.globals.providers) @@ -12,16 +13,16 @@ def get_face_analyser(): return FACE_ANALYSER -def get_face_single(img_data): - face = get_face_analyser().get(img_data) +def get_face_single(image_data) -> Any: + face = get_face_analyser().get(image_data) try: - return sorted(face, key=lambda x: x.bbox[0])[0] - except IndexError: + return min(face, key=lambda x: x.bbox[0]) + except ValueError: return None -def get_face_many(img_data): +def get_face_many(image_data) -> Any: try: - return get_face_analyser().get(img_data) + return get_face_analyser().get(image_data) except IndexError: return None diff --git a/roop/core.py b/roop/core.py index fb86d54..48bd2d7 100755 --- a/roop/core.py +++ b/roop/core.py @@ -2,10 +2,11 @@ import os import sys -os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # single thread doubles performance of gpu-mode - needs to be set before torch import if any(arg.startswith('--gpu-vendor') for arg in sys.argv): os.environ['OMP_NUM_THREADS'] = '1' +# reduce tensorflow log level +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' import warnings from typing import List import platform @@ -20,15 +21,18 @@ from opennsfw2 import predict_video_frames, predict_image import cv2 import roop.globals -from roop.swapper import process_video, process_img, process_faces +import roop.ui as ui +from roop.swapper import process_video, process_img from roop.utilities import has_image_extention, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frames_paths, restore_audio, create_temp, move_temp, clean_temp from roop.analyser import get_face_single -import roop.ui as ui + +if 'ROCMExecutionProvider' in roop.globals.providers: + del torch warnings.simplefilter(action='ignore', category=FutureWarning) -def handle_parse(): - global args + +def parse_args() -> None: signal.signal(signal.SIGINT, lambda signal_number, frame: destroy()) parser = argparse.ArgumentParser() parser.add_argument('-f', '--face', help='use this face', dest='source_path') @@ -45,6 +49,9 @@ def handle_parse(): args = parser.parse_known_args()[0] + roop.globals.source_path = args.source_path + roop.globals.target_path = args.target_path + roop.globals.output_path = args.output_path roop.globals.headless = args.source_path or args.target_path or args.output_path roop.globals.keep_fps = args.keep_fps roop.globals.keep_audio = args.keep_audio @@ -76,8 +83,8 @@ def limit_resources(): gpus = tensorflow.config.experimental.list_physical_devices('GPU') for gpu in gpus: tensorflow.config.experimental.set_memory_growth(gpu, True) - if args.max_memory: - memory = args.max_memory * 1024 * 1024 * 1024 + if roop.globals.max_memory: + memory = roop.globals.max_memory * 1024 * 1024 * 1024 if str(platform.system()).lower() == 'windows': import ctypes kernel32 = ctypes.windll.kernel32 @@ -102,58 +109,22 @@ def pre_check(): if 'ROCMExecutionProvider' not in roop.globals.providers: quit('You are using --gpu=amd flag but ROCM is not available or properly installed on your system.') if roop.globals.gpu_vendor == 'nvidia': - CUDA_VERSION = torch.version.cuda - CUDNN_VERSION = torch.backends.cudnn.version() if not torch.cuda.is_available(): quit('You are using --gpu=nvidia flag but CUDA is not available or properly installed on your system.') - if CUDA_VERSION > '11.8': - quit(f'CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8') - if CUDA_VERSION < '11.4': - quit(f'CUDA version {CUDA_VERSION} is not supported - please upgrade to 11.8') - if CUDNN_VERSION < 8220: - quit(f'CUDNN version {CUDNN_VERSION} is not supported - please upgrade to 8.9.1') - if CUDNN_VERSION > 8910: - quit(f'CUDNN version {CUDNN_VERSION} is not supported - please downgrade to 8.9.1') - - -def get_video_frame(video_path, frame_number = 1): - cap = cv2.VideoCapture(video_path) - amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) - cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1)) - if not cap.isOpened(): - status('Error opening video file') - return - ret, frame = cap.read() - if ret: - return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - cap.release() - - -def preview_video(video_path): - cap = cv2.VideoCapture(video_path) - if not cap.isOpened(): - status('Error opening video file') - return 0 - amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) - ret, frame = cap.read() - if ret: - frame = get_video_frame(video_path) - - cap.release() - return (amount_of_frames, frame) - - -def status(message: str): - value = 'Status: ' + message - print(value) - if not roop.globals.headless: - ui.update_status_label(value) + if torch.version.cuda > '11.8': + quit(f'CUDA version {torch.version.cuda} is not supported - please downgrade to 11.8') + if torch.version.cuda < '11.4': + quit(f'CUDA version {torch.version.cuda} is not supported - please upgrade to 11.8') + if torch.backends.cudnn.version() < 8220: + quit(f'CUDNN version { torch.backends.cudnn.version()} is not supported - please upgrade to 8.9.1') + if torch.backends.cudnn.version() > 8910: + quit(f'CUDNN version { torch.backends.cudnn.version()} is not supported - please downgrade to 8.9.1') def conditional_process_video(source_path: str, frame_paths: List[str]) -> None: pool_amount = len(frame_paths) // roop.globals.cpu_cores if pool_amount > 2 and roop.globals.cpu_cores > 1 and roop.globals.gpu_vendor is None: - status('Pool-Swapping in progress...') + update_status('Pool-Swapping in progress...') global POOL POOL = multiprocessing.Pool(roop.globals.cpu_cores, maxtasksperchild=1) pools = [] @@ -162,129 +133,89 @@ def conditional_process_video(source_path: str, frame_paths: List[str]) -> None: pools.append(pool) for pool in pools: pool.get() - POOL.join() POOL.close() + POOL.join() else: - status('Swapping in progress...') - process_video(args.source_path, frame_paths) + update_status('Swapping in progress...') + process_video(roop.globals.source_path, frame_paths) -def start(preview_callback = None) -> None: - if not args.source_path or not os.path.isfile(args.source_path): - status('Please select an image containing a face.') +def update_status(message: str): + value = 'Status: ' + message + print(value) + if not roop.globals.headless: + ui.update_status(value) + + +def start() -> None: + if not roop.globals.source_path or not os.path.isfile(roop.globals.source_path): + update_status('Please select an image containing a face.') return - elif not args.target_path or not os.path.isfile(args.target_path): - status('Please select a video/image target!') + elif not roop.globals.target_path or not os.path.isfile(roop.globals.target_path): + update_status('Please select a video/image target!') return - test_face = get_face_single(cv2.imread(args.source_path)) + test_face = get_face_single(cv2.imread(roop.globals.source_path)) if not test_face: - status('No face detected in source image. Please try with another one!') + update_status('No face detected in source image. Please try with another one!') return # process image to image - if has_image_extention(args.target_path): - if predict_image(args.target_path) > 0.85: + if has_image_extention(roop.globals.target_path): + if predict_image(roop.globals.target_path) > 0.85: destroy() - process_img(args.source_path, args.target_path, args.output_path) - if is_image(args.target_path): - status('Swapping to image succeed!') + process_img(roop.globals.source_path, roop.globals.target_path, roop.globals.output_path) + if is_image(roop.globals.target_path): + update_status('Swapping to image succeed!') else: - status('Swapping to image failed!') + update_status('Swapping to image failed!') return # process image to videos - seconds, probabilities = predict_video_frames(video_path=args.target_path, frame_interval=100) + seconds, probabilities = predict_video_frames(video_path=roop.globals.target_path, frame_interval=100) if any(probability > 0.85 for probability in probabilities): destroy() - status('Creating temp resources...') - create_temp(args.target_path) - status('Extracting frames...') - extract_frames(args.target_path) - frame_paths = get_temp_frames_paths(args.target_path) - conditional_process_video(args.source_path, frame_paths) + update_status('Creating temp resources...') + create_temp(roop.globals.target_path) + update_status('Extracting frames...') + extract_frames(roop.globals.target_path) + frame_paths = get_temp_frames_paths(roop.globals.target_path) + conditional_process_video(roop.globals.source_path, frame_paths) # prevent memory leak using ffmpeg with cuda - if args.gpu_vendor == 'nvidia': + if roop.globals.gpu_vendor == 'nvidia': torch.cuda.empty_cache() if roop.globals.keep_fps: - status('Detecting fps...') - fps = detect_fps(args.source_path) - status(f'Creating video with {fps} fps...') - create_video(args.target_path, fps) + update_status('Detecting fps...') + fps = detect_fps(roop.globals.source_path) + update_status(f'Creating video with {fps} fps...') + create_video(roop.globals.target_path, fps) else: - status('Creating video with 30 fps...') - create_video(args.target_path, 30) + update_status('Creating video with 30 fps...') + create_video(roop.globals.target_path, 30) if roop.globals.keep_audio: if roop.globals.keep_fps: - status('Restoring audio...') + update_status('Restoring audio...') else: - status('Restoring audio might cause issues as fps are not kept...') - restore_audio(args.target_path, args.output_path) + update_status('Restoring audio might cause issues as fps are not kept...') + restore_audio(roop.globals.target_path, roop.globals.output_path) else: - move_temp(args.target_path, args.output_path) - clean_temp(args.target_path) - if is_video(args.target_path): - status('Swapping to video succeed!') + move_temp(roop.globals.target_path, roop.globals.output_path) + clean_temp(roop.globals.target_path) + if is_video(roop.globals.target_path): + update_status('Swapping to video succeed!') else: - status('Swapping to video failed!') - - -def select_face_handler(path: str): - args.source_path = path - - -def select_target_handler(path: str): - args.target_path = path - return preview_video(args.target_path) - - -def toggle_all_faces_handler(value: int): - roop.globals.all_faces = True if value == 1 else False - - -def toggle_fps_limit_handler(value: int): - args.keep_fps = int(value != 1) - - -def toggle_keep_frames_handler(value: int): - args.keep_frames = value - - -def save_file_handler(path: str): - args.output_path = path - - -def create_test_preview(frame_number): - return process_faces( - get_face_single(cv2.imread(args.source_path)), - get_video_frame(args.target_path, frame_number) - ) + update_status('Swapping to video failed!') def destroy() -> None: - clean_temp(args.target_path) + if roop.globals.target_path: + clean_temp(roop.globals.target_path) quit() def run() -> None: - global all_faces, keep_frames, limit_fps - handle_parse() + parse_args() pre_check() limit_resources() if roop.globals.headless: start() else: - window = ui.init( - { - 'all_faces': args.all_faces, - 'keep_fps': args.keep_fps, - 'keep_frames': args.keep_frames - }, - select_face_handler, - select_target_handler, - toggle_all_faces_handler, - toggle_fps_limit_handler, - toggle_keep_frames_handler, - save_file_handler, - start, - get_video_frame, - create_test_preview - ) + window = ui.init(start) window.mainloop() diff --git a/roop/globals.py b/roop/globals.py index 100c193..c872571 100644 --- a/roop/globals.py +++ b/roop/globals.py @@ -1,5 +1,8 @@ import onnxruntime +source_path = None +target_path = None +output_path = None keep_fps = None keep_audio = None keep_frames = None @@ -7,6 +10,7 @@ all_faces = None cpu_cores = None gpu_threads = None gpu_vendor = None +max_memory = None headless = None log_level = 'error' providers = onnxruntime.get_available_providers() diff --git a/roop/ui.py b/roop/ui.py index bbca8bf..b83678e 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -1,11 +1,16 @@ import tkinter as tk from typing import Any, Callable, Tuple + +import cv2 from PIL import Image, ImageTk, ImageOps import webbrowser from tkinter import filedialog from tkinter.filedialog import asksaveasfilename import threading +import roop.globals +from roop.analyser import get_face_single +from roop.swapper import process_faces from roop.utilities import is_image max_preview_size = 800 @@ -213,23 +218,12 @@ def preview_target(frame): target_label.image = photo_img -def update_status_label(value): +def update_status(value): status_label["text"] = value window.update() -def init( - initial_values: dict, - select_face_handler: Callable[[str], None], - select_target_handler: Callable[[str], Tuple[int, Any]], - toggle_all_faces_handler: Callable[[int], None], - toggle_fps_limit_handler: Callable[[int], None], - toggle_keep_frames_handler: Callable[[int], None], - save_file_handler: Callable[[str], None], - start: Callable[[], None], - get_video_frame: Callable[[str, int], None], - create_test_preview: Callable[[int], Any], -): +def init(start: Callable[[], None]): global window, preview, preview_visible, face_label, target_label, status_label window = tk.Tk() @@ -274,22 +268,23 @@ def init( target_button.place(x=360,y=320,width=180,height=80) # All faces checkbox - all_faces = tk.IntVar(None, initial_values['all_faces']) + all_faces = tk.IntVar(None, roop.globals.all_faces) all_faces_checkbox = create_check(window, "Process all faces in frame", all_faces, toggle_all_faces(toggle_all_faces_handler, all_faces)) all_faces_checkbox.place(x=60,y=500,width=240,height=31) # FPS limit checkbox - limit_fps = tk.IntVar(None, not initial_values['keep_fps']) + limit_fps = tk.IntVar(None, not roop.globals.keep_fps) fps_checkbox = create_check(window, "Limit FPS to 30", limit_fps, toggle_fps_limit(toggle_fps_limit_handler, limit_fps)) fps_checkbox.place(x=60,y=475,width=240,height=31) # Keep frames checkbox - keep_frames = tk.IntVar(None, initial_values['keep_frames']) + keep_frames = tk.IntVar(None, roop.globals.keep_frames) frames_checkbox = create_check(window, "Keep frames dir", keep_frames, toggle_keep_frames(toggle_keep_frames_handler, keep_frames)) frames_checkbox.place(x=60,y=450,width=240,height=31) # Start button - start_button = create_button(window, "Start", lambda: [save_file(save_file_handler, target_path.get()), preview_thread(lambda: start(update_preview))]) + #start_button = create_button(window, "Start", lambda: [save_file(save_file_handler, target_path.get()), preview_thread(lambda: start(update_preview))]) + start_button = create_button(window, "Start", lambda: [save_file(save_file_handler, target_path.get()), start]) start_button.place(x=170,y=560,width=120,height=49) # Preview button @@ -301,3 +296,63 @@ def init( status_label.place(x=10,y=640,width=580,height=30) return window + + +def get_video_frame(video_path, frame_number = 1): + cap = cv2.VideoCapture(video_path) + amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) + cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1)) + if not cap.isOpened(): + update_status('Error opening video file') + return + ret, frame = cap.read() + if ret: + return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + cap.release() + + +def preview_video(video_path): + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + update_status('Error opening video file') + return 0 + amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) + ret, frame = cap.read() + if ret: + frame = get_video_frame(video_path) + + cap.release() + return (amount_of_frames, frame) + + +def select_face_handler(path: str): + roop.globals.source_path = path + + +def select_target_handler(target_path: str) -> None: + roop.globals.target_path = target_path + return preview_video(roop.globals.target_path) + + +def toggle_all_faces_handler(value: int): + roop.globals.all_faces = True if value == 1 else False + + +def toggle_fps_limit_handler(value: int): + roop.globals.keep_fps = int(value != 1) + + +def toggle_keep_frames_handler(value: int): + roop.globals.keep_frames = value + + +def save_file_handler(path: str): + roop.globals.output_path = path + + +def create_test_preview(frame_number): + return process_faces( + get_face_single(cv2.imread(roop.globals.source_path)), + get_video_frame(roop.globals.target_path, frame_number) + ) + From f77df69553fc3a04a31593dacb2dfd7f78c33809 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 7 Jun 2023 00:00:26 +0200 Subject: [PATCH 26/40] Refactor UI --- roop/analyser.py | 4 +- roop/core.py | 12 +- roop/globals.py | 2 +- roop/swapper.py | 16 +- roop/ui.py | 441 +++++++++++++--------------------------------- roop/utilities.py | 15 +- 6 files changed, 145 insertions(+), 345 deletions(-) diff --git a/roop/analyser.py b/roop/analyser.py index c2899e7..f8ec369 100644 --- a/roop/analyser.py +++ b/roop/analyser.py @@ -13,7 +13,7 @@ def get_face_analyser() -> Any: return FACE_ANALYSER -def get_face_single(image_data) -> Any: +def get_one_face(image_data) -> Any: face = get_face_analyser().get(image_data) try: return min(face, key=lambda x: x.bbox[0]) @@ -21,7 +21,7 @@ def get_face_single(image_data) -> Any: return None -def get_face_many(image_data) -> Any: +def get_many_faces(image_data) -> Any: try: return get_face_analyser().get(image_data) except IndexError: diff --git a/roop/core.py b/roop/core.py index 48bd2d7..845ea45 100755 --- a/roop/core.py +++ b/roop/core.py @@ -24,7 +24,7 @@ import roop.globals import roop.ui as ui from roop.swapper import process_video, process_img from roop.utilities import has_image_extention, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frames_paths, restore_audio, create_temp, move_temp, clean_temp -from roop.analyser import get_face_single +from roop.analyser import get_one_face if 'ROCMExecutionProvider' in roop.globals.providers: del torch @@ -41,11 +41,11 @@ def parse_args() -> None: parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) parser.add_argument('--keep-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True) 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('--many-faces', help='swap every face in the frame', dest='many_faces', action='store_true', default=False) parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) - parser.add_argument('--gpu-vendor', help='select your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia']) + parser.add_argument('--gpu-vendor', help='select your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'nvidia']) args = parser.parse_known_args()[0] @@ -56,7 +56,7 @@ def parse_args() -> None: roop.globals.keep_fps = args.keep_fps roop.globals.keep_audio = args.keep_audio roop.globals.keep_frames = args.keep_frames - roop.globals.all_faces = args.all_faces + roop.globals.many_faces = args.many_faces if args.cpu_cores: roop.globals.cpu_cores = int(args.cpu_cores) @@ -154,7 +154,7 @@ def start() -> None: elif not roop.globals.target_path or not os.path.isfile(roop.globals.target_path): update_status('Please select a video/image target!') return - test_face = get_face_single(cv2.imread(roop.globals.source_path)) + test_face = get_one_face(cv2.imread(roop.globals.source_path)) if not test_face: update_status('No face detected in source image. Please try with another one!') return @@ -217,5 +217,5 @@ def run() -> None: if roop.globals.headless: start() else: - window = ui.init(start) + window = ui.init(start, destroy) window.mainloop() diff --git a/roop/globals.py b/roop/globals.py index c872571..6680e63 100644 --- a/roop/globals.py +++ b/roop/globals.py @@ -6,7 +6,7 @@ output_path = None keep_fps = None keep_audio = None keep_frames = None -all_faces = None +many_faces = None cpu_cores = None gpu_threads = None gpu_vendor = None diff --git a/roop/swapper.py b/roop/swapper.py index 5f2eb1f..d001055 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -5,7 +5,7 @@ import cv2 import insightface import threading import roop.globals -from roop.analyser import get_face_single, get_face_many +from roop.analyser import get_one_face, get_many_faces FACE_SWAPPER = None THREAD_LOCK = threading.Lock() @@ -27,20 +27,20 @@ def swap_face_in_frame(source_face, target_face, frame): def process_faces(source_face, target_frame): - if roop.globals.all_faces: - many_faces = get_face_many(target_frame) + if roop.globals.many_faces: + many_faces = get_many_faces(target_frame) if many_faces: for face in many_faces: target_frame = swap_face_in_frame(source_face, face, target_frame) else: - face = get_face_single(target_frame) + face = get_one_face(target_frame) if face: target_frame = swap_face_in_frame(source_face, face, target_frame) return target_frame def process_frames(source_img, frame_paths, progress=None): - source_face = get_face_single(cv2.imread(source_img)) + source_face = get_one_face(cv2.imread(source_img)) for frame_path in frame_paths: frame = cv2.imread(frame_path) try: @@ -77,9 +77,9 @@ def multi_process_frame(source_img, frame_paths, progress): def process_img(source_img, target_path, output_file): frame = cv2.imread(target_path) - face = get_face_single(frame) - source_face = get_face_single(cv2.imread(source_img)) - result = get_face_swapper().get(frame, face, source_face, paste_back=True) + target_frame = get_one_face(frame) + source_face = get_one_face(cv2.imread(source_img)) + result = get_face_swapper().get(frame, target_frame, source_face, paste_back=True) cv2.imwrite(output_file, result) diff --git a/roop/ui.py b/roop/ui.py index b83678e..5c84d64 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -1,358 +1,157 @@ import tkinter as tk -from typing import Any, Callable, Tuple +from tkinter import filedialog +from typing import Callable, Any import cv2 from PIL import Image, ImageTk, ImageOps -import webbrowser -from tkinter import filedialog -from tkinter.filedialog import asksaveasfilename -import threading import roop.globals -from roop.analyser import get_face_single -from roop.swapper import process_faces from roop.utilities import is_image -max_preview_size = 800 +PRIMARY_COLOR = '#2d3436' +SECONDARY_COLOR = '#74b9ff' +TERTIARY_COLOR = '#f1c40f' +ACCENT_COLOR = '#2ecc71' +WINDOW_HEIGHT = 700 +WINDOW_WIDTH = 600 +MAX_PREVIEW_SIZE = 800 -def create_preview(parent): - global preview_image_frame, preview_frame_slider, test_button +def init(start: Callable, destroy: Callable): + global WINDOW, source_label, target_label, status_label - preview_window = tk.Toplevel(parent) - # Override close button - preview_window.protocol("WM_DELETE_WINDOW", hide_preview) - preview_window.withdraw() - preview_window.title("Preview") - preview_window.configure(bg="red") - preview_window.resizable(width=False, height=False) + WINDOW = tk.Tk() + WINDOW.minsize(WINDOW_WIDTH, WINDOW_HEIGHT) + WINDOW.title('roop') + WINDOW.configure(bg=PRIMARY_COLOR) + WINDOW.option_add('*Font', ('Arial', 11)) - frame = tk.Frame(preview_window, background="#2d3436") - frame.pack(fill='both', side='left', expand='True') - - # Preview image - preview_image_frame = tk.Label(frame) - preview_image_frame.pack(side='top') + source_label = tk.Label(bg=PRIMARY_COLOR) + source_label.place(relx=0.1, rely=0.1, relwidth=0.3, relheight=0.25) - # Bottom frame - buttons_frame = tk.Frame(frame, background="#2d3436") - buttons_frame.pack(fill='both', side='bottom') + target_label = tk.Label(bg=PRIMARY_COLOR) + target_label.place(relx=0.6, rely=0.1, relwidth=0.3, relheight=0.25) - current_frame = tk.IntVar() - preview_frame_slider = tk.Scale( - buttons_frame, - from_=0, - to=0, - orient='horizontal', - variable=current_frame - ) - preview_frame_slider.pack(fill='both', side='left', expand='True') + source_button = create_primary_button(WINDOW, 'Select a face', lambda: select_source_path()) + source_button.place(relx=0.1, rely=0.4, relwidth=0.3, relheight=0.1) - test_button = tk.Button(buttons_frame, text="Test", bg="#f1c40f", relief="flat", width=15, borderwidth=0, highlightthickness=0) - test_button.pack(side='right', fill='y') - return preview_window + target_button = create_primary_button(WINDOW, 'Select a target', lambda: select_target_path()) + target_button.place(relx=0.6, rely=0.4, relwidth=0.3, relheight=0.1) + + keep_fps_value = tk.BooleanVar(value=roop.globals.keep_fps) + keep_fps_checkbox = create_checkbox(WINDOW, 'Limit to 30 fps', keep_fps_value, lambda: setattr(roop.globals, 'keep_fps', not roop.globals.keep_fps)) + keep_fps_checkbox.place(relx=0.1, rely=0.6) + + keep_frames_value = tk.BooleanVar(value=roop.globals.keep_frames) + keep_frames_checkbox = create_checkbox(WINDOW, 'Keep frames dir', keep_frames_value, lambda: setattr(roop.globals, 'keep_frames', keep_frames_value.get())) + keep_frames_checkbox.place(relx=0.1, rely=0.65) + + keep_audio_value = tk.BooleanVar(value=roop.globals.keep_audio) + keep_audio_checkbox = create_checkbox(WINDOW, 'Keep original audio', keep_frames_value, lambda: setattr(roop.globals, 'keep_audio', keep_audio_value.get())) + keep_audio_checkbox.place(relx=0.6, rely=0.6) + + many_faces_value = tk.BooleanVar(value=roop.globals.many_faces) + many_faces_checkbox = create_checkbox(WINDOW, 'Replace all faces', many_faces_value, lambda: setattr(roop.globals, 'many_faces', keep_audio_value.get())) + many_faces_checkbox.place(relx=0.6, rely=0.65) + + start_button = create_secondary_button(WINDOW, 'Start', lambda: select_output_path(start)) + start_button.place(relx=0.15, rely=0.75, relwidth=0.2, relheight=0.05) + + stop_button = create_secondary_button(WINDOW, 'Destroy', lambda: destroy()) + stop_button.place(relx=0.4, rely=0.75, relwidth=0.2, relheight=0.05) + + preview_button = create_secondary_button(WINDOW, 'Preview', lambda: None) + preview_button.place(relx=0.65, rely=0.75, relwidth=0.2, relheight=0.05) + preview_button.config(state='disabled') + + status_label = tk.Label(WINDOW, justify='center', text='Status: UI under heavy development, more features will soon be (re)added', fg=ACCENT_COLOR, bg=PRIMARY_COLOR) + status_label.place(relx=0.1, rely=0.9) + + return WINDOW -def show_preview(): - preview.deiconify() - preview_visible.set(True) - - -def hide_preview(): - preview.withdraw() - preview_visible.set(False) - - -def set_preview_handler(test_handler): - test_button.config(command = test_handler) - - -def init_slider(frames_count, change_handler): - preview_frame_slider.configure(to=frames_count, command=lambda value: change_handler(preview_frame_slider.get())) - preview_frame_slider.set(0) - - -def update_preview(frame): - img = Image.fromarray(frame) - img = ImageOps.contain(img, (max_preview_size, max_preview_size), Image.LANCZOS) - photo_img = ImageTk.PhotoImage(img) - preview_image_frame.configure(image=photo_img) - preview_image_frame.image = photo_img - - -def select_face(select_face_handler: Callable[[str], None]): - if select_face_handler: - path = filedialog.askopenfilename(title="Select a face") - preview_face(path) - return select_face_handler(path) - return None - - -def update_slider_handler(get_video_frame, video_path): - return lambda frame_number: update_preview(get_video_frame(video_path, frame_number)) - - -def test_preview(create_test_preview): - frame = create_test_preview(preview_frame_slider.get()) - update_preview(frame) - - -def update_slider(get_video_frame, create_test_preview, video_path, frames_amount): - init_slider(frames_amount, update_slider_handler(get_video_frame, video_path)) - set_preview_handler(lambda: preview_thread(lambda: test_preview(create_test_preview))) - - -def analyze_target(select_target_handler: Callable[[str], Tuple[int, Any]], target_path: tk.StringVar, frames_amount: tk.IntVar): - path = filedialog.askopenfilename(title="Select a target") - target_path.set(path) - amount, frame = select_target_handler(path) - frames_amount.set(amount) - preview_target(frame) - update_preview(frame) - - -def select_target(select_target_handler: Callable[[str], Tuple[int, Any]], target_path: tk.StringVar, frames_amount: tk.IntVar): - if select_target_handler: - analyze_target(select_target_handler, target_path, frames_amount) - - -def save_file(save_file_handler: Callable[[str], None], target_path: str): - filename, ext = 'output.mp4', '.mp4' - - if is_image(target_path): - filename, ext = 'output.png', '.png' - - if save_file_handler: - return save_file_handler(asksaveasfilename(initialfile=filename, defaultextension=ext, filetypes=[("All Files","*.*"),("Videos","*.mp4")])) - return None - - -def toggle_all_faces(toggle_all_faces_handler: Callable[[int], None], variable: tk.IntVar): - if toggle_all_faces_handler: - return lambda: toggle_all_faces_handler(variable.get()) - return None - - -def toggle_fps_limit(toggle_all_faces_handler: Callable[[int], None], variable: tk.IntVar): - if toggle_all_faces_handler: - return lambda: toggle_all_faces_handler(variable.get()) - return None - - -def toggle_keep_frames(toggle_keep_frames_handler: Callable[[int], None], variable: tk.IntVar): - if toggle_keep_frames_handler: - return lambda: toggle_keep_frames_handler(variable.get()) - return None - - -def create_button(parent, text, command): +def create_primary_button(parent: Any, text: str, command: Callable) -> tk.Button: return tk.Button( - parent, - text=text, + parent, + text=text, command=command, - bg="#f1c40f", - relief="flat", - borderwidth=0, + bg=PRIMARY_COLOR, + fg=SECONDARY_COLOR, + relief='flat', + highlightthickness=4, + highlightbackground=SECONDARY_COLOR, + activebackground=SECONDARY_COLOR, + borderwidth=4 + ) + + +def create_secondary_button(parent: Any, text: str, command: Callable) -> tk.Button: + return tk.Button( + parent, + text=text, + command=command, + bg=TERTIARY_COLOR, + relief='flat', + borderwidth=0, highlightthickness=0 ) -def create_background_button(parent, text, command): - button = create_button(parent, text, command) - button.configure( - bg="#2d3436", - fg="#74b9ff", - highlightthickness=4, - highlightbackground="#74b9ff", - activebackground="#74b9ff", - borderwidth=4 - ) - return button - - -def create_check(parent, text, variable, command): +def create_checkbox(parent: Any, text: str, variable: tk.BooleanVar, command: Callable) -> tk.Checkbutton: return tk.Checkbutton( - parent, - anchor="w", - relief="groove", - activebackground="#2d3436", - activeforeground="#74b9ff", - selectcolor="black", - text=text, - fg="#dfe6e9", - borderwidth=0, - highlightthickness=0, - bg="#2d3436", - variable=variable, - command=command + parent, + text=text, + variable=variable, + command=command, + relief='flat', + bg=PRIMARY_COLOR, + activebackground=PRIMARY_COLOR, + activeforeground=SECONDARY_COLOR, + selectcolor=PRIMARY_COLOR, + fg=SECONDARY_COLOR, + borderwidth=0, + highlightthickness=0 ) -def preview_thread(thread_function): - threading.Thread(target=thread_function).start() +def update_status(text: str) -> None: + status_label['text'] = text + WINDOW.update() -def open_preview_window(get_video_frame, target_path): - if preview_visible.get(): - hide_preview() +def select_source_path(): + path = filedialog.askopenfilename(title='Select a face') + if is_image(path): + roop.globals.source_path = path + image = render_frame_image(roop.globals.source_path) + source_label.configure(image=image) + source_label.image = image else: - show_preview() - if target_path: - frame = get_video_frame(target_path) - update_preview(frame) + roop.globals.source_path = None + source_label.configure(image=None) + source_label.image = None -def preview_face(path): - img = Image.open(path) - img = ImageOps.fit(img, (180, 180), Image.LANCZOS) - photo_img = ImageTk.PhotoImage(img) - face_label.configure(image=photo_img) - face_label.image = photo_img +def select_target_path(): + path = filedialog.askopenfilename(title='Select a target') + if is_image(path): + roop.globals.target_path = path + image = render_frame_image(roop.globals.target_path) + target_label.configure(image=image) + target_label.image = image + else: + roop.globals.target_path = None + target_label.configure(image=None) + target_label.image = None -def preview_target(frame): - img = Image.fromarray(frame) - img = ImageOps.fit(img, (180, 180), Image.LANCZOS) - photo_img = ImageTk.PhotoImage(img) - target_label.configure(image=photo_img) - target_label.image = photo_img +def select_output_path(start): + roop.globals.output_path = filedialog.askdirectory(title='Select a target') + start() -def update_status(value): - status_label["text"] = value - window.update() - - -def init(start: Callable[[], None]): - global window, preview, preview_visible, face_label, target_label, status_label - - window = tk.Tk() - window.geometry("600x700") - window.title("roop") - window.configure(bg="#2d3436") - window.resizable(width=False, height=False) - - preview_visible = tk.BooleanVar(window, False) - target_path = tk.StringVar() - frames_amount = tk.IntVar() - - # Preview window - preview = create_preview(window) - - # Contact information - support_link = tk.Label(window, text="Donate to project <3", fg="#fd79a8", bg="#2d3436", cursor="hand2", font=("Arial", 8)) - support_link.place(x=180,y=20,width=250,height=30) - support_link.bind("", lambda e: webbrowser.open("https://github.com/sponsors/s0md3v")) - - left_frame = tk.Frame(window) - left_frame.place(x=60, y=100, width=180, height=180) - face_label = tk.Label(left_frame) - face_label.pack(fill='both', side='top', expand=True) - - right_frame = tk.Frame(window) - right_frame.place(x=360, y=100, width=180, height=180) - target_label = tk.Label(right_frame) - target_label.pack(fill='both', side='top', expand=True) - - # Select a face button - face_button = create_background_button(window, "Select a face", lambda: [ - select_face(select_face_handler) - ]) - face_button.place(x=60,y=320,width=180,height=80) - - # Select a target button - target_button = create_background_button(window, "Select a target", lambda: [ - select_target(select_target_handler, target_path, frames_amount), - update_slider(get_video_frame, create_test_preview, target_path.get(), frames_amount.get()) - ]) - target_button.place(x=360,y=320,width=180,height=80) - - # All faces checkbox - all_faces = tk.IntVar(None, roop.globals.all_faces) - all_faces_checkbox = create_check(window, "Process all faces in frame", all_faces, toggle_all_faces(toggle_all_faces_handler, all_faces)) - all_faces_checkbox.place(x=60,y=500,width=240,height=31) - - # FPS limit checkbox - limit_fps = tk.IntVar(None, not roop.globals.keep_fps) - fps_checkbox = create_check(window, "Limit FPS to 30", limit_fps, toggle_fps_limit(toggle_fps_limit_handler, limit_fps)) - fps_checkbox.place(x=60,y=475,width=240,height=31) - - # Keep frames checkbox - keep_frames = tk.IntVar(None, roop.globals.keep_frames) - frames_checkbox = create_check(window, "Keep frames dir", keep_frames, toggle_keep_frames(toggle_keep_frames_handler, keep_frames)) - frames_checkbox.place(x=60,y=450,width=240,height=31) - - # Start button - #start_button = create_button(window, "Start", lambda: [save_file(save_file_handler, target_path.get()), preview_thread(lambda: start(update_preview))]) - start_button = create_button(window, "Start", lambda: [save_file(save_file_handler, target_path.get()), start]) - start_button.place(x=170,y=560,width=120,height=49) - - # Preview button - preview_button = create_button(window, "Preview", lambda: open_preview_window(get_video_frame, target_path.get())) - preview_button.place(x=310,y=560,width=120,height=49) - - # Status label - status_label = tk.Label(window, width=580, justify="center", text="Status: waiting for input...", fg="#2ecc71", bg="#2d3436") - status_label.place(x=10,y=640,width=580,height=30) - - return window - - -def get_video_frame(video_path, frame_number = 1): - cap = cv2.VideoCapture(video_path) - amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) - cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1)) - if not cap.isOpened(): - update_status('Error opening video file') - return - ret, frame = cap.read() - if ret: - return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - cap.release() - - -def preview_video(video_path): - cap = cv2.VideoCapture(video_path) - if not cap.isOpened(): - update_status('Error opening video file') - return 0 - amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) - ret, frame = cap.read() - if ret: - frame = get_video_frame(video_path) - - cap.release() - return (amount_of_frames, frame) - - -def select_face_handler(path: str): - roop.globals.source_path = path - - -def select_target_handler(target_path: str) -> None: - roop.globals.target_path = target_path - return preview_video(roop.globals.target_path) - - -def toggle_all_faces_handler(value: int): - roop.globals.all_faces = True if value == 1 else False - - -def toggle_fps_limit_handler(value: int): - roop.globals.keep_fps = int(value != 1) - - -def toggle_keep_frames_handler(value: int): - roop.globals.keep_frames = value - - -def save_file_handler(path: str): - roop.globals.output_path = path - - -def create_test_preview(frame_number): - return process_faces( - get_face_single(cv2.imread(roop.globals.source_path)), - get_video_frame(roop.globals.target_path, frame_number) - ) +def render_frame_image(image_path: str) -> ImageTk.PhotoImage: + image = Image.open(image_path) + image = ImageOps.fit(image, (200, 200), Image.LANCZOS) + return ImageTk.PhotoImage(image) diff --git a/roop/utilities.py b/roop/utilities.py index 487cbb7..38cc888 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -70,13 +70,13 @@ def clean_temp(target_path: str) -> None: def has_image_extention(image_path: str) -> bool: - return image_path.lower().endswith(('png', 'jpg', 'jpeg', 'bmp')) + return image_path.lower().endswith(('png', 'jpg', 'jpeg')) -def is_image(path: str) -> bool: - if os.path.isfile(path): +def is_image(image_path: str) -> bool: + if image_path and os.path.isfile(image_path): try: - image = Image.open(path) + image = Image.open(image_path) image.verify() return True except Exception: @@ -84,10 +84,11 @@ def is_image(path: str) -> bool: return False -def is_video(path: str) -> bool: +def is_video(video_path: str) -> bool: try: - run_ffmpeg(['-v', 'error', '-i', path, '-f', 'null', '-']) - return True + if video_path and os.path.isfile(video_path): + run_ffmpeg(['-v', 'error', '-i', video_path, '-f', 'null', '-']) + return True except subprocess.CalledProcessError: pass return False From e555d98cd8044919c314e46ad6ff46baa2fc6939 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 7 Jun 2023 00:48:17 +0200 Subject: [PATCH 27/40] Introduce render_video_preview() --- roop/core.py | 8 ++++---- roop/ui.py | 46 +++++++++++++++++++++++++++++++++------------- roop/utilities.py | 19 ++++++++++++------- 3 files changed, 49 insertions(+), 24 deletions(-) diff --git a/roop/core.py b/roop/core.py index 845ea45..191791b 100755 --- a/roop/core.py +++ b/roop/core.py @@ -35,8 +35,8 @@ warnings.simplefilter(action='ignore', category=FutureWarning) def parse_args() -> None: signal.signal(signal.SIGINT, lambda signal_number, frame: destroy()) parser = argparse.ArgumentParser() - parser.add_argument('-f', '--face', help='use this face', dest='source_path') - parser.add_argument('-t', '--target', help='replace this face', dest='target_path') + parser.add_argument('-f', '--face', help='use a face image', dest='source_path') + parser.add_argument('-t', '--target', help='replace image or video with face', dest='target_path') 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-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True) @@ -149,10 +149,10 @@ def update_status(message: str): def start() -> None: if not roop.globals.source_path or not os.path.isfile(roop.globals.source_path): - update_status('Please select an image containing a face.') + update_status('Select an image that contains a face.') return elif not roop.globals.target_path or not os.path.isfile(roop.globals.target_path): - update_status('Please select a video/image target!') + update_status('Select an image or video target!') return test_face = get_one_face(cv2.imread(roop.globals.source_path)) if not test_face: diff --git a/roop/ui.py b/roop/ui.py index 5c84d64..b1a9a43 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -1,12 +1,12 @@ +import os import tkinter as tk from tkinter import filedialog from typing import Callable, Any import cv2 from PIL import Image, ImageTk, ImageOps - import roop.globals -from roop.utilities import is_image +from roop.utilities import is_image, is_video PRIMARY_COLOR = '#2d3436' SECONDARY_COLOR = '#74b9ff' @@ -120,10 +120,10 @@ def update_status(text: str) -> None: def select_source_path(): - path = filedialog.askopenfilename(title='Select a face') - if is_image(path): - roop.globals.source_path = path - image = render_frame_image(roop.globals.source_path) + source_path = filedialog.askopenfilename(title='Select an face image') + if is_image(source_path): + roop.globals.source_path = source_path + image = render_image_preview(roop.globals.source_path) source_label.configure(image=image) source_label.image = image else: @@ -133,12 +133,17 @@ def select_source_path(): def select_target_path(): - path = filedialog.askopenfilename(title='Select a target') - if is_image(path): - roop.globals.target_path = path - image = render_frame_image(roop.globals.target_path) + target_path = filedialog.askopenfilename(title='Select an image or video target') + if is_image(target_path): + roop.globals.target_path = target_path + image = render_image_preview(roop.globals.target_path) target_label.configure(image=image) target_label.image = image + elif is_video(target_path): + roop.globals.target_path = target_path + video_frame = render_video_preview(target_path) + target_label.configure(image=video_frame) + target_label.image = video_frame else: roop.globals.target_path = None target_label.configure(image=None) @@ -146,12 +151,27 @@ def select_target_path(): def select_output_path(start): - roop.globals.output_path = filedialog.askdirectory(title='Select a target') - start() + output_path = filedialog.askopenfilename(title='Save to output file') + if os.path.isfile(output_path): + roop.globals.output_path = output_path + start() -def render_frame_image(image_path: str) -> ImageTk.PhotoImage: +def render_image_preview(image_path: str) -> ImageTk.PhotoImage: image = Image.open(image_path) image = ImageOps.fit(image, (200, 200), Image.LANCZOS) return ImageTk.PhotoImage(image) + +def render_video_preview(target_path: str) -> ImageTk.PhotoImage: + capture = cv2.VideoCapture(target_path) + total_frames = capture.get(cv2.CAP_PROP_FRAME_COUNT) + capture.set(cv2.CAP_PROP_POS_FRAMES, total_frames / 2) + has_frame, frame = capture.read() + if has_frame: + image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + image = ImageOps.fit(image, (200, 200), Image.LANCZOS) + return ImageTk.PhotoImage(image) + capture.release() + cv2.destroyAllWindows() + diff --git a/roop/utilities.py b/roop/utilities.py index 38cc888..e05e323 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -5,6 +5,8 @@ import subprocess from pathlib import Path from typing import List, Any +import cv2 + import roop.globals from PIL import Image @@ -74,7 +76,7 @@ def has_image_extention(image_path: str) -> bool: def is_image(image_path: str) -> bool: - if image_path and os.path.isfile(image_path): + if os.path.isfile(image_path): try: image = Image.open(image_path) image.verify() @@ -85,10 +87,13 @@ def is_image(image_path: str) -> bool: def is_video(video_path: str) -> bool: - try: - if video_path and os.path.isfile(video_path): - run_ffmpeg(['-v', 'error', '-i', video_path, '-f', 'null', '-']) - return True - except subprocess.CalledProcessError: - pass + if os.path.isfile(video_path): + try: + capture = cv2.VideoCapture(video_path) + if capture.isOpened(): + is_video, _ = capture.read() + capture.release() + return is_video + except Exception: + pass return False From 80f3870228dca1947f9aff0e02798713e6ca26c0 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 7 Jun 2023 09:03:46 +0200 Subject: [PATCH 28/40] Add preview back part1 --- roop/ui.py | 105 ++++++++++++++++++++++++++++++++++++----------------- 1 file changed, 71 insertions(+), 34 deletions(-) diff --git a/roop/ui.py b/roop/ui.py index b1a9a43..ca9e78a 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -1,7 +1,7 @@ import os import tkinter as tk from tkinter import filedialog -from typing import Callable, Any +from typing import Callable, Any, Tuple import cv2 from PIL import Image, ImageTk, ImageOps @@ -14,60 +14,89 @@ TERTIARY_COLOR = '#f1c40f' ACCENT_COLOR = '#2ecc71' WINDOW_HEIGHT = 700 WINDOW_WIDTH = 600 -MAX_PREVIEW_SIZE = 800 +PREVIEW_HEIGHT = 700 +PREVIEW_WIDTH = 1200 -def init(start: Callable, destroy: Callable): - global WINDOW, source_label, target_label, status_label +def init(start: Callable, destroy: Callable) -> tk.Tk: + global ROOT, PREVIEW - WINDOW = tk.Tk() - WINDOW.minsize(WINDOW_WIDTH, WINDOW_HEIGHT) - WINDOW.title('roop') - WINDOW.configure(bg=PRIMARY_COLOR) - WINDOW.option_add('*Font', ('Arial', 11)) + ROOT = create_root(start, destroy) + PREVIEW = create_preview(ROOT) - source_label = tk.Label(bg=PRIMARY_COLOR) + return ROOT + + +def create_root(start: Callable, destroy: Callable) -> tk.Tk: + global source_label, target_label, status_label + + root = tk.Tk() + root.minsize(WINDOW_WIDTH, WINDOW_HEIGHT) + root.title('roop') + root.configure(bg=PRIMARY_COLOR) + root.option_add('*Font', ('Arial', 11)) + + source_label = tk.Label(root, bg=PRIMARY_COLOR) source_label.place(relx=0.1, rely=0.1, relwidth=0.3, relheight=0.25) - target_label = tk.Label(bg=PRIMARY_COLOR) + target_label = tk.Label(root, bg=PRIMARY_COLOR) target_label.place(relx=0.6, rely=0.1, relwidth=0.3, relheight=0.25) - source_button = create_primary_button(WINDOW, 'Select a face', lambda: select_source_path()) + source_button = create_primary_button(root, 'Select a face', lambda: select_source_path()) source_button.place(relx=0.1, rely=0.4, relwidth=0.3, relheight=0.1) - target_button = create_primary_button(WINDOW, 'Select a target', lambda: select_target_path()) + target_button = create_primary_button(root, 'Select a target', lambda: select_target_path()) target_button.place(relx=0.6, rely=0.4, relwidth=0.3, relheight=0.1) keep_fps_value = tk.BooleanVar(value=roop.globals.keep_fps) - keep_fps_checkbox = create_checkbox(WINDOW, 'Limit to 30 fps', keep_fps_value, lambda: setattr(roop.globals, 'keep_fps', not roop.globals.keep_fps)) + keep_fps_checkbox = create_checkbox(root, 'Limit to 30 fps', keep_fps_value, lambda: setattr(roop.globals, 'keep_fps', not roop.globals.keep_fps)) keep_fps_checkbox.place(relx=0.1, rely=0.6) keep_frames_value = tk.BooleanVar(value=roop.globals.keep_frames) - keep_frames_checkbox = create_checkbox(WINDOW, 'Keep frames dir', keep_frames_value, lambda: setattr(roop.globals, 'keep_frames', keep_frames_value.get())) + keep_frames_checkbox = create_checkbox(root, 'Keep frames dir', keep_frames_value, lambda: setattr(roop.globals, 'keep_frames', keep_frames_value.get())) keep_frames_checkbox.place(relx=0.1, rely=0.65) keep_audio_value = tk.BooleanVar(value=roop.globals.keep_audio) - keep_audio_checkbox = create_checkbox(WINDOW, 'Keep original audio', keep_frames_value, lambda: setattr(roop.globals, 'keep_audio', keep_audio_value.get())) + keep_audio_checkbox = create_checkbox(root, 'Keep original audio', keep_frames_value, lambda: setattr(roop.globals, 'keep_audio', keep_audio_value.get())) keep_audio_checkbox.place(relx=0.6, rely=0.6) many_faces_value = tk.BooleanVar(value=roop.globals.many_faces) - many_faces_checkbox = create_checkbox(WINDOW, 'Replace all faces', many_faces_value, lambda: setattr(roop.globals, 'many_faces', keep_audio_value.get())) + many_faces_checkbox = create_checkbox(root, 'Replace all faces', many_faces_value, lambda: setattr(roop.globals, 'many_faces', keep_audio_value.get())) many_faces_checkbox.place(relx=0.6, rely=0.65) - start_button = create_secondary_button(WINDOW, 'Start', lambda: select_output_path(start)) + start_button = create_secondary_button(root, 'Start', lambda: select_output_path(start)) start_button.place(relx=0.15, rely=0.75, relwidth=0.2, relheight=0.05) - stop_button = create_secondary_button(WINDOW, 'Destroy', lambda: destroy()) + stop_button = create_secondary_button(root, 'Destroy', lambda: destroy()) stop_button.place(relx=0.4, rely=0.75, relwidth=0.2, relheight=0.05) - preview_button = create_secondary_button(WINDOW, 'Preview', lambda: None) + preview_button = create_secondary_button(root, 'Preview', lambda: toggle_preview()) preview_button.place(relx=0.65, rely=0.75, relwidth=0.2, relheight=0.05) - preview_button.config(state='disabled') - status_label = tk.Label(WINDOW, justify='center', text='Status: UI under heavy development, more features will soon be (re)added', fg=ACCENT_COLOR, bg=PRIMARY_COLOR) + status_label = tk.Label(root, justify='center', text='Status: UI under heavy development, more features will soon be (re)added', fg=ACCENT_COLOR, bg=PRIMARY_COLOR) status_label.place(relx=0.1, rely=0.9) - return WINDOW + return root + + +def create_preview(parent) -> tk.Toplevel: + global preview_label + + preview = tk.Toplevel(parent) + preview.withdraw() + preview.title('Preview') + preview.configure(bg=PRIMARY_COLOR) + preview.option_add('*Font', ('Arial', 11)) + preview.minsize(PREVIEW_WIDTH, PREVIEW_HEIGHT) + + preview_label = tk.Label(preview, bg=PRIMARY_COLOR) + preview_label.pack(fill='both', expand=True) + + frame_value = tk.IntVar() + frame_slider = tk.Scale(preview, orient='horizontal', variable=frame_value) + frame_slider.pack(fill='x') + + return preview def create_primary_button(parent: Any, text: str, command: Callable) -> tk.Button: @@ -116,14 +145,14 @@ def create_checkbox(parent: Any, text: str, variable: tk.BooleanVar, command: Ca def update_status(text: str) -> None: status_label['text'] = text - WINDOW.update() + ROOT.update() def select_source_path(): source_path = filedialog.askopenfilename(title='Select an face image') if is_image(source_path): roop.globals.source_path = source_path - image = render_image_preview(roop.globals.source_path) + image = render_image_preview(roop.globals.source_path, (200, 200)) source_label.configure(image=image) source_label.image = image else: @@ -141,7 +170,7 @@ def select_target_path(): target_label.image = image elif is_video(target_path): roop.globals.target_path = target_path - video_frame = render_video_preview(target_path) + video_frame = render_video_preview(target_path, (200, 200)) target_label.configure(image=video_frame) target_label.image = video_frame else: @@ -157,21 +186,29 @@ def select_output_path(start): start() -def render_image_preview(image_path: str) -> ImageTk.PhotoImage: +def render_image_preview(image_path: str, dimensions: Tuple[int, int] = None) -> ImageTk.PhotoImage: image = Image.open(image_path) - image = ImageOps.fit(image, (200, 200), Image.LANCZOS) + if dimensions: + image = ImageOps.fit(image, dimensions, Image.LANCZOS) return ImageTk.PhotoImage(image) -def render_video_preview(target_path: str) -> ImageTk.PhotoImage: - capture = cv2.VideoCapture(target_path) - total_frames = capture.get(cv2.CAP_PROP_FRAME_COUNT) - capture.set(cv2.CAP_PROP_POS_FRAMES, total_frames / 2) +def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame: int = 1) -> ImageTk.PhotoImage: + capture = cv2.VideoCapture(video_path) + if frame: + capture.set(cv2.CAP_PROP_POS_FRAMES, frame) has_frame, frame = capture.read() if has_frame: - image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) - image = ImageOps.fit(image, (200, 200), Image.LANCZOS) + if dimensions: + image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + image = ImageOps.fit(image, dimensions, Image.LANCZOS) return ImageTk.PhotoImage(image) capture.release() cv2.destroyAllWindows() + +def toggle_preview(): + if PREVIEW.state() == 'normal': + PREVIEW.withdraw() + else: + PREVIEW.deiconify() From eeae9e46f492cbc15dc63db96e0ca972e53a56d4 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 7 Jun 2023 09:49:32 +0200 Subject: [PATCH 29/40] Add preview back part2, Introduce --video-quality for CLI --- roop/core.py | 2 ++ roop/globals.py | 1 + roop/ui.py | 10 ++++++---- roop/utilities.py | 6 +++--- 4 files changed, 12 insertions(+), 7 deletions(-) diff --git a/roop/core.py b/roop/core.py index 191791b..ab8dd07 100755 --- a/roop/core.py +++ b/roop/core.py @@ -42,6 +42,7 @@ def parse_args() -> None: parser.add_argument('--keep-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--many-faces', help='swap every face in the frame', dest='many_faces', action='store_true', default=False) + parser.add_argument('--video-quality', help='adjust video quality of output file', dest='video_quality', type=int, default=10) parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) @@ -57,6 +58,7 @@ def parse_args() -> None: roop.globals.keep_audio = args.keep_audio roop.globals.keep_frames = args.keep_frames roop.globals.many_faces = args.many_faces + roop.globals.video_quality = args.video_quality if args.cpu_cores: roop.globals.cpu_cores = int(args.cpu_cores) diff --git a/roop/globals.py b/roop/globals.py index 6680e63..f280566 100644 --- a/roop/globals.py +++ b/roop/globals.py @@ -7,6 +7,7 @@ keep_fps = None keep_audio = None keep_frames = None many_faces = None +video_quality = None cpu_cores = None gpu_threads = None gpu_vendor = None diff --git a/roop/ui.py b/roop/ui.py index ca9e78a..c0d68c8 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -35,6 +35,7 @@ def create_root(start: Callable, destroy: Callable) -> tk.Tk: root.title('roop') root.configure(bg=PRIMARY_COLOR) root.option_add('*Font', ('Arial', 11)) + root.protocol('WM_DELETE_WINDOW', lambda: destroy()) source_label = tk.Label(root, bg=PRIMARY_COLOR) source_label.place(relx=0.1, rely=0.1, relwidth=0.3, relheight=0.25) @@ -73,7 +74,7 @@ def create_root(start: Callable, destroy: Callable) -> tk.Tk: preview_button = create_secondary_button(root, 'Preview', lambda: toggle_preview()) preview_button.place(relx=0.65, rely=0.75, relwidth=0.2, relheight=0.05) - status_label = tk.Label(root, justify='center', text='Status: UI under heavy development, more features will soon be (re)added', fg=ACCENT_COLOR, bg=PRIMARY_COLOR) + status_label = tk.Label(root, justify='center', text='Status: None', fg=ACCENT_COLOR, bg=PRIMARY_COLOR) status_label.place(relx=0.1, rely=0.9) return root @@ -88,6 +89,7 @@ def create_preview(parent) -> tk.Toplevel: preview.configure(bg=PRIMARY_COLOR) preview.option_add('*Font', ('Arial', 11)) preview.minsize(PREVIEW_WIDTH, PREVIEW_HEIGHT) + preview.protocol('WM_DELETE_WINDOW', lambda: toggle_preview()) preview_label = tk.Label(preview, bg=PRIMARY_COLOR) preview_label.pack(fill='both', expand=True) @@ -180,8 +182,8 @@ def select_target_path(): def select_output_path(start): - output_path = filedialog.askopenfilename(title='Save to output file') - if os.path.isfile(output_path): + output_path = filedialog.asksaveasfilename(title='Save to output file', initialfile='output.mp4') + if output_path and os.path.isfile(output_path): roop.globals.output_path = output_path start() @@ -207,7 +209,7 @@ def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, fr cv2.destroyAllWindows() -def toggle_preview(): +def toggle_preview() -> None: if PREVIEW.state() == 'normal': PREVIEW.withdraw() else: diff --git a/roop/utilities.py b/roop/utilities.py index e05e323..27afbe6 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -35,7 +35,7 @@ def extract_frames(target_path: str) -> None: def create_video(target_path: str, fps: int) -> None: - run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', '7', '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) + run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', roop.globals.video_quality, '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) def restore_audio(target_path: str, output_path: str) -> None: @@ -76,7 +76,7 @@ def has_image_extention(image_path: str) -> bool: def is_image(image_path: str) -> bool: - if os.path.isfile(image_path): + if image_path and os.path.isfile(image_path): try: image = Image.open(image_path) image.verify() @@ -87,7 +87,7 @@ def is_image(image_path: str) -> bool: def is_video(video_path: str) -> bool: - if os.path.isfile(video_path): + if video_path and os.path.isfile(video_path): try: capture = cv2.VideoCapture(video_path) if capture.isOpened(): From 9c66fd9712685b4342946aab8aa4a12dcf0aa2bc Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 7 Jun 2023 15:40:31 +0200 Subject: [PATCH 30/40] Get the preview working --- roop/capturer.py | 12 ++++++++++++ roop/ui.py | 36 ++++++++++++++++++++++++++---------- roop/utilities.py | 7 ++++--- 3 files changed, 42 insertions(+), 13 deletions(-) create mode 100644 roop/capturer.py diff --git a/roop/capturer.py b/roop/capturer.py new file mode 100644 index 0000000..ecebf49 --- /dev/null +++ b/roop/capturer.py @@ -0,0 +1,12 @@ +import cv2 + + +def get_video_frame(video_path: str, frame_number: int = 1): + capture = cv2.VideoCapture(video_path) + frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT) + capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1)) + has_frame, frame = capture.read() + capture.release() + if has_frame: + return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + return None diff --git a/roop/ui.py b/roop/ui.py index c0d68c8..b7aa848 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -6,6 +6,9 @@ from typing import Callable, Any, Tuple import cv2 from PIL import Image, ImageTk, ImageOps import roop.globals +from roop.analyser import get_one_face +from roop.capturer import get_video_frame +from roop.swapper import process_faces from roop.utilities import is_image, is_video PRIMARY_COLOR = '#2d3436' @@ -81,7 +84,7 @@ def create_root(start: Callable, destroy: Callable) -> tk.Tk: def create_preview(parent) -> tk.Toplevel: - global preview_label + global preview_label, preview_scale preview = tk.Toplevel(parent) preview.withdraw() @@ -94,9 +97,8 @@ def create_preview(parent) -> tk.Toplevel: preview_label = tk.Label(preview, bg=PRIMARY_COLOR) preview_label.pack(fill='both', expand=True) - frame_value = tk.IntVar() - frame_slider = tk.Scale(preview, orient='horizontal', variable=frame_value) - frame_slider.pack(fill='x') + preview_scale = tk.Scale(preview, orient='horizontal', command=lambda frame_value: update_preview(int(frame_value))) + preview_scale.pack(fill='x') return preview @@ -183,7 +185,7 @@ def select_target_path(): def select_output_path(start): output_path = filedialog.asksaveasfilename(title='Save to output file', initialfile='output.mp4') - if output_path and os.path.isfile(output_path): + if output_path: roop.globals.output_path = output_path start() @@ -195,15 +197,15 @@ def render_image_preview(image_path: str, dimensions: Tuple[int, int] = None) -> return ImageTk.PhotoImage(image) -def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame: int = 1) -> ImageTk.PhotoImage: +def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame_number: int = 1) -> ImageTk.PhotoImage: capture = cv2.VideoCapture(video_path) - if frame: - capture.set(cv2.CAP_PROP_POS_FRAMES, frame) + if frame_number: + capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number) has_frame, frame = capture.read() if has_frame: + image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) if dimensions: - image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) - image = ImageOps.fit(image, dimensions, Image.LANCZOS) + image = ImageOps.fit(image, dimensions, Image.LANCZOS) return ImageTk.PhotoImage(image) capture.release() cv2.destroyAllWindows() @@ -213,4 +215,18 @@ def toggle_preview() -> None: if PREVIEW.state() == 'normal': PREVIEW.withdraw() else: + update_preview(1) PREVIEW.deiconify() + + +def update_preview(frame_number: int) -> None: + if roop.globals.source_path and roop.globals.target_path and frame_number: + video_frame = process_faces( + get_one_face(cv2.imread(roop.globals.source_path)), + get_video_frame(roop.globals.target_path, frame_number) + ) + img = Image.fromarray(video_frame) + img = ImageOps.contain(img, (PREVIEW_WIDTH, PREVIEW_HEIGHT), Image.LANCZOS) + img = ImageTk.PhotoImage(img) + preview_label.configure(image=img) + preview_label.image = img diff --git a/roop/utilities.py b/roop/utilities.py index 27afbe6..8746a75 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -35,7 +35,7 @@ def extract_frames(target_path: str) -> None: def create_video(target_path: str, fps: int) -> None: - run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', roop.globals.video_quality, '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) + run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) def restore_audio(target_path: str, output_path: str) -> None: @@ -67,8 +67,9 @@ def move_temp(target_path: str, output_path: str) -> None: def clean_temp(target_path: str) -> None: - if not roop.globals.keep_frames: - shutil.rmtree(get_temp_directory_path(target_path)) + temp_directory_path = get_temp_directory_path(target_path) + if not roop.globals.keep_frames and os.path.isdir(temp_directory_path): + shutil.rmtree(temp_directory_path) def has_image_extention(image_path: str) -> bool: From e37fa75522bbcfa4b21ff9f49c09bc21fa38ba9d Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 7 Jun 2023 19:07:42 +0200 Subject: [PATCH 31/40] Couple if minor UI fixes --- roop/ui.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/roop/ui.py b/roop/ui.py index b7aa848..6c64b21 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -1,10 +1,10 @@ -import os import tkinter as tk from tkinter import filedialog from typing import Callable, Any, Tuple import cv2 from PIL import Image, ImageTk, ImageOps + import roop.globals from roop.analyser import get_one_face from roop.capturer import get_video_frame @@ -17,8 +17,8 @@ TERTIARY_COLOR = '#f1c40f' ACCENT_COLOR = '#2ecc71' WINDOW_HEIGHT = 700 WINDOW_WIDTH = 600 -PREVIEW_HEIGHT = 700 -PREVIEW_WIDTH = 1200 +PREVIEW_MAX_HEIGHT = 700 +PREVIEW_MAX_WIDTH = 1200 def init(start: Callable, destroy: Callable) -> tk.Tk: @@ -61,11 +61,11 @@ def create_root(start: Callable, destroy: Callable) -> tk.Tk: keep_frames_checkbox.place(relx=0.1, rely=0.65) keep_audio_value = tk.BooleanVar(value=roop.globals.keep_audio) - keep_audio_checkbox = create_checkbox(root, 'Keep original audio', keep_frames_value, lambda: setattr(roop.globals, 'keep_audio', keep_audio_value.get())) + keep_audio_checkbox = create_checkbox(root, 'Keep original audio', keep_audio_value, lambda: setattr(roop.globals, 'keep_audio', keep_audio_value.get())) keep_audio_checkbox.place(relx=0.6, rely=0.6) many_faces_value = tk.BooleanVar(value=roop.globals.many_faces) - many_faces_checkbox = create_checkbox(root, 'Replace all faces', many_faces_value, lambda: setattr(roop.globals, 'many_faces', keep_audio_value.get())) + many_faces_checkbox = create_checkbox(root, 'Replace all faces', many_faces_value, lambda: setattr(roop.globals, 'many_faces', many_faces_value.get())) many_faces_checkbox.place(relx=0.6, rely=0.65) start_button = create_secondary_button(root, 'Start', lambda: select_output_path(start)) @@ -91,8 +91,8 @@ def create_preview(parent) -> tk.Toplevel: preview.title('Preview') preview.configure(bg=PRIMARY_COLOR) preview.option_add('*Font', ('Arial', 11)) - preview.minsize(PREVIEW_WIDTH, PREVIEW_HEIGHT) preview.protocol('WM_DELETE_WINDOW', lambda: toggle_preview()) + preview.resizable(width=False, height=False) preview_label = tk.Label(preview, bg=PRIMARY_COLOR) preview_label.pack(fill='both', expand=True) @@ -184,7 +184,10 @@ def select_target_path(): def select_output_path(start): - output_path = filedialog.asksaveasfilename(title='Save to output file', initialfile='output.mp4') + if is_image(roop.globals.target_path): + output_path = filedialog.asksaveasfilename(title='Save image output', initialfile='output.png') + elif is_video(roop.globals.target_path): + output_path = filedialog.asksaveasfilename(title='Save video output', initialfile='output.mp4') if output_path: roop.globals.output_path = output_path start() @@ -225,8 +228,8 @@ def update_preview(frame_number: int) -> None: get_one_face(cv2.imread(roop.globals.source_path)), get_video_frame(roop.globals.target_path, frame_number) ) - img = Image.fromarray(video_frame) - img = ImageOps.contain(img, (PREVIEW_WIDTH, PREVIEW_HEIGHT), Image.LANCZOS) - img = ImageTk.PhotoImage(img) - preview_label.configure(image=img) - preview_label.image = img + image = Image.fromarray(video_frame) + image = ImageOps.contain(image, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), Image.LANCZOS) + image = ImageTk.PhotoImage(image) + preview_label.configure(image=image) + preview_label.image = image From 60128d3e96efd307dadd56a838a4a57283fc6fc1 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 7 Jun 2023 20:01:42 +0200 Subject: [PATCH 32/40] Add video encoder via CLI --- README.md | 19 ++++++++++++------- roop/core.py | 4 +++- roop/globals.py | 1 + roop/utilities.py | 2 +- 4 files changed, 17 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 9d19aa9..7de722c 100644 --- a/README.md +++ b/README.md @@ -34,23 +34,28 @@ Additional command line arguments are given below: ``` options: -h, --help show this help message and exit - -f SOURCE_IMG, --face SOURCE_IMG - use this face + -f SOURCE_PATH, --face SOURCE_PATH + use a face image -t TARGET_PATH, --target TARGET_PATH - replace this face - -o OUTPUT_FILE, --output OUTPUT_FILE + replace image or video with face + -o OUTPUT_PATH, --output OUTPUT_PATH save output to this file --keep-fps maintain original fps + --keep-audio maintain original audio --keep-frames keep frames directory - --all-faces swap all faces in frame + --many-faces swap every face in the frame + --video-encoder VIDEO_ENCODER + adjust output video encoder + --video-quality VIDEO_QUALITY + adjust output video quality --max-memory MAX_MEMORY maximum amount of RAM in GB to be used --cpu-cores CPU_CORES number of CPU cores to use --gpu-threads GPU_THREADS number of threads to be use for the GPU - --gpu-vendor {apple,amd,intel,nvidia} - choice your GPU vendor + --gpu-vendor {apple,amd,nvidia} + select your GPU vendor ``` Looking for a CLI mode? Using the -f/--face argument will make the program in cli mode. diff --git a/roop/core.py b/roop/core.py index ab8dd07..906151f 100755 --- a/roop/core.py +++ b/roop/core.py @@ -42,7 +42,8 @@ def parse_args() -> None: parser.add_argument('--keep-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--many-faces', help='swap every face in the frame', dest='many_faces', action='store_true', default=False) - parser.add_argument('--video-quality', help='adjust video quality of output file', dest='video_quality', type=int, default=10) + parser.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264') + parser.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=10) parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) @@ -58,6 +59,7 @@ def parse_args() -> None: roop.globals.keep_audio = args.keep_audio roop.globals.keep_frames = args.keep_frames roop.globals.many_faces = args.many_faces + roop.globals.video_encoder = args.video_encoder roop.globals.video_quality = args.video_quality if args.cpu_cores: diff --git a/roop/globals.py b/roop/globals.py index f280566..4ad9085 100644 --- a/roop/globals.py +++ b/roop/globals.py @@ -7,6 +7,7 @@ keep_fps = None keep_audio = None keep_frames = None many_faces = None +video_encoder = None video_quality = None cpu_cores = None gpu_threads = None diff --git a/roop/utilities.py b/roop/utilities.py index 8746a75..da4dd84 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -35,7 +35,7 @@ def extract_frames(target_path: str) -> None: def create_video(target_path: str, fps: int) -> None: - run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) + run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', roop.globals.video_encoder, '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) def restore_audio(target_path: str, output_path: str) -> None: From b1c2d414fff3f9c2327476292ad43390170a92af Mon Sep 17 00:00:00 2001 From: henryruhs Date: Thu, 8 Jun 2023 12:58:29 +0200 Subject: [PATCH 33/40] Change default video quality, Integrate recent directories for UI --- roop/core.py | 2 +- roop/ui.py | 19 +++++++++++++++---- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/roop/core.py b/roop/core.py index 906151f..1621210 100755 --- a/roop/core.py +++ b/roop/core.py @@ -43,7 +43,7 @@ def parse_args() -> None: parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--many-faces', help='swap every face in the frame', dest='many_faces', action='store_true', default=False) parser.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264') - parser.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=10) + parser.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18) parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) diff --git a/roop/ui.py b/roop/ui.py index 6c64b21..18d7ff9 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -1,3 +1,4 @@ +import os import tkinter as tk from tkinter import filedialog from typing import Callable, Any, Tuple @@ -19,6 +20,9 @@ WINDOW_HEIGHT = 700 WINDOW_WIDTH = 600 PREVIEW_MAX_HEIGHT = 700 PREVIEW_MAX_WIDTH = 1200 +RECENT_DIRECTORY_SOURCE = None +RECENT_DIRECTORY_TARGET = None +RECENT_DIRECTORY_OUTPUT = None def init(start: Callable, destroy: Callable) -> tk.Tk: @@ -153,9 +157,11 @@ def update_status(text: str) -> None: def select_source_path(): - source_path = filedialog.askopenfilename(title='Select an face image') + global RECENT_DIRECTORY_SOURCE + source_path = filedialog.askopenfilename(title='Select an face image', initialdir=RECENT_DIRECTORY_SOURCE) if is_image(source_path): roop.globals.source_path = source_path + RECENT_DIRECTORY_SOURCE = os.path.dirname(roop.globals.source_path) image = render_image_preview(roop.globals.source_path, (200, 200)) source_label.configure(image=image) source_label.image = image @@ -166,14 +172,17 @@ def select_source_path(): def select_target_path(): - target_path = filedialog.askopenfilename(title='Select an image or video target') + global RECENT_DIRECTORY_TARGET + target_path = filedialog.askopenfilename(title='Select an image or video target', initialdir=RECENT_DIRECTORY_TARGET) if is_image(target_path): roop.globals.target_path = target_path + RECENT_DIRECTORY_TARGET = os.path.dirname(roop.globals.target_path) image = render_image_preview(roop.globals.target_path) target_label.configure(image=image) target_label.image = image elif is_video(target_path): roop.globals.target_path = target_path + RECENT_DIRECTORY_TARGET = os.path.dirname(roop.globals.target_path) video_frame = render_video_preview(target_path, (200, 200)) target_label.configure(image=video_frame) target_label.image = video_frame @@ -184,12 +193,14 @@ def select_target_path(): def select_output_path(start): + global RECENT_DIRECTORY_OUTPUT if is_image(roop.globals.target_path): - output_path = filedialog.asksaveasfilename(title='Save image output', initialfile='output.png') + output_path = filedialog.asksaveasfilename(title='Save image output', initialfile='output.png', initialdir=RECENT_DIRECTORY_OUTPUT) elif is_video(roop.globals.target_path): - output_path = filedialog.asksaveasfilename(title='Save video output', initialfile='output.mp4') + output_path = filedialog.asksaveasfilename(title='Save video output', initialfile='output.mp4', initialdir=RECENT_DIRECTORY_OUTPUT) if output_path: roop.globals.output_path = output_path + RECENT_DIRECTORY_OUTPUT = os.path.dirname(roop.globals.output_path) start() From 703346da79f1c23abc738cc6a9f59eba9627ed8c Mon Sep 17 00:00:00 2001 From: henryruhs Date: Thu, 8 Jun 2023 17:27:20 +0200 Subject: [PATCH 34/40] Move temporary files to temp/{target-name} --- .gitignore | 2 +- roop/utilities.py | 38 ++++++++++++++++++++++++++------------ 2 files changed, 27 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index 09916c4..3e5e7f2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,3 @@ .idea +temp __pycache__ -*.onnx \ No newline at end of file diff --git a/roop/utilities.py b/roop/utilities.py index da4dd84..ce55969 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -3,12 +3,14 @@ import os import shutil import subprocess from pathlib import Path -from typing import List, Any - +from typing import List import cv2 +from PIL import Image import roop.globals -from PIL import Image + +TEMP_FILE = 'temp.mp4' +TEMP_DIRECTORY = 'temp' def run_ffmpeg(args: List) -> None: @@ -16,7 +18,7 @@ def run_ffmpeg(args: List) -> None: commands.extend(args) try: subprocess.check_output(commands, stderr=subprocess.STDOUT) - except Exception as exception: + except Exception: pass @@ -31,33 +33,41 @@ def detect_fps(source_path: str) -> int: def extract_frames(target_path: str) -> None: - run_ffmpeg(['-i', target_path, get_temp_directory_path(target_path) + os.sep + '%04d.png']) + temp_directory_path = get_temp_directory_path(target_path) + run_ffmpeg(['-i', target_path, os.path.join(temp_directory_path, '%04d.png')]) def create_video(target_path: str, fps: int) -> None: - run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', roop.globals.video_encoder, '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) + temp_directory_path = get_temp_directory_path(target_path) + run_ffmpeg(['-i', os.path.join(temp_directory_path, '%04d.png'), '-framerate', str(fps), '-c:v', roop.globals.video_encoder, '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) def restore_audio(target_path: str, output_path: str) -> None: - run_ffmpeg(['-i', get_temp_file_path(target_path), '-i', target_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path]) + temp_file_path = get_temp_file_path(target_path) + run_ffmpeg(['-i', temp_file_path, '-i', target_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path]) if not os.path.isfile(output_path): move_temp(target_path, output_path) def get_temp_frames_paths(target_path: str) -> List: - return glob.glob(get_temp_directory_path(target_path) + os.sep + '*.png') + temp_directory_path = get_temp_directory_path(target_path) + return glob.glob(os.path.join(temp_directory_path, '*.png')) def get_temp_directory_path(target_path: str) -> str: - return os.path.dirname(target_path) + os.sep + 'temp' + filename, _ = os.path.splitext(os.path.basename(target_path)) + target_name = os.path.dirname(target_path) + return os.path.join(target_name, TEMP_DIRECTORY, filename) def get_temp_file_path(target_path: str) -> str: - return get_temp_directory_path(target_path) + os.sep + 'temp.mp4' + temp_directory_path = get_temp_directory_path(target_path) + return os.path.join(temp_directory_path, TEMP_FILE) def create_temp(target_path: str) -> None: - Path(get_temp_directory_path(target_path)).mkdir(exist_ok=True) + temp_directory_path = get_temp_directory_path(target_path) + Path(temp_directory_path).mkdir(parents=True, exist_ok=True) def move_temp(target_path: str, output_path: str) -> None: @@ -68,8 +78,12 @@ def move_temp(target_path: str, output_path: str) -> None: def clean_temp(target_path: str) -> None: temp_directory_path = get_temp_directory_path(target_path) - if not roop.globals.keep_frames and os.path.isdir(temp_directory_path): + parent_directory_path = os.path.dirname(temp_directory_path) + parent_directory_name = os.path.basename(parent_directory_path) + if not roop.globals.keep_frames and os.path.isdir(temp_directory_path): shutil.rmtree(temp_directory_path) + if not os.listdir(parent_directory_path) and parent_directory_name == TEMP_DIRECTORY: + os.rmdir(parent_directory_path) def has_image_extention(image_path: str) -> bool: From 6cd16c9f351007006cd974dc43b265be7e93b11a Mon Sep 17 00:00:00 2001 From: henryruhs Date: Fri, 9 Jun 2023 08:25:53 +0200 Subject: [PATCH 35/40] Fix fps detection --- roop/core.py | 2 +- roop/utilities.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/roop/core.py b/roop/core.py index 1621210..c2db1cf 100755 --- a/roop/core.py +++ b/roop/core.py @@ -187,7 +187,7 @@ def start() -> None: torch.cuda.empty_cache() if roop.globals.keep_fps: update_status('Detecting fps...') - fps = detect_fps(roop.globals.source_path) + fps = detect_fps(roop.globals.target_path) update_status(f'Creating video with {fps} fps...') create_video(roop.globals.target_path, fps) else: diff --git a/roop/utilities.py b/roop/utilities.py index ce55969..68d508a 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -22,8 +22,8 @@ def run_ffmpeg(args: List) -> None: pass -def detect_fps(source_path: str) -> int: - command = ['ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries', 'stream=r_frame_rate', '-of', 'default=noprint_wrappers=1:nokey=1', source_path] +def detect_fps(target_path: str) -> int: + command = ['ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries', 'stream=r_frame_rate', '-of', 'default=noprint_wrappers=1:nokey=1', target_path] output = subprocess.check_output(command).decode().strip() try: return int(eval(output)) From b4ebfa4122a57939f6b396acc553663cb2b6ac30 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Fri, 9 Jun 2023 08:38:06 +0200 Subject: [PATCH 36/40] Rename method --- roop/core.py | 4 ++-- roop/swapper.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/roop/core.py b/roop/core.py index c2db1cf..b4ddab0 100755 --- a/roop/core.py +++ b/roop/core.py @@ -22,7 +22,7 @@ import cv2 import roop.globals import roop.ui as ui -from roop.swapper import process_video, process_img +from roop.swapper import process_video, process_image from roop.utilities import has_image_extention, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frames_paths, restore_audio, create_temp, move_temp, clean_temp from roop.analyser import get_one_face @@ -166,7 +166,7 @@ def start() -> None: if has_image_extention(roop.globals.target_path): if predict_image(roop.globals.target_path) > 0.85: destroy() - process_img(roop.globals.source_path, roop.globals.target_path, roop.globals.output_path) + process_image(roop.globals.source_path, roop.globals.target_path, roop.globals.output_path) if is_image(roop.globals.target_path): update_status('Swapping to image succeed!') else: diff --git a/roop/swapper.py b/roop/swapper.py index d001055..8f4e4ac 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -75,7 +75,7 @@ def multi_process_frame(source_img, frame_paths, progress): thread.join() -def process_img(source_img, target_path, output_file): +def process_image(source_img, target_path, output_file): frame = cv2.imread(target_path) target_frame = get_one_face(frame) source_face = get_one_face(cv2.imread(source_img)) From 755a5e5a3e503a13aa4a8b6ee6becd1e37ea5e9d Mon Sep 17 00:00:00 2001 From: henryruhs Date: Fri, 9 Jun 2023 09:28:02 +0200 Subject: [PATCH 37/40] Introduce suggest methods for args defaults, output mode and core/threads count via postfix --- roop/core.py | 45 ++++++++++++++++++++++----------------------- roop/swapper.py | 14 +++++++------- 2 files changed, 29 insertions(+), 30 deletions(-) diff --git a/roop/core.py b/roop/core.py index b4ddab0..a2ae5a1 100755 --- a/roop/core.py +++ b/roop/core.py @@ -45,8 +45,8 @@ def parse_args() -> None: parser.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264') parser.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18) parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) - parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) - parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) + parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=suggest_cpu_cores()) + parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=suggest_gpu_threads()) parser.add_argument('--gpu-vendor', help='select your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'nvidia']) args = parser.parse_known_args()[0] @@ -61,20 +61,8 @@ def parse_args() -> None: roop.globals.many_faces = args.many_faces roop.globals.video_encoder = args.video_encoder roop.globals.video_quality = args.video_quality - - 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 = int(args.gpu_threads) - - # gpu thread fix for amd - if args.gpu_vendor == 'amd': - roop.globals.gpu_threads = 1 + roop.globals.cpu_cores = args.cpu_cores + roop.globals.gpu_threads = args.gpu_threads if args.gpu_vendor: roop.globals.gpu_vendor = args.gpu_vendor @@ -82,7 +70,19 @@ def parse_args() -> None: roop.globals.providers = ['CPUExecutionProvider'] -def limit_resources(): +def suggest_gpu_threads() -> int: + if 'ROCMExecutionProvider' in roop.globals.providers: + return 2 + return 8 + + +def suggest_cpu_cores() -> int: + if sys.platform == 'darwin': + return 2 + return int(max(psutil.cpu_count() / 2, 1)) + + +def limit_resources() -> None: # prevent tensorflow memory leak gpus = tensorflow.config.experimental.list_physical_devices('GPU') for gpu in gpus: @@ -98,7 +98,7 @@ def limit_resources(): resource.setrlimit(resource.RLIMIT_DATA, (memory, memory)) -def pre_check(): +def pre_check() -> None: if sys.version_info < (3, 9): quit('Python version is not supported - please upgrade to 3.9 or higher.') if not shutil.which('ffmpeg'): @@ -128,23 +128,21 @@ def pre_check(): def conditional_process_video(source_path: str, frame_paths: List[str]) -> None: pool_amount = len(frame_paths) // roop.globals.cpu_cores if pool_amount > 2 and roop.globals.cpu_cores > 1 and roop.globals.gpu_vendor is None: - update_status('Pool-Swapping in progress...') global POOL POOL = multiprocessing.Pool(roop.globals.cpu_cores, maxtasksperchild=1) pools = [] for i in range(0, len(frame_paths), pool_amount): - pool = POOL.apply_async(process_video, args=(source_path, frame_paths[i:i + pool_amount])) + pool = POOL.apply_async(process_video, args=(source_path, frame_paths[i:i + pool_amount], 'cpu')) pools.append(pool) for pool in pools: pool.get() POOL.close() POOL.join() else: - update_status('Swapping in progress...') - process_video(roop.globals.source_path, frame_paths) + process_video(roop.globals.source_path, frame_paths, 'gpu') -def update_status(message: str): +def update_status(message: str) -> None: value = 'Status: ' + message print(value) if not roop.globals.headless: @@ -181,6 +179,7 @@ def start() -> None: update_status('Extracting frames...') extract_frames(roop.globals.target_path) frame_paths = get_temp_frames_paths(roop.globals.target_path) + update_status('Swapping in progress...') conditional_process_video(roop.globals.source_path, frame_paths) # prevent memory leak using ffmpeg with cuda if roop.globals.gpu_vendor == 'nvidia': diff --git a/roop/swapper.py b/roop/swapper.py index 8f4e4ac..1152423 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -1,4 +1,3 @@ - import os from tqdm import tqdm import cv2 @@ -83,12 +82,13 @@ def process_image(source_img, target_path, output_file): cv2.imwrite(output_file, result) -def process_video(source_img, frame_paths): - do_multi = roop.globals.gpu_vendor is not None and roop.globals.gpu_threads > 1 +def process_video(source_path, frame_paths, mode: str): progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' total = len(frame_paths) with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress: - if do_multi: - multi_process_frame(source_img, frame_paths, progress) - else: - process_frames(source_img, frame_paths, progress) + if mode == 'cpu': + progress.set_postfix({'mode': mode, 'cores': roop.globals.cpu_cores}) + process_frames(source_path, frame_paths, progress) + elif mode == 'gpu': + progress.set_postfix({'mode': mode, 'threads': roop.globals.gpu_threads}) + multi_process_frame(source_path, frame_paths, progress) From 583aee95cbc7d31d71ce0ebc40615b726a67054a Mon Sep 17 00:00:00 2001 From: henryruhs Date: Fri, 9 Jun 2023 11:10:20 +0200 Subject: [PATCH 38/40] Fix max_memory and output memory in progress bar too --- roop/core.py | 13 ++++++++++--- roop/globals.py | 2 +- roop/swapper.py | 4 ++-- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/roop/core.py b/roop/core.py index a2ae5a1..8aaab4a 100755 --- a/roop/core.py +++ b/roop/core.py @@ -44,7 +44,7 @@ def parse_args() -> None: parser.add_argument('--many-faces', help='swap every face in the frame', dest='many_faces', action='store_true', default=False) parser.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264') parser.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18) - parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) + parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int, default=suggest_max_memory()) parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=suggest_cpu_cores()) parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=suggest_gpu_threads()) parser.add_argument('--gpu-vendor', help='select your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'nvidia']) @@ -61,6 +61,7 @@ def parse_args() -> None: roop.globals.many_faces = args.many_faces roop.globals.video_encoder = args.video_encoder roop.globals.video_quality = args.video_quality + roop.globals.max_memory = args.max_memory roop.globals.cpu_cores = args.cpu_cores roop.globals.gpu_threads = args.gpu_threads @@ -70,6 +71,12 @@ def parse_args() -> None: roop.globals.providers = ['CPUExecutionProvider'] +def suggest_max_memory() -> int: + if platform.system().lower() == 'darwin': + return 4 + return 16 + + def suggest_gpu_threads() -> int: if 'ROCMExecutionProvider' in roop.globals.providers: return 2 @@ -77,7 +84,7 @@ def suggest_gpu_threads() -> int: def suggest_cpu_cores() -> int: - if sys.platform == 'darwin': + if platform.system().lower() == 'darwin': return 2 return int(max(psutil.cpu_count() / 2, 1)) @@ -89,7 +96,7 @@ def limit_resources() -> None: tensorflow.config.experimental.set_memory_growth(gpu, True) if roop.globals.max_memory: memory = roop.globals.max_memory * 1024 * 1024 * 1024 - if str(platform.system()).lower() == 'windows': + if platform.system().lower() == 'windows': import ctypes kernel32 = ctypes.windll.kernel32 kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory)) diff --git a/roop/globals.py b/roop/globals.py index 4ad9085..6446325 100644 --- a/roop/globals.py +++ b/roop/globals.py @@ -9,10 +9,10 @@ keep_frames = None many_faces = None video_encoder = None video_quality = None +max_memory = None cpu_cores = None gpu_threads = None gpu_vendor = None -max_memory = None headless = None log_level = 'error' providers = onnxruntime.get_available_providers() diff --git a/roop/swapper.py b/roop/swapper.py index 1152423..ddc4404 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -87,8 +87,8 @@ def process_video(source_path, frame_paths, mode: str): total = len(frame_paths) with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress: if mode == 'cpu': - progress.set_postfix({'mode': mode, 'cores': roop.globals.cpu_cores}) + progress.set_postfix({'mode': mode, 'cores': roop.globals.cpu_cores, 'memory': roop.globals.max_memory}) process_frames(source_path, frame_paths, progress) elif mode == 'gpu': - progress.set_postfix({'mode': mode, 'threads': roop.globals.gpu_threads}) + progress.set_postfix({'mode': mode, 'threads': roop.globals.gpu_threads, 'memory': roop.globals.max_memory}) multi_process_frame(source_path, frame_paths, progress) From 27a76eea60ccc4bfba4e6b8a2e11d40cc9710b99 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Fri, 9 Jun 2023 11:41:42 +0200 Subject: [PATCH 39/40] Turns out mac has a different memory unit --- roop/core.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/roop/core.py b/roop/core.py index 8aaab4a..dbea8db 100755 --- a/roop/core.py +++ b/roop/core.py @@ -95,7 +95,9 @@ def limit_resources() -> None: for gpu in gpus: tensorflow.config.experimental.set_memory_growth(gpu, True) if roop.globals.max_memory: - memory = roop.globals.max_memory * 1024 * 1024 * 1024 + memory = roop.globals.max_memory * 1024 ** 3 + if platform.system().lower() == 'darwin': + memory = roop.globals.max_memory * 1024 ** 6 if platform.system().lower() == 'windows': import ctypes kernel32 = ctypes.windll.kernel32 From 1a7db694cf2deb3598c45bc7787ba1501c109016 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Fri, 9 Jun 2023 12:39:50 +0200 Subject: [PATCH 40/40] Add typing to swapper --- roop/swapper.py | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/roop/swapper.py b/roop/swapper.py index ddc4404..8bc060b 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -1,4 +1,6 @@ import os +from typing import Any + from tqdm import tqdm import cv2 import insightface @@ -10,7 +12,7 @@ FACE_SWAPPER = None THREAD_LOCK = threading.Lock() -def get_face_swapper(): +def get_face_swapper() -> None: global FACE_SWAPPER with THREAD_LOCK: if FACE_SWAPPER is None: @@ -19,13 +21,13 @@ def get_face_swapper(): return FACE_SWAPPER -def swap_face_in_frame(source_face, target_face, frame): +def swap_face_in_frame(source_face: Any, target_face: Any, frame: Any) -> None: if target_face: return get_face_swapper().get(frame, target_face, source_face, paste_back=True) return frame -def process_faces(source_face, target_frame): +def process_faces(source_face: Any, target_frame: Any) -> Any: if roop.globals.many_faces: many_faces = get_many_faces(target_frame) if many_faces: @@ -38,8 +40,8 @@ def process_faces(source_face, target_frame): return target_frame -def process_frames(source_img, frame_paths, progress=None): - source_face = get_one_face(cv2.imread(source_img)) +def process_frames(source_path: str, frame_paths: [str], progress=None) -> None: + source_face = get_one_face(cv2.imread(source_path)) for frame_path in frame_paths: frame = cv2.imread(frame_path) try: @@ -52,15 +54,14 @@ def process_frames(source_img, frame_paths, progress=None): progress.update(1) -def multi_process_frame(source_img, frame_paths, progress): +def multi_process_frame(source_img, frame_paths, progress) -> None: threads = [] - num_threads = roop.globals.gpu_threads - num_frames_per_thread = len(frame_paths) // num_threads - remaining_frames = len(frame_paths) % num_threads + frames_per_thread = len(frame_paths) // roop.globals.gpu_threads + remaining_frames = len(frame_paths) % roop.globals.gpu_threads start_index = 0 # create threads by frames - for _ in range(num_threads): - end_index = start_index + num_frames_per_thread + for _ in range(roop.globals.gpu_threads): + end_index = start_index + frames_per_thread if remaining_frames > 0: end_index += 1 remaining_frames -= 1 @@ -74,15 +75,15 @@ def multi_process_frame(source_img, frame_paths, progress): thread.join() -def process_image(source_img, target_path, output_file): +def process_image(source_path: str, target_path: str, output_file) -> None: frame = cv2.imread(target_path) target_frame = get_one_face(frame) - source_face = get_one_face(cv2.imread(source_img)) + source_face = get_one_face(cv2.imread(source_path)) result = get_face_swapper().get(frame, target_frame, source_face, paste_back=True) cv2.imwrite(output_file, result) -def process_video(source_path, frame_paths, mode: str): +def process_video(source_path: str, frame_paths: [str], mode: str) -> None: progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' total = len(frame_paths) with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress: