diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py index a8ca32c1..a0d14de1 100644 --- a/facefusion/ffmpeg.py +++ b/facefusion/ffmpeg.py @@ -70,6 +70,33 @@ def open_ffmpeg(commands : List[Command]) -> subprocess.Popen[bytes]: return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.PIPE) +def open_video_encoder(target_path : str, temp_video_fps : Fps, temp_video_resolution : Resolution, output_video_resolution : Resolution, output_video_fps : Fps) -> subprocess.Popen[bytes]: + output_video_encoder = state_manager.get_item('output_video_encoder') + output_video_quality = state_manager.get_item('output_video_quality') + output_video_preset = state_manager.get_item('output_video_preset') + temp_video_path = get_temp_file_path(target_path) + temp_video_format = cast(VideoFormat, get_file_format(temp_video_path)) + + output_video_encoder = fix_video_encoder(temp_video_format, output_video_encoder) + commands = ffmpeg_builder.chain( + ffmpeg_builder.set_output_format('rawvideo'), + ffmpeg_builder.enforce_pixel_format('bgr24'), + ffmpeg_builder.set_media_resolution(pack_resolution(temp_video_resolution)), + ffmpeg_builder.set_input_fps(temp_video_fps), + ffmpeg_builder.set_input('pipe:0'), + ffmpeg_builder.set_media_resolution(pack_resolution(output_video_resolution)), + ffmpeg_builder.set_video_encoder(output_video_encoder), + ffmpeg_builder.set_video_tag(output_video_encoder, temp_video_format), + ffmpeg_builder.set_video_quality(output_video_encoder, output_video_quality), + ffmpeg_builder.set_video_preset(output_video_encoder, output_video_preset), + ffmpeg_builder.set_video_fps(output_video_fps), + ffmpeg_builder.set_pixel_format(output_video_encoder), + ffmpeg_builder.force_output(temp_video_path) + ) + commands = ffmpeg_builder.run(commands) + return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.DEVNULL) + + def log_debug(process : subprocess.Popen[bytes]) -> None: _, stderr = process.communicate() errors = stderr.decode().split(os.linesep) diff --git a/facefusion/workflows/image_to_video.py b/facefusion/workflows/image_to_video.py index e51528f2..e3dca160 100644 --- a/facefusion/workflows/image_to_video.py +++ b/facefusion/workflows/image_to_video.py @@ -1,5 +1,6 @@ -from concurrent.futures import ThreadPoolExecutor, as_completed +from concurrent.futures import ThreadPoolExecutor from functools import partial +from typing import Tuple import cv2 import numpy @@ -12,10 +13,10 @@ from facefusion.common_helper import get_first, get_middle from facefusion.content_analyser import analyse_video from facefusion.filesystem import filter_audio_paths, is_video from facefusion.processors.core import get_processors_modules -from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_frame_pattern, move_temp_file +from facefusion.temp_helper import clear_temp_directory, create_temp_directory, move_temp_file from facefusion.time_helper import calculate_end_time -from facefusion.types import ErrorCode, Resolution -from facefusion.vision import conditional_merge_vision_mask, detect_video_resolution, extract_vision_mask, pack_resolution, read_static_images, read_static_video_frame, restrict_trim_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames, write_image +from facefusion.types import ErrorCode, Resolution, VisionFrame +from facefusion.vision import conditional_merge_vision_mask, detect_video_resolution, extract_vision_mask, pack_resolution, read_static_images, read_static_video_frame, restrict_trim_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames from facefusion.workflows.core import is_process_stopping @@ -24,7 +25,6 @@ def process(start_time : float) -> ErrorCode: [ setup, process_video, - merge_frames, restore_audio, partial(finalize_video, start_time) ] @@ -77,33 +77,43 @@ def process_video() -> ErrorCode: trim_frame_start, trim_frame_end = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) output_video_resolution = scale_resolution(detect_video_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_video_scale')) temp_video_resolution = restrict_video_resolution(state_manager.get_item('target_path'), output_video_resolution) + temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) frame_range = range(trim_frame_start, trim_frame_end) if frame_range: + video_encoder = ffmpeg.open_video_encoder(state_manager.get_item('target_path'), temp_video_fps, temp_video_resolution, output_video_resolution, state_manager.get_item('output_video_fps')) + with tqdm(total = len(frame_range), desc = translator.get('processing'), unit = 'frame', ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: progress.set_postfix(execution_providers = state_manager.get_item('execution_providers')) with ThreadPoolExecutor(max_workers = state_manager.get_item('execution_thread_count')) as executor: futures = [] + frame_look_ahead = state_manager.get_item('execution_thread_count') * 2 for frame_number in frame_range: - future = executor.submit(process_temp_frame, frame_number, temp_video_resolution) - futures.append(future) + futures.append(executor.submit(process_stream_frame, frame_number, temp_video_resolution)) - for future in as_completed(futures): - if is_process_stopping(): - for __future__ in futures: - __future__.cancel() + if len(futures) >= frame_look_ahead and is_process_stopping() is False: + _, vision_frame = futures.pop(0).result() + video_encoder.stdin.write(vision_frame.tobytes()) + progress.update() - if not future.cancelled(): - future.result() + for future in futures: + if is_process_stopping() is False: + _, vision_frame = future.result() + video_encoder.stdin.write(vision_frame.tobytes()) progress.update() for processor_module in get_processors_modules(state_manager.get_item('processors')): processor_module.post_process() + video_encoder.stdin.close() + video_encoder.wait() + if is_process_stopping(): return 4 + if not video_encoder.returncode == 0: + return 1 else: logger.error(translator.get('temp_frames_not_found'), __name__) return 1 @@ -157,7 +167,7 @@ def restore_audio() -> ErrorCode: return 0 -def process_temp_frame(frame_number : int, temp_video_resolution : Resolution) -> bool: +def process_stream_frame(frame_number : int, temp_video_resolution : Resolution) -> Tuple[int, VisionFrame]: trim_frame_start, _ = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) reference_vision_frame = read_static_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) source_vision_frames = read_static_images(state_manager.get_item('source_paths')) @@ -165,7 +175,6 @@ def process_temp_frame(frame_number : int, temp_video_resolution : Resolution) - target_vision_frames = select_video_frames(state_manager.get_item('target_path'), frame_number, state_manager.get_item('target_frame_amount')) target_vision_frame = get_middle(target_vision_frames) temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) - temp_frame_path = get_temp_frame_pattern(state_manager.get_item('target_path'), format(frame_number, '08d')) temp_vision_frame = target_vision_frame.copy() if not (target_vision_frame.shape[1], target_vision_frame.shape[0]) == temp_video_resolution: @@ -193,7 +202,7 @@ def process_temp_frame(frame_number : int, temp_video_resolution : Resolution) - }) temp_vision_frame = conditional_merge_vision_mask(temp_vision_frame, temp_vision_mask) - return write_image(temp_frame_path, temp_vision_frame) + return frame_number, numpy.ascontiguousarray(temp_vision_frame[:, :, :3]) def finalize_video(start_time : float) -> ErrorCode: