Stream video frames to avoid disk I/O

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):

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'
 ]
 
 

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)
+

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)
+    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)

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",