From cf814bbcee3c161838297a1eb1192aeeee50cb33 Mon Sep 17 00:00:00 2001 From: Christoph9211 <123782754+Christoph9211@users.noreply.github.com> Date: Wed, 4 Jun 2025 00:31:38 -0500 Subject: [PATCH 1/6] Stream video frames to avoid disk I/O --- modules/core.py | 73 +++++++++++++++++++++-- modules/processors/frame/core.py | 3 +- modules/processors/frame/face_enhancer.py | 5 ++ modules/processors/frame/face_swapper.py | 19 +++++- modules/utilities.py | 32 ++++++++++ 5 files changed, 125 insertions(+), 7 deletions(-) diff --git a/modules/core.py b/modules/core.py index b6ef9b8..9ea867b 100644 --- a/modules/core.py +++ b/modules/core.py @@ -19,7 +19,24 @@ import modules.globals import modules.metadata import modules.ui as ui from modules.processors.frame.core import get_frame_processors_modules -from modules.utilities import has_image_extension, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frame_paths, restore_audio, create_temp, move_temp, clean_temp, normalize_output_path +from modules.utilities import ( + has_image_extension, + is_image, + is_video, + detect_fps, + create_video, + extract_frames, + get_temp_frame_paths, + restore_audio, + create_temp, + move_temp, + clean_temp, + normalize_output_path, + start_ffmpeg_writer, + get_temp_output_path, +) +import cv2 +from tqdm import tqdm if 'ROCMExecutionProvider' in modules.globals.execution_providers: del torch @@ -175,6 +192,45 @@ def update_status(message: str, scope: str = 'DLC.CORE') -> None: if not modules.globals.headless: ui.update_status(message) + +def stream_video() -> None: + capture = cv2.VideoCapture(modules.globals.target_path) + if not capture.isOpened(): + update_status('Failed to open video file.') + return + fps = capture.get(cv2.CAP_PROP_FPS) if modules.globals.keep_fps else 30.0 + width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) + + update_status('Creating temp resources...') + create_temp(modules.globals.target_path) + temp_output_path = get_temp_output_path(modules.globals.target_path) + writer = start_ffmpeg_writer(width, height, fps, temp_output_path) + + progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' + with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress: + progress.set_postfix({'execution_providers': modules.globals.execution_providers, 'execution_threads': modules.globals.execution_threads, 'max_memory': modules.globals.max_memory}) + while True: + ret, frame = capture.read() + if not ret: + break + for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): + frame = frame_processor.process_frame_stream(modules.globals.source_path, frame) + writer.stdin.write(frame.tobytes()) + progress.update(1) + + capture.release() + writer.stdin.close() + writer.wait() + + if modules.globals.keep_audio: + update_status('Restoring audio...') + restore_audio(modules.globals.target_path, modules.globals.output_path) + else: + move_temp(modules.globals.target_path, modules.globals.output_path) + clean_temp(modules.globals.target_path) + def start() -> None: for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): if not frame_processor.pre_start(): @@ -202,10 +258,17 @@ def start() -> None: return if not modules.globals.map_faces: - update_status('Creating temp resources...') - create_temp(modules.globals.target_path) - update_status('Extracting frames...') - extract_frames(modules.globals.target_path) + stream_video() + if is_video(modules.globals.target_path): + update_status('Processing to video succeed!') + else: + update_status('Processing to video failed!') + return + + update_status('Creating temp resources...') + create_temp(modules.globals.target_path) + update_status('Extracting frames...') + extract_frames(modules.globals.target_path) temp_frame_paths = get_temp_frame_paths(modules.globals.target_path) for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): diff --git a/modules/processors/frame/core.py b/modules/processors/frame/core.py index 7d76704..0b6f8af 100644 --- a/modules/processors/frame/core.py +++ b/modules/processors/frame/core.py @@ -14,7 +14,8 @@ FRAME_PROCESSORS_INTERFACE = [ 'pre_start', 'process_frame', 'process_image', - 'process_video' + 'process_video', + 'process_frame_stream' ] diff --git a/modules/processors/frame/face_enhancer.py b/modules/processors/frame/face_enhancer.py index 4e1fdff..c418c84 100644 --- a/modules/processors/frame/face_enhancer.py +++ b/modules/processors/frame/face_enhancer.py @@ -107,3 +107,8 @@ def process_frame_v2(temp_frame: Frame) -> Frame: if target_face: temp_frame = enhance_face(temp_frame) return temp_frame + + +def process_frame_stream(source_path: str, frame: Frame) -> Frame: + return process_frame(None, frame) + diff --git a/modules/processors/frame/face_swapper.py b/modules/processors/frame/face_swapper.py index b09e600..c04ef1b 100644 --- a/modules/processors/frame/face_swapper.py +++ b/modules/processors/frame/face_swapper.py @@ -255,4 +255,21 @@ def process_video(source_path: str, temp_frame_paths: List[str]) -> None: if modules.globals.map_faces and modules.globals.many_faces: update_status('Many faces enabled. Using first source image (if applicable in v2). Processing...', NAME) # The core processing logic is delegated, which is good. - modules.processors.frame.core.process_video(source_path, temp_frame_paths, process_frames) \ No newline at end of file + modules.processors.frame.core.process_video(source_path, temp_frame_paths, process_frames) + + +STREAM_SOURCE_FACE = None + + +def process_frame_stream(source_path: str, frame: Frame) -> Frame: + global STREAM_SOURCE_FACE + if not modules.globals.map_faces: + if STREAM_SOURCE_FACE is None: + source_img = cv2.imread(source_path) + if source_img is not None: + STREAM_SOURCE_FACE = get_one_face(source_img) + if STREAM_SOURCE_FACE is not None: + return process_frame(STREAM_SOURCE_FACE, frame) + return frame + else: + return process_frame_v2(frame) diff --git a/modules/utilities.py b/modules/utilities.py index fe17997..13bc93b 100644 --- a/modules/utilities.py +++ b/modules/utilities.py @@ -38,6 +38,38 @@ def run_ffmpeg(args: List[str]) -> bool: return False +def start_ffmpeg_writer(width: int, height: int, fps: float, output_path: str) -> subprocess.Popen: + commands = [ + "ffmpeg", + "-hide_banner", + "-hwaccel", + "auto", + "-loglevel", + modules.globals.log_level, + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s", + f"{width}x{height}", + "-r", + str(fps), + "-i", + "-", + "-c:v", + modules.globals.video_encoder, + "-crf", + str(modules.globals.video_quality), + "-pix_fmt", + "yuv420p", + "-vf", + "colorspace=bt709:iall=bt601-6-625:fast=1", + "-y", + output_path, + ] + return subprocess.Popen(commands, stdin=subprocess.PIPE) + + def detect_fps(target_path: str) -> float: command = [ "ffprobe", From df31fa808307b69ea034534692cbdf993b165a13 Mon Sep 17 00:00:00 2001 From: Christoph9211 Date: Sun, 8 Jun 2025 15:45:55 -0500 Subject: [PATCH 2/6] typo fix --- modules/processors/frame/face_swapper.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/modules/processors/frame/face_swapper.py b/modules/processors/frame/face_swapper.py index c04ef1b..f1c65cd 100644 --- a/modules/processors/frame/face_swapper.py +++ b/modules/processors/frame/face_swapper.py @@ -123,12 +123,12 @@ def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame: if is_image(modules.globals.target_path): if modules.globals.many_faces: source_face = default_source_face() - for map_entry in modules.globals.souce_target_map: # Renamed 'map' to 'map_entry' + for map_entry in modules.globals.source_target_map: # Renamed 'map' to 'map_entry' target_face = map_entry['target']['face'] temp_frame = swap_face(source_face, target_face, temp_frame) elif not modules.globals.many_faces: - for map_entry in modules.globals.souce_target_map: # Renamed 'map' to 'map_entry' + for map_entry in modules.globals.source_target_map: # Renamed 'map' to 'map_entry' if "source" in map_entry: source_face = map_entry['source']['face'] target_face = map_entry['target']['face'] @@ -137,7 +137,7 @@ def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame: elif is_video(modules.globals.target_path): if modules.globals.many_faces: source_face = default_source_face() - for map_entry in modules.globals.souce_target_map: # Renamed 'map' to 'map_entry' + for map_entry in modules.globals.source_target_map: # Renamed 'map' to 'map_entry' target_frame = [f for f in map_entry['target_faces_in_frame'] if f['location'] == temp_frame_path] for frame in target_frame: @@ -145,7 +145,7 @@ def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame: temp_frame = swap_face(source_face, target_face, temp_frame) elif not modules.globals.many_faces: - for map_entry in modules.globals.souce_target_map: # Renamed 'map' to 'map_entry' + for map_entry in modules.globals.source_target_map: # Renamed 'map' to 'map_entry' if "source" in map_entry: target_frame = [f for f in map_entry['target_faces_in_frame'] if f['location'] == temp_frame_path] source_face = map_entry['source']['face'] From 45f284adb6f24cbad0987a1fa67e5f7c076f7bd6 Mon Sep 17 00:00:00 2001 From: Christoph9211 Date: Sun, 8 Jun 2025 16:54:47 -0500 Subject: [PATCH 3/6] Refactor variable names in process_frame_v2 for clarity --- modules/processors/frame/face_swapper.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/modules/processors/frame/face_swapper.py b/modules/processors/frame/face_swapper.py index f1c65cd..f38046d 100644 --- a/modules/processors/frame/face_swapper.py +++ b/modules/processors/frame/face_swapper.py @@ -123,32 +123,32 @@ def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame: if is_image(modules.globals.target_path): if modules.globals.many_faces: source_face = default_source_face() - for map_entry in modules.globals.source_target_map: # Renamed 'map' to 'map_entry' - target_face = map_entry['target']['face'] + for map in modules.globals.source_target_map: + target_face = map['target']['face'] temp_frame = swap_face(source_face, target_face, temp_frame) elif not modules.globals.many_faces: - for map_entry in modules.globals.source_target_map: # Renamed 'map' to 'map_entry' - if "source" in map_entry: - source_face = map_entry['source']['face'] - target_face = map_entry['target']['face'] + for map in modules.globals.source_target_map: + if "source" in map: + source_face = map['source']['face'] + target_face = map['target']['face'] temp_frame = swap_face(source_face, target_face, temp_frame) elif is_video(modules.globals.target_path): if modules.globals.many_faces: source_face = default_source_face() - for map_entry in modules.globals.source_target_map: # Renamed 'map' to 'map_entry' - target_frame = [f for f in map_entry['target_faces_in_frame'] if f['location'] == temp_frame_path] + for map in modules.globals.source_target_map: + target_frame = [f for f in map['target_faces_in_frame'] if f['location'] == temp_frame_path] for frame in target_frame: for target_face in frame['faces']: temp_frame = swap_face(source_face, target_face, temp_frame) elif not modules.globals.many_faces: - for map_entry in modules.globals.source_target_map: # Renamed 'map' to 'map_entry' - if "source" in map_entry: - target_frame = [f for f in map_entry['target_faces_in_frame'] if f['location'] == temp_frame_path] - source_face = map_entry['source']['face'] + for map in modules.globals.source_target_map: + if "source" in map: + target_frame = [f for f in map['target_faces_in_frame'] if f['location'] == temp_frame_path] + source_face = map['source']['face'] for frame in target_frame: for target_face in frame['faces']: From 11d2e93d72224f636989357bd05d29e36c2953b2 Mon Sep 17 00:00:00 2001 From: Christoph9211 Date: Sun, 8 Jun 2025 17:00:55 -0500 Subject: [PATCH 4/6] Improve process_frame_stream logic and ensure safe argument handling in start_ffmpeg_writer --- modules/processors/frame/face_swapper.py | 23 +++++++++++++---------- modules/utilities.py | 7 ++++--- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/modules/processors/frame/face_swapper.py b/modules/processors/frame/face_swapper.py index f38046d..cea0b59 100644 --- a/modules/processors/frame/face_swapper.py +++ b/modules/processors/frame/face_swapper.py @@ -263,13 +263,16 @@ STREAM_SOURCE_FACE = None def process_frame_stream(source_path: str, frame: Frame) -> Frame: global STREAM_SOURCE_FACE - if not modules.globals.map_faces: - if STREAM_SOURCE_FACE is None: - source_img = cv2.imread(source_path) - if source_img is not None: - STREAM_SOURCE_FACE = get_one_face(source_img) - if STREAM_SOURCE_FACE is not None: - return process_frame(STREAM_SOURCE_FACE, frame) - return frame - else: - return process_frame_v2(frame) + if modules.globals.map_faces: + result = process_frame_v2(frame) + if result is not None: + return result + else: + return frame # Fallback to original frame if process_frame_v2 returns None + if STREAM_SOURCE_FACE is None: + source_img = cv2.imread(source_path) + if source_img is not None: + STREAM_SOURCE_FACE = get_one_face(source_img) + if STREAM_SOURCE_FACE is not None: + return process_frame(STREAM_SOURCE_FACE, frame) + return frame diff --git a/modules/utilities.py b/modules/utilities.py index 13bc93b..a146b20 100644 --- a/modules/utilities.py +++ b/modules/utilities.py @@ -39,13 +39,14 @@ def run_ffmpeg(args: List[str]) -> bool: def start_ffmpeg_writer(width: int, height: int, fps: float, output_path: str) -> subprocess.Popen: + # Pass all arguments as a list to avoid shell injection commands = [ "ffmpeg", "-hide_banner", "-hwaccel", "auto", "-loglevel", - modules.globals.log_level, + str(modules.globals.log_level), "-f", "rawvideo", "-pix_fmt", @@ -57,7 +58,7 @@ def start_ffmpeg_writer(width: int, height: int, fps: float, output_path: str) - "-i", "-", "-c:v", - modules.globals.video_encoder, + str(modules.globals.video_encoder), "-crf", str(modules.globals.video_quality), "-pix_fmt", @@ -65,7 +66,7 @@ def start_ffmpeg_writer(width: int, height: int, fps: float, output_path: str) - "-vf", "colorspace=bt709:iall=bt601-6-625:fast=1", "-y", - output_path, + str(output_path), ] return subprocess.Popen(commands, stdin=subprocess.PIPE) From 216fac1122a16012e4ef4f3328afed6b45e8c9bd Mon Sep 17 00:00:00 2001 From: Christoph9211 Date: Sun, 8 Jun 2025 17:01:33 -0500 Subject: [PATCH 5/6] Handle non-zero exit status in ffmpeg writer to improve error reporting --- modules/core.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modules/core.py b/modules/core.py index 9ea867b..f66be6d 100644 --- a/modules/core.py +++ b/modules/core.py @@ -222,7 +222,9 @@ def stream_video() -> None: capture.release() writer.stdin.close() - writer.wait() + exit_code = writer.wait() + if exit_code != 0: + raise RuntimeError(f"ffmpeg writer exited with non-zero status: {exit_code}") if modules.globals.keep_audio: update_status('Restoring audio...') From f7fdaaec20c10edefd4815e7ce9158ef13c59736 Mon Sep 17 00:00:00 2001 From: Christoph9211 Date: Sun, 8 Jun 2025 17:08:08 -0500 Subject: [PATCH 6/6] Refactor device selection logic in get_face_enhancer for improved clarity and compatibility with Mac OS --- modules/processors/frame/face_enhancer.py | 39 ++++++----------------- 1 file changed, 9 insertions(+), 30 deletions(-) diff --git a/modules/processors/frame/face_enhancer.py b/modules/processors/frame/face_enhancer.py index e219f12..c418c84 100644 --- a/modules/processors/frame/face_enhancer.py +++ b/modules/processors/frame/face_enhancer.py @@ -48,17 +48,6 @@ def pre_start() -> bool: return True -TENSORRT_AVAILABLE = False -try: - import torch_tensorrt - TENSORRT_AVAILABLE = True -except ImportError as im: - print(f"TensorRT is not available: {im}") - pass -except Exception as e: - print(f"TensorRT is not available: {e}") - pass - def get_face_enhancer() -> Any: global FACE_ENHANCER @@ -66,26 +55,16 @@ def get_face_enhancer() -> Any: if FACE_ENHANCER is None: model_path = os.path.join(models_dir, "GFPGANv1.4.pth") - selected_device = None - device_priority = [] + match platform.system(): + case "Darwin": # Mac OS + if torch.backends.mps.is_available(): + mps_device = torch.device("mps") + FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=mps_device) # type: ignore[attr-defined] + else: + FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1) # type: ignore[attr-defined] + case _: # Other OS + FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1) # type: ignore[attr-defined] - if TENSORRT_AVAILABLE and torch.cuda.is_available(): - selected_device = torch.device("cuda") - device_priority.append("TensorRT+CUDA") - elif torch.cuda.is_available(): - selected_device = torch.device("cuda") - device_priority.append("CUDA") - elif torch.backends.mps.is_available() and platform.system() == "Darwin": - selected_device = torch.device("mps") - device_priority.append("MPS") - elif not torch.cuda.is_available(): - selected_device = torch.device("cpu") - device_priority.append("CPU") - - FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=selected_device) - - # for debug: - print(f"Selected device: {selected_device} and device priority: {device_priority}") return FACE_ENHANCER