From 583aee95cbc7d31d71ce0ebc40615b726a67054a Mon Sep 17 00:00:00 2001
From: henryruhs <info@henryruhs.com>
Date: Fri, 9 Jun 2023 11:10:20 +0200
Subject: [PATCH] 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)