Fix max_memory and output memory in progress bar too

This commit is contained in:
henryruhs 2023-06-09 11:10:20 +02:00
parent 755a5e5a3e
commit 583aee95cb
3 changed files with 13 additions and 6 deletions

View File

@ -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('--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-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('--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('--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-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']) 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.many_faces = args.many_faces
roop.globals.video_encoder = args.video_encoder roop.globals.video_encoder = args.video_encoder
roop.globals.video_quality = args.video_quality roop.globals.video_quality = args.video_quality
roop.globals.max_memory = args.max_memory
roop.globals.cpu_cores = args.cpu_cores roop.globals.cpu_cores = args.cpu_cores
roop.globals.gpu_threads = args.gpu_threads roop.globals.gpu_threads = args.gpu_threads
@ -70,6 +71,12 @@ def parse_args() -> None:
roop.globals.providers = ['CPUExecutionProvider'] roop.globals.providers = ['CPUExecutionProvider']
def suggest_max_memory() -> int:
if platform.system().lower() == 'darwin':
return 4
return 16
def suggest_gpu_threads() -> int: def suggest_gpu_threads() -> int:
if 'ROCMExecutionProvider' in roop.globals.providers: if 'ROCMExecutionProvider' in roop.globals.providers:
return 2 return 2
@ -77,7 +84,7 @@ def suggest_gpu_threads() -> int:
def suggest_cpu_cores() -> int: def suggest_cpu_cores() -> int:
if sys.platform == 'darwin': if platform.system().lower() == 'darwin':
return 2 return 2
return int(max(psutil.cpu_count() / 2, 1)) return int(max(psutil.cpu_count() / 2, 1))
@ -89,7 +96,7 @@ def limit_resources() -> None:
tensorflow.config.experimental.set_memory_growth(gpu, True) tensorflow.config.experimental.set_memory_growth(gpu, True)
if roop.globals.max_memory: if roop.globals.max_memory:
memory = roop.globals.max_memory * 1024 * 1024 * 1024 memory = roop.globals.max_memory * 1024 * 1024 * 1024
if str(platform.system()).lower() == 'windows': if platform.system().lower() == 'windows':
import ctypes import ctypes
kernel32 = ctypes.windll.kernel32 kernel32 = ctypes.windll.kernel32
kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory)) kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))

View File

@ -9,10 +9,10 @@ keep_frames = None
many_faces = None many_faces = None
video_encoder = None video_encoder = None
video_quality = None video_quality = None
max_memory = None
cpu_cores = None cpu_cores = None
gpu_threads = None gpu_threads = None
gpu_vendor = None gpu_vendor = None
max_memory = None
headless = None headless = None
log_level = 'error' log_level = 'error'
providers = onnxruntime.get_available_providers() providers = onnxruntime.get_available_providers()

View File

@ -87,8 +87,8 @@ def process_video(source_path, frame_paths, mode: str):
total = len(frame_paths) total = len(frame_paths)
with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress: with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
if mode == 'cpu': 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) process_frames(source_path, frame_paths, progress)
elif mode == 'gpu': 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) multi_process_frame(source_path, frame_paths, progress)