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