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)