diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f3e00cc5..4cd98e0f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -46,7 +46,7 @@ jobs: - name: Checkout uses: actions/checkout@v4 - name: Set up FFmpeg - uses: FedericoCarboni/setup-ffmpeg@v3 + uses: AnimMouse/setup-ffmpeg@v1 - name: Set up Python 3.12 uses: actions/setup-python@v5 with: diff --git a/facefusion.ini b/facefusion.ini index af63451b..ae4021ec 100644 --- a/facefusion.ini +++ b/facefusion.ini @@ -1,6 +1,3 @@ -[workflow] -workflow_mode = - [paths] temp_path = jobs_path = @@ -54,6 +51,7 @@ voice_extractor_model = trim_frame_start = trim_frame_end = temp_frame_format = +temp_pixel_format = [frame_distribution] target_frame_amount = @@ -71,6 +69,10 @@ output_video_quality = output_video_scale = output_video_fps = +[workflow] +workflow_mode = +workflow_strategy = + [processors] processors = age_modifier_model = diff --git a/facefusion/apis/stream_manager.py b/facefusion/apis/stream_manager.py index aa7ac130..62268ee7 100644 --- a/facefusion/apis/stream_manager.py +++ b/facefusion/apis/stream_manager.py @@ -14,13 +14,15 @@ from facefusion.apis.stream_audio import receive_audio_frames, run_audio_encode_ from facefusion.apis.stream_video import receive_video_frames, run_video_encode_loop from facefusion.libraries import datachannel as datachannel_module from facefusion.types import AudioCodec, AudioFrame, BufferPack, PeerConnection, RtcPeer, RtcPeerAudio, SdpAnswer, SdpOffer, SessionId, Time, VideoCodec, VisionFrame +from facefusion.vision import read_static_images async def process_image(websocket : WebSocket) -> None: capture_vision_frame = await anext(receive_vision_frames(websocket), None) if numpy.any(capture_vision_frame): - output_vision_frame = streamer.process_stream_frame(capture_vision_frame) + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) + output_vision_frame = streamer.process_stream_frame(source_vision_frames, capture_vision_frame) is_success, output_frame_buffer = cv2.imencode('.jpg', output_vision_frame) if is_success: diff --git a/facefusion/apis/stream_video.py b/facefusion/apis/stream_video.py index ae502478..b593786b 100644 --- a/facefusion/apis/stream_video.py +++ b/facefusion/apis/stream_video.py @@ -1,15 +1,16 @@ from concurrent.futures import Future, ThreadPoolExecutor from functools import partial from queue import Queue -from typing import Optional, Tuple +from typing import List, Optional, Tuple import cv2 import numpy -from facefusion import rtc, streamer +from facefusion import rtc, state_manager, streamer from facefusion.apis.stream_event import create_receive_event from facefusion.codecs import aom_decoder, aom_encoder, vpx_decoder, vpx_encoder from facefusion.types import AomDecoder, AomEncoder, BitRate, Buffer, BufferPack, Resolution, RtcPeer, RtcPeerVideo, Time, VideoCodec, VisionFrame, VpxDecoder, VpxEncoder +from facefusion.vision import read_static_images def run_video_encode_loop(rtc_peer : RtcPeer, video_queue : Queue[Tuple[Time, Future[BufferPack]]]) -> None: @@ -59,8 +60,9 @@ def receive_video_frames(rtc_peer_video : RtcPeerVideo, video_queue : Queue[Tupl video_track = rtc_peer_video.get('receiver_track') video_codec = rtc_peer_video.get('codec') video_decoder = create_video_decoder(video_codec) + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) - video_frame_handler = partial(handle_video_frame, video_codec, video_decoder, video_queue, video_executor) + video_frame_handler = partial(handle_video_frame, source_vision_frames, video_codec, video_decoder, video_queue, video_executor) receive_event = create_receive_event(video_track, video_frame_handler) receive_event.wait() @@ -70,8 +72,8 @@ def receive_video_frames(rtc_peer_video : RtcPeerVideo, video_queue : Queue[Tupl destroy_video_decoder(video_codec, video_decoder) -def process_video_frame(input_vision_frame : VisionFrame) -> BufferPack: - output_vision_frame = streamer.process_stream_frame(input_vision_frame) +def process_video_frame(source_vision_frames : List[VisionFrame], input_vision_frame : VisionFrame) -> BufferPack: + output_vision_frame = streamer.process_stream_frame(source_vision_frames, input_vision_frame) output_resolution : Resolution = (output_vision_frame.shape[1], output_vision_frame.shape[0]) output_buffer = cv2.cvtColor(output_vision_frame, cv2.COLOR_BGR2YUV_I420).tobytes() return BufferPack(buffer = output_buffer, resolution = output_resolution) @@ -165,10 +167,10 @@ def update_video_encoder_bitrate(video_codec : VideoCodec, video_encoder : VpxEn return False -def handle_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, video_queue : Queue[Tuple[Time, Future[BufferPack]]], video_executor : ThreadPoolExecutor, video_buffer : Buffer, video_timestamp : int) -> None: +def handle_video_frame(source_vision_frames : List[VisionFrame], video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, video_queue : Queue[Tuple[Time, Future[BufferPack]]], video_executor : ThreadPoolExecutor, video_buffer : Buffer, video_timestamp : int) -> None: vision_frame = decode_video_frame(video_codec, video_decoder, video_buffer) if numpy.any(vision_frame) and video_queue.qsize() < video_queue.maxsize: - video_future = video_executor.submit(process_video_frame, vision_frame) + video_future = video_executor.submit(process_video_frame, source_vision_frames, vision_frame) video_time = rtc.convert_timestamp_to_time(video_codec, video_timestamp) video_queue.put((video_time, video_future)) diff --git a/facefusion/args_helper.py b/facefusion/args_helper.py index fffa627e..2c6c2f5a 100644 --- a/facefusion/args_helper.py +++ b/facefusion/args_helper.py @@ -11,7 +11,6 @@ from facefusion.vision import detect_video_fps def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('command', args.get('command')) - apply_state_item('workflow_mode', args.get('workflow_mode')) apply_state_item('temp_path', args.get('temp_path')) apply_state_item('jobs_path', args.get('jobs_path')) apply_state_item('source_paths', args.get('source_paths')) @@ -48,6 +47,7 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('trim_frame_start', args.get('trim_frame_start')) apply_state_item('trim_frame_end', args.get('trim_frame_end')) apply_state_item('temp_frame_format', args.get('temp_frame_format')) + apply_state_item('temp_pixel_format', args.get('temp_pixel_format')) apply_state_item('target_frame_amount', args.get('target_frame_amount')) apply_state_item('output_image_quality', args.get('output_image_quality')) apply_state_item('output_image_scale', args.get('output_image_scale')) @@ -64,6 +64,8 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: output_video_fps = normalize_fps(args.get('output_video_fps')) or detect_video_fps(args.get('target_path')) apply_state_item('output_video_fps', output_video_fps) + apply_state_item('workflow_mode', args.get('workflow_mode')) + apply_state_item('workflow_strategy', args.get('workflow_strategy')) available_processors = [ get_file_name(file_path) for file_path in resolve_file_paths('facefusion/processors/modules') ] apply_state_item('processors', args.get('processors')) diff --git a/facefusion/choices.py b/facefusion/choices.py index 8ffe3b83..4b9316e2 100755 --- a/facefusion/choices.py +++ b/facefusion/choices.py @@ -2,7 +2,7 @@ import logging from typing import List, Sequence, get_args from facefusion.common_helper import create_float_range, create_int_range -from facefusion.types import Angle, ApiSecurityStrategy, AudioEncoder, AudioFormat, AudioSet, BenchmarkMode, BenchmarkResolution, BenchmarkSet, DownloadProvider, DownloadProviderSet, DownloadScope, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskArea, FaceMaskAreaSet, FaceMaskRegion, FaceMaskRegionSet, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorGender, FaceSelectorMode, FaceSelectorOrder, FaceSelectorRace, Gender, ImageEncoder, ImageFormat, ImageSet, JobStatus, LogLevel, LogLevelSet, Race, Score, TempFrameFormat, VideoEncoder, VideoFormat, VideoMemoryStrategy, VideoPreset, VideoSet, VoiceExtractorModel, WorkflowMode +from facefusion.types import Angle, ApiSecurityStrategy, AudioEncoder, AudioFormat, AudioSet, BenchmarkMode, BenchmarkResolution, BenchmarkSet, DownloadProvider, DownloadProviderSet, DownloadScope, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskArea, FaceMaskAreaSet, FaceMaskRegion, FaceMaskRegionSet, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorGender, FaceSelectorMode, FaceSelectorOrder, FaceSelectorRace, Gender, ImageEncoder, ImageFormat, ImageSet, JobStatus, LogLevel, LogLevelSet, Race, Score, TempFrameFormat, TempPixelFormat, VideoEncoder, VideoFormat, VideoMemoryStrategy, VideoPreset, VideoSet, VoiceExtractorModel, WorkflowMode, WorkflowStrategy face_detector_set : FaceDetectorSet =\ { @@ -48,6 +48,7 @@ face_mask_regions : List[FaceMaskRegion] = list(get_args(FaceMaskRegion)) voice_extractor_models : List[VoiceExtractorModel] = list(get_args(VoiceExtractorModel)) workflow_modes : List[WorkflowMode] = [ 'auto', 'audio-to-image:frames', 'audio-to-image:video', 'image-to-image', 'image-to-video', 'image-to-video:frames' ] +workflow_strategies : List[WorkflowStrategy] = list(get_args(WorkflowStrategy)) audio_set : AudioSet =\ { @@ -82,6 +83,7 @@ audio_formats : List[AudioFormat] = list(get_args(AudioFormat)) image_formats : List[ImageFormat] = list(get_args(ImageFormat)) video_formats : List[VideoFormat] = list(get_args(VideoFormat)) temp_frame_formats : List[TempFrameFormat] = list(get_args(TempFrameFormat)) +temp_pixel_formats : List[TempPixelFormat] = list(get_args(TempPixelFormat)) audio_encoders : List[AudioEncoder] = list(get_args(AudioEncoder)) image_encoders : List[ImageEncoder] = list(get_args(ImageEncoder)) diff --git a/facefusion/content_analyser.py b/facefusion/content_analyser.py index cf21f06f..6caf43fe 100644 --- a/facefusion/content_analyser.py +++ b/facefusion/content_analyser.py @@ -4,12 +4,12 @@ from typing import Tuple import numpy from tqdm import tqdm -from facefusion import inference_manager, state_manager, translator +from facefusion import inference_manager, state_manager, translator, video_manager from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url from facefusion.filesystem import resolve_relative_path from facefusion.thread_helper import conditional_thread_semaphore from facefusion.types import Detection, DownloadScope, DownloadSet, Fps, InferencePool, ModelSet, VisionFrame -from facefusion.vision import detect_video_fps, fit_contain_frame, read_image, read_video_frame +from facefusion.vision import detect_video_fps, fit_contain_frame, is_vision_frame, read_image STREAM_COUNTER = 0 @@ -158,17 +158,21 @@ def analyse_image(image_path : str) -> bool: def analyse_video(video_path : str, trim_frame_start : int, trim_frame_end : int) -> bool: video_fps = detect_video_fps(video_path) frame_range = range(trim_frame_start, trim_frame_end) + video_reader = video_manager.get_reader(video_path, 'analyse_video') rate = 0.0 total = 0 counter = 0 + if trim_frame_start > 0: + video_manager.seek_video_reader(video_reader, trim_frame_start) + with tqdm(total = len(frame_range), desc = translator.get('analysing'), unit = 'frame', ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: for frame_number in frame_range: - if frame_number % int(video_fps) == 0: - vision_frame = read_video_frame(video_path, frame_number) + vision_frame = video_manager.read_video_frame(video_reader) - if numpy.any(vision_frame): + if frame_number % int(video_fps) == 0: + if is_vision_frame(vision_frame): total += 1 if analyse_frame(vision_frame): diff --git a/facefusion/core.py b/facefusion/core.py index 22077c62..000d3b31 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -91,20 +91,17 @@ def pre_check() -> bool: logger.error(translator.get('python_not_supported').format(version = '3.10'), __name__) return False - if not shutil.which('curl'): - logger.error(translator.get('curl_not_installed'), __name__) - return False - - if not shutil.which('ffmpeg'): - logger.error(translator.get('ffmpeg_not_installed'), __name__) - return False + for dependency in [ 'curl', 'ffmpeg', 'ffprobe' ]: + if not shutil.which(dependency): + logger.error(translator.get('dependency_not_installed').format(dependency = dependency), __name__) + return False return True def common_pre_check() -> bool: content_analyser_content = inspect.getsource(content_analyser).encode() - return hash_helper.create_hash(content_analyser_content) == '975d67d6' + return hash_helper.create_hash(content_analyser_content) == '3c6ce25e' def processors_pre_check() -> bool: diff --git a/facefusion/execution.py b/facefusion/execution.py index 30cab3c0..8684023a 100644 --- a/facefusion/execution.py +++ b/facefusion/execution.py @@ -12,6 +12,16 @@ from facefusion.types import ExecutionProvider, InferenceOptionSet, InferencePro onnxruntime.set_default_logger_severity(3) +@lru_cache() +def get_onnxruntime_version() -> Tuple[int, int, int]: + version_split = onnxruntime.__version__.split('.') + major_version = int(version_split[0]) + minor_version = int(version_split[1]) + patch_version = int(version_split[2].split('+')[0]) + + return major_version, minor_version, patch_version + + def has_execution_provider(execution_provider : ExecutionProvider) -> bool: return execution_provider in get_available_execution_providers() diff --git a/facefusion/face_creator.py b/facefusion/face_creator.py index fb4a8a81..e347bf9f 100644 --- a/facefusion/face_creator.py +++ b/facefusion/face_creator.py @@ -10,6 +10,7 @@ from facefusion.face_helper import apply_nms, average_points, convert_to_face_la from facefusion.face_landmarker import detect_face_landmark, estimate_face_landmark_68_5 from facefusion.face_recognizer import calculate_face_embedding from facefusion.types import BoundingBox, Face, FaceLandmark5, FaceLandmarkSet, FaceScoreSet, Score, VisionFrame +from facefusion.vision import is_vision_frame def create_faces(vision_frame : VisionFrame, bounding_boxes : List[BoundingBox], face_scores : List[Score], face_landmarks_5 : List[FaceLandmark5]) -> List[Face]: @@ -73,7 +74,7 @@ def get_many_faces(vision_frames : List[VisionFrame]) -> List[Face]: many_faces : List[Face] = [] for vision_frame in vision_frames: - if numpy.any(vision_frame): + if is_vision_frame(vision_frame): all_bounding_boxes = [] all_face_scores = [] all_face_landmarks_5 = [] diff --git a/facefusion/face_store.py b/facefusion/face_store.py index 022210a6..8042735d 100644 --- a/facefusion/face_store.py +++ b/facefusion/face_store.py @@ -1,16 +1,15 @@ import threading from typing import List, Optional -import numpy - from facefusion.hash_helper import create_hash from facefusion.types import Face, FaceStore, VisionFrame +from facefusion.vision import is_vision_frame FACE_STORE : FaceStore = {} def get_faces(vision_frame : VisionFrame) -> Optional[List[Face]]: - if numpy.any(vision_frame): + if is_vision_frame(vision_frame): vision_hash = create_hash(vision_frame.tobytes()) if FACE_STORE.get(vision_hash): @@ -20,7 +19,7 @@ def get_faces(vision_frame : VisionFrame) -> Optional[List[Face]]: def set_faces(vision_frame : VisionFrame, faces : List[Face]) -> None: - if numpy.any(vision_frame): + if is_vision_frame(vision_frame): vision_hash = create_hash(vision_frame.tobytes()) FACE_STORE.setdefault(vision_hash, { @@ -29,7 +28,7 @@ def set_faces(vision_frame : VisionFrame, faces : List[Face]) -> None: def resolve_lock(vision_frame : VisionFrame) -> threading.Lock: - if numpy.any(vision_frame): + if is_vision_frame(vision_frame): vision_hash = create_hash(vision_frame.tobytes()) return FACE_STORE.setdefault(vision_hash, { diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py index 26d5a784..0b97c6da 100644 --- a/facefusion/ffmpeg.py +++ b/facefusion/ffmpeg.py @@ -7,11 +7,10 @@ from typing import List, Optional, cast from tqdm import tqdm import facefusion.choices -from facefusion import ffmpeg_builder, logger, process_manager, state_manager, translator +from facefusion import ffmpeg_builder, ffprobe, logger, process_manager, state_manager, translator, vision from facefusion.filesystem import get_file_format, remove_file from facefusion.temp_helper import get_temp_file_path, get_temp_frames_pattern -from facefusion.types import ApiSecurityStrategy, AudioEncoder, Buffer, Command, EncoderSet, Fps, Resolution, SampleRate, UpdateProgress, VideoEncoder, VideoFormat -from facefusion.vision import detect_video_duration, detect_video_fps, pack_resolution, predict_video_frame_total +from facefusion.types import ApiSecurityStrategy, AudioEncoder, Buffer, Command, EncoderSet, Fps, Resolution, SampleRate, UpdateProgress, VideoEncoder, VideoFormat, VideoReaderMetadata def run_ffmpeg_with_progress(commands : List[Command], update_progress : UpdateProgress) -> subprocess.Popen[Buffer]: @@ -74,7 +73,52 @@ def run_ffmpeg(commands : List[Command]) -> subprocess.Popen[Buffer]: def open_ffmpeg(commands : List[Command]) -> subprocess.Popen[Buffer]: commands = ffmpeg_builder.run(commands) - return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.PIPE) + return subprocess.Popen(commands, stdin = subprocess.PIPE, stderr = subprocess.DEVNULL, stdout = subprocess.PIPE) + + +def create_video_reader(video_path : str, frame_number : int, video_metadata : VideoReaderMetadata) -> subprocess.Popen[Buffer]: + commands = ffmpeg_builder.chain( + ffmpeg_builder.seek_to(frame_number / video_metadata.get('fps')), + ffmpeg_builder.set_input(video_path), + ffmpeg_builder.restrict_color_transfer(video_metadata.get('color_transfer')), + ffmpeg_builder.prevent_frame_drop(), + ffmpeg_builder.enforce_pixel_format('bgr24'), + ffmpeg_builder.set_output_format('rawvideo'), + ffmpeg_builder.cast_stream() + ) + + return open_ffmpeg(commands) + + +def create_video_writer(target_path : str, temp_video_fps : Fps, temp_video_resolution : Resolution, output_video_resolution : Resolution, output_video_fps : Fps) -> subprocess.Popen[Buffer]: + 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(state_manager.get_temp_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(state_manager.get_item('temp_pixel_format')), + ffmpeg_builder.set_media_resolution(vision.pack_resolution(temp_video_resolution)), + ffmpeg_builder.set_input_fps(temp_video_fps), + ffmpeg_builder.set_input('pipe:0'), + ffmpeg_builder.set_media_resolution(vision.pack_resolution(output_video_resolution)), + ffmpeg_builder.set_video_encoder(output_video_encoder), + ffmpeg_builder.set_thread_count(16), + 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.concat( + ffmpeg_builder.set_video_fps(output_video_fps), + ffmpeg_builder.convert_color_space('bt709') + ), + ffmpeg_builder.set_pixel_format(output_video_encoder), + ffmpeg_builder.force_output(temp_video_path) + ) + + return open_ffmpeg(commands) def log_debug(process : subprocess.Popen[Buffer]) -> None: @@ -123,14 +167,18 @@ def get_static_available_encoder_set() -> EncoderSet: def extract_frames(target_path : str, output_path : str, temp_video_resolution : Resolution, temp_video_fps : Fps, trim_frame_start : int, trim_frame_end : int) -> bool: - extract_frame_total = predict_video_frame_total(target_path, temp_video_fps, trim_frame_start, trim_frame_end) + color_transfer = ffprobe.extract_static_video_metadata(target_path).get('color_transfer') + extract_frame_total = vision.predict_video_frame_total(target_path, temp_video_fps, trim_frame_start, trim_frame_end) temp_frames_pattern = get_temp_frames_pattern(state_manager.get_temp_path(), output_path, state_manager.get_item('temp_frame_format'), '%08d') commands = ffmpeg_builder.chain( ffmpeg_builder.set_input(target_path), - ffmpeg_builder.set_media_resolution(pack_resolution(temp_video_resolution)), + ffmpeg_builder.set_media_resolution(vision.pack_resolution(temp_video_resolution)), ffmpeg_builder.set_frame_quality(0), ffmpeg_builder.enforce_pixel_format('rgb24'), - ffmpeg_builder.select_frame_range(trim_frame_start, trim_frame_end, temp_video_fps), + ffmpeg_builder.concat( + ffmpeg_builder.select_frame_range(trim_frame_start, trim_frame_end, temp_video_fps), + ffmpeg_builder.restrict_color_transfer(color_transfer) + ), ffmpeg_builder.prevent_frame_drop(), ffmpeg_builder.set_start_number(trim_frame_start), ffmpeg_builder.set_output(temp_frames_pattern) @@ -150,7 +198,7 @@ def spawn_frames(target_path : str, output_path : str, temp_video_resolution : R ffmpeg_builder.set_input(target_path), ffmpeg_builder.set_video_duration(duration), ffmpeg_builder.set_video_fps(temp_video_fps), - ffmpeg_builder.set_media_resolution(pack_resolution(temp_video_resolution)), + ffmpeg_builder.set_media_resolution(vision.pack_resolution(temp_video_resolution)), ffmpeg_builder.set_output(temp_frames_pattern) ) @@ -163,7 +211,7 @@ def copy_image(target_path : str, output_path : str, temp_image_resolution : Res temp_image_path = get_temp_file_path(state_manager.get_temp_path(), output_path) commands = ffmpeg_builder.chain( ffmpeg_builder.set_input(target_path), - ffmpeg_builder.set_media_resolution(pack_resolution(temp_image_resolution)), + ffmpeg_builder.set_media_resolution(vision.pack_resolution(temp_image_resolution)), ffmpeg_builder.set_image_quality(target_path, 100), ffmpeg_builder.force_output(temp_image_path) ) @@ -175,7 +223,7 @@ def finalize_image(output_path : str, output_image_resolution : Resolution) -> b temp_image_path = get_temp_file_path(state_manager.get_temp_path(), output_path) commands = ffmpeg_builder.chain( ffmpeg_builder.set_input(temp_image_path), - ffmpeg_builder.set_media_resolution(pack_resolution(output_image_resolution)), + ffmpeg_builder.set_media_resolution(vision.pack_resolution(output_image_resolution)), ffmpeg_builder.set_image_quality(output_path, output_image_quality), ffmpeg_builder.force_output(output_path) ) @@ -204,10 +252,10 @@ def restore_audio(target_path : str, output_path : str, trim_frame_start : int, output_audio_encoder = state_manager.get_item('output_audio_encoder') output_audio_quality = state_manager.get_item('output_audio_quality') output_audio_volume = state_manager.get_item('output_audio_volume') - target_video_fps = detect_video_fps(target_path) + target_video_fps = vision.detect_video_fps(target_path) temp_video_path = get_temp_file_path(state_manager.get_temp_path(), output_path) temp_video_format = cast(VideoFormat, get_file_format(output_path)) - temp_video_duration = detect_video_duration(temp_video_path) + temp_video_duration = vision.detect_video_duration(temp_video_path) output_video_format = cast(VideoFormat, get_file_format(output_path)) output_audio_encoder = fix_audio_encoder(temp_video_format, output_audio_encoder) @@ -234,7 +282,7 @@ def replace_audio(audio_path : str, output_path : str) -> bool: output_audio_volume = state_manager.get_item('output_audio_volume') temp_video_path = get_temp_file_path(state_manager.get_temp_path(), output_path) temp_video_format = cast(VideoFormat, get_file_format(output_path)) - temp_video_duration = detect_video_duration(temp_video_path) + temp_video_duration = vision.detect_video_duration(temp_video_path) output_video_format = cast(VideoFormat, get_file_format(output_path)) output_audio_encoder = fix_audio_encoder(temp_video_format, output_audio_encoder) @@ -256,7 +304,7 @@ def merge_video(target_path : str, output_path : str, temp_video_fps : Fps, outp 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') - merge_frame_total = predict_video_frame_total(target_path, output_video_fps, trim_frame_start, trim_frame_end) + merge_frame_total = vision.predict_video_frame_total(target_path, output_video_fps, trim_frame_start, trim_frame_end) temp_video_path = get_temp_file_path(state_manager.get_temp_path(), output_path) temp_video_format = cast(VideoFormat, get_file_format(output_path)) temp_frames_pattern = get_temp_frames_pattern(state_manager.get_temp_path(), output_path, state_manager.get_item('temp_frame_format'), '%08d') @@ -266,14 +314,15 @@ def merge_video(target_path : str, output_path : str, temp_video_fps : Fps, outp ffmpeg_builder.set_input_fps(temp_video_fps), ffmpeg_builder.set_start_number(trim_frame_start), ffmpeg_builder.set_input(temp_frames_pattern), - ffmpeg_builder.set_media_resolution(pack_resolution(output_video_resolution)), + ffmpeg_builder.set_media_resolution(vision.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.concat( ffmpeg_builder.set_video_fps(output_video_fps), - ffmpeg_builder.keep_video_alpha(output_video_encoder) + ffmpeg_builder.keep_video_alpha(output_video_encoder), + ffmpeg_builder.convert_color_space('bt709') ), ffmpeg_builder.set_pixel_format(output_video_encoder), ffmpeg_builder.force_output(temp_video_path) diff --git a/facefusion/ffmpeg_builder.py b/facefusion/ffmpeg_builder.py index 3c20ae01..4a1d17c3 100644 --- a/facefusion/ffmpeg_builder.py +++ b/facefusion/ffmpeg_builder.py @@ -5,7 +5,7 @@ from typing import List, Optional import numpy from facefusion.filesystem import get_file_format -from facefusion.types import AudioEncoder, Command, CommandSet, Duration, Fps, SampleRate, StreamMode, VideoEncoder, VideoFormat, VideoPreset +from facefusion.types import AudioEncoder, ColorSpace, ColorTransfer, Command, CommandSet, Duration, Fps, SampleRate, StreamMode, VideoEncoder, VideoFormat, VideoPreset def run(commands : List[Command]) -> List[Command]: @@ -47,6 +47,10 @@ def set_input(input_path : str) -> List[Command]: return [ '-i', input_path ] +def seek_to(time : float) -> List[Command]: + return [ '-ss', str(time) ] + + def set_input_fps(input_fps : Fps) -> List[Command]: return [ '-r', str(input_fps) ] @@ -63,6 +67,14 @@ def force_output(output_path : str) -> List[Command]: return [ '-y', output_path ] +def set_output_format(output_format : str) -> List[Command]: + return [ '-f', output_format ] + + +def set_thread_count(thread_count : int) -> List[Command]: + return [ '-threads', str(thread_count) ] + + def set_loop() -> List[Command]: return [ '-loop', '1' ] @@ -103,6 +115,16 @@ def set_pixel_format(video_encoder : VideoEncoder) -> List[Command]: return [ '-pix_fmt', 'yuv420p' ] +def restrict_color_transfer(color_transfer : ColorTransfer) -> List[Command]: + if color_transfer in [ 'smpte2084', 'arib-std-b67' ]: + return [ '-vf', 'scale=out_primaries=bt709:out_transfer=bt709:intent=perceptual' ] + return [] + + +def convert_color_space(color_space : ColorSpace) -> List[Command]: + return [ '-vf', 'scale=out_color_matrix=' + color_space + ':out_range=tv,setparams=colorspace=' + color_space + ':color_primaries=' + color_space + ':color_trc=' + color_space ] + + def set_frame_quality(frame_quality : int) -> List[Command]: return [ '-q:v', str(frame_quality) ] diff --git a/facefusion/ffprobe.py b/facefusion/ffprobe.py index 3d01365c..c0cf7191 100644 --- a/facefusion/ffprobe.py +++ b/facefusion/ffprobe.py @@ -1,4 +1,5 @@ import subprocess +from functools import lru_cache from typing import Dict, List from facefusion import ffprobe_builder @@ -10,16 +11,9 @@ def run_ffprobe(commands : List[Command]) -> subprocess.Popen[Buffer]: return subprocess.Popen(commands, stderr = subprocess.PIPE, stdout = subprocess.PIPE) -def probe_entries(media_path : str, entries : List[str]) -> Dict[str, str]: +def parse_entries(output : Buffer) -> Dict[str, str]: media_entries = {} - commands = ffprobe_builder.chain( - ffprobe_builder.show_entries(entries), - ffprobe_builder.format_to_key_value(), - ffprobe_builder.set_input(media_path) - ) - output, _ = run_ffprobe(commands).communicate() - if output: lines = output.decode().strip().splitlines() @@ -31,14 +25,58 @@ def probe_entries(media_path : str, entries : List[str]) -> Dict[str, str]: return media_entries -def extract_audio_metadata(audio_path : str) -> AudioMetadata: - audio_entries = probe_entries(audio_path, [ 'duration', 'sample_rate', 'channels', 'bit_rate' ]) +def probe_audio_entries(audio_path : str, entries : List[str]) -> Dict[str, str]: + commands = ffprobe_builder.chain( + ffprobe_builder.select_stream('a:0'), + ffprobe_builder.show_stream_entries(entries), + ffprobe_builder.format_to_key_value(), + ffprobe_builder.set_input(audio_path) + ) - duration = float(audio_entries.get('duration')) + output, _ = run_ffprobe(commands).communicate() + + return parse_entries(output) + + +def probe_video_entries(video_path : str, entries : List[str]) -> Dict[str, str]: + commands = ffprobe_builder.chain( + ffprobe_builder.select_stream('v:0'), + ffprobe_builder.show_stream_entries(entries), + ffprobe_builder.format_to_key_value(), + ffprobe_builder.set_input(video_path) + ) + + output, _ = run_ffprobe(commands).communicate() + + return parse_entries(output) + + +def probe_format_entries(media_path : str, entries : List[str]) -> Dict[str, str]: + commands = ffprobe_builder.chain( + ffprobe_builder.show_format_entries(entries), + ffprobe_builder.format_to_key_value(), + ffprobe_builder.set_input(media_path) + ) + + output, _ = run_ffprobe(commands).communicate() + + return parse_entries(output) + + +@lru_cache(maxsize = 128) +def extract_static_audio_metadata(audio_path : str) -> AudioMetadata: + return extract_audio_metadata(audio_path) + + +def extract_audio_metadata(audio_path : str) -> AudioMetadata: + audio_entries = probe_audio_entries(audio_path, [ 'sample_rate', 'channels' ]) + format_entries = probe_format_entries(audio_path, [ 'duration', 'bit_rate' ]) + + duration = float(format_entries.get('duration')) sample_rate = int(audio_entries.get('sample_rate')) - frame_total = int(duration * sample_rate) + frame_total = round(duration * sample_rate) channel_total = int(audio_entries.get('channels')) - bit_rate = int(audio_entries.get('bit_rate')) + bit_rate = int(format_entries.get('bit_rate')) audio_metadata : AudioMetadata =\ { @@ -52,15 +90,22 @@ def extract_audio_metadata(audio_path : str) -> AudioMetadata: return audio_metadata -def extract_video_metadata(video_path : str) -> VideoMetadata: - video_entries = probe_entries(video_path, [ 'duration', 'width', 'height', 'r_frame_rate', 'bit_rate' ]) +@lru_cache(maxsize = 128) +def extract_static_video_metadata(video_path : str) -> VideoMetadata: + return extract_video_metadata(video_path) - duration = float(video_entries.get('duration')) + +def extract_video_metadata(video_path : str) -> VideoMetadata: + video_entries = probe_video_entries(video_path, [ 'width', 'height', 'r_frame_rate', 'color_transfer' ]) + format_entries = probe_format_entries(video_path, [ 'duration', 'bit_rate' ]) + + duration = float(format_entries.get('duration')) fps = extract_video_fps(video_entries.get('r_frame_rate')) - frame_total = int(duration * fps) + frame_total = round(duration * fps) width = int(video_entries.get('width')) height = int(video_entries.get('height')) - bit_rate = int(video_entries.get('bit_rate')) + bit_rate = int(format_entries.get('bit_rate')) + color_transfer = video_entries.get('color_transfer', 'unknown') video_metadata : VideoMetadata =\ { @@ -68,7 +113,8 @@ def extract_video_metadata(video_path : str) -> VideoMetadata: 'frame_total' : frame_total, 'fps' : fps, 'resolution' : (width, height), - 'bit_rate' : bit_rate + 'bit_rate' : bit_rate, + 'color_transfer' : color_transfer } return video_metadata diff --git a/facefusion/ffprobe_builder.py b/facefusion/ffprobe_builder.py index f1cbccc2..22e7d146 100644 --- a/facefusion/ffprobe_builder.py +++ b/facefusion/ffprobe_builder.py @@ -13,10 +13,18 @@ def chain(*commands : List[Command]) -> List[Command]: return list(itertools.chain(*commands)) -def show_entries(entries : List[str]) -> List[Command]: +def select_stream(stream : str) -> List[Command]: + return [ '-select_streams', stream ] + + +def show_stream_entries(entries : List[str]) -> List[Command]: return [ '-show_entries', 'stream=' + ','.join(entries) ] +def show_format_entries(entries : List[str]) -> List[Command]: + return [ '-show_entries', 'format=' + ','.join(entries) ] + + def format_to_key_value() -> List[Command]: return [ '-of', 'default=noprint_wrappers=1' ] diff --git a/facefusion/frame_store.py b/facefusion/frame_store.py new file mode 100644 index 00000000..843ce21f --- /dev/null +++ b/facefusion/frame_store.py @@ -0,0 +1,35 @@ +from facefusion.types import FrameStoreSet, VisionFrame, VisionFrameSet + +FRAME_STORE_SET : FrameStoreSet = {} + + +def get_frame_store(id : str) -> VisionFrameSet: + if id not in FRAME_STORE_SET: + FRAME_STORE_SET[id] = {} + + return FRAME_STORE_SET.get(id) + + +def set_frame(id : str, frame_number : int, vision_frame : VisionFrame) -> None: + frame_store = get_frame_store(id) + frame_store[frame_number] = vision_frame + + +def select_frame_set(id : str, frame_start : int, frame_end : int) -> VisionFrameSet: + frame_store = get_frame_store(id) + frame_set = {} + + for frame_number in range(frame_start, frame_end + 1): + if frame_number in frame_store: + frame_set[frame_number] = frame_store.get(frame_number) + + return frame_set + + +def reduce_frames(id : str, frame_min : int, frame_max : int) -> None: + FRAME_STORE_SET[id] = select_frame_set(id, frame_min, frame_max) + + +def clear_frames(id : str) -> None: + if id in FRAME_STORE_SET: + del FRAME_STORE_SET[id] diff --git a/facefusion/inference_manager.py b/facefusion/inference_manager.py index 0741bf20..57210cf5 100644 --- a/facefusion/inference_manager.py +++ b/facefusion/inference_manager.py @@ -9,7 +9,7 @@ from onnxruntime import InferenceSession from facefusion import logger, process_manager, state_manager, translator from facefusion.app_context import detect_app_context from facefusion.common_helper import is_windows -from facefusion.execution import create_inference_providers, has_execution_provider +from facefusion.execution import create_inference_providers, get_onnxruntime_version, has_execution_provider from facefusion.exit_helper import fatal_exit from facefusion.filesystem import get_file_name, is_file from facefusion.time_helper import calculate_end_time @@ -25,17 +25,21 @@ INFERENCE_POOL_SET : InferencePoolSet =\ def get_inference_pool(module_name : str, model_names : List[str], model_source_set : DownloadSet) -> InferencePool: while process_manager.is_checking(): sleep(0.5) + execution_device_ids = state_manager.get_item('execution_device_ids') execution_providers = state_manager.get_item('execution_providers') + has_arena_leak = has_execution_provider('cuda') and get_onnxruntime_version() > (1, 24, 4) app_context = detect_app_context() for execution_device_id in execution_device_ids: inference_context = get_inference_context(module_name, model_names, execution_device_id, execution_providers) - if app_context == 'cli' and INFERENCE_POOL_SET.get('api').get(inference_context): - INFERENCE_POOL_SET['cli'][inference_context] = INFERENCE_POOL_SET.get('api').get(inference_context) - if app_context == 'api' and INFERENCE_POOL_SET.get('cli').get(inference_context): - INFERENCE_POOL_SET['api'][inference_context] = INFERENCE_POOL_SET.get('cli').get(inference_context) + if not has_arena_leak: + if app_context == 'cli' and INFERENCE_POOL_SET.get('api').get(inference_context): + INFERENCE_POOL_SET['cli'][inference_context] = INFERENCE_POOL_SET.get('api').get(inference_context) + if app_context == 'api' and INFERENCE_POOL_SET.get('cli').get(inference_context): + INFERENCE_POOL_SET['api'][inference_context] = INFERENCE_POOL_SET.get('cli').get(inference_context) + if not INFERENCE_POOL_SET.get(app_context).get(inference_context): inference_providers = resolve_static_inference_providers(module_name, execution_device_id) INFERENCE_POOL_SET[app_context][inference_context] = create_inference_pool(model_source_set, inference_providers) @@ -49,6 +53,7 @@ def create_inference_pool(model_source_set : DownloadSet, inference_providers : for model_name in model_source_set.keys(): model_path = model_source_set.get(model_name).get('path') + if is_file(model_path): inference_pool[model_name] = create_inference_session(model_path, inference_providers) @@ -65,6 +70,7 @@ def clear_inference_pool(module_name : str, model_names : List[str]) -> None: for execution_device_id in execution_device_ids: inference_context = get_inference_context(module_name, model_names, execution_device_id, execution_providers) + if INFERENCE_POOL_SET.get(app_context).get(inference_context): del INFERENCE_POOL_SET[app_context][inference_context] @@ -93,10 +99,23 @@ def resolve_static_inference_providers(module_name : str, execution_device_id : module = importlib.import_module(module_name) execution_providers = state_manager.get_item('execution_providers') - if hasattr(module, 'resolve_inference_providers'): - inference_providers = getattr(module, 'resolve_inference_providers')() + if hasattr(module, 'override_inference_providers'): + override_inference_providers = getattr(module, 'override_inference_providers')() + + if override_inference_providers: + return override_inference_providers + + if hasattr(module, 'adjust_inference_providers'): + adjust_inference_providers = getattr(module, 'adjust_inference_providers')() + + if adjust_inference_providers: + inference_providers = create_inference_providers(execution_device_id, execution_providers) + + for adjust_inference_provider in adjust_inference_providers: + for inference_provider in inference_providers: + if inference_provider[0] == adjust_inference_provider[0] and inference_provider[1]: + inference_provider[1].update(adjust_inference_provider[1]) - if inference_providers: return inference_providers return create_inference_providers(execution_device_id, execution_providers) diff --git a/facefusion/installer.py b/facefusion/installer.py index 7bbcaa7c..36f39fb5 100644 --- a/facefusion/installer.py +++ b/facefusion/installer.py @@ -19,17 +19,18 @@ LOCALES =\ } ONNXRUNTIME_SET =\ { - 'default': ('onnxruntime', '1.26.0') + 'default': ('onnxruntime', '1.28.0') } if is_windows() or is_linux(): - ONNXRUNTIME_SET['cuda'] = ('onnxruntime-gpu', '1.26.0') + ONNXRUNTIME_SET['cuda@12'] = ('onnxruntime-gpu', '1.24.4') + ONNXRUNTIME_SET['cuda@13'] = ('onnxruntime-gpu', '1.28.0') ONNXRUNTIME_SET['openvino'] = ('onnxruntime-openvino', '1.24.1') if is_windows(): ONNXRUNTIME_SET['directml'] = ('onnxruntime-directml', '1.24.4') - ONNXRUNTIME_SET['qnn'] = ('onnxruntime-qnn', '1.24.4') + ONNXRUNTIME_SET['qnn'] = ('onnxruntime-qnn', '2.4.0') if is_linux(): - ONNXRUNTIME_SET['migraphx'] = ('onnxruntime-migraphx', '1.25.0') - ONNXRUNTIME_SET['rocm'] = ('onnxruntime-rocm', '1.22.2.post1') + ONNXRUNTIME_SET['migraphx'] = ('onnxruntime-migraphx', '1.27.1') + ONNXRUNTIME_SET['rocm'] = ('onnxruntime-rocm', '1.22.2.post3') def cli() -> None: @@ -54,6 +55,9 @@ def run(program : ArgumentParser) -> None: sys.stdout.write(LOCALES.get('conda_not_activated') + os.linesep) sys.exit(1) + for onnxruntime_package, _ in ONNXRUNTIME_SET.values(): + subprocess.call([ shutil.which('pip'), 'uninstall', onnxruntime_package, '-y', '-q' ], stderr = subprocess.DEVNULL) + commands = [ shutil.which('pip'), 'install' ] if args.force_reinstall: @@ -63,12 +67,10 @@ def run(program : ArgumentParser) -> None: for line in file.readlines(): __line__ = line.strip() + if not __line__.startswith('onnxruntime'): commands.append(__line__) onnxruntime_name, onnxruntime_version = ONNXRUNTIME_SET.get(args.onnxruntime) commands.append(onnxruntime_name + '==' + onnxruntime_version) - - subprocess.call([ shutil.which('pip'), 'uninstall', 'onnxruntime', onnxruntime_name, '-y', '-q' ]) - subprocess.call(commands) diff --git a/facefusion/locales.py b/facefusion/locales.py index 41531ea9..a3f11795 100644 --- a/facefusion/locales.py +++ b/facefusion/locales.py @@ -6,8 +6,7 @@ LOCALES : Locales =\ { 'conda_not_activated': 'conda is not activated', 'python_not_supported': 'python version is not supported, upgrade to {version} or higher', - 'curl_not_installed': 'curl is not installed', - 'ffmpeg_not_installed': 'ffmpeg is not installed', + 'dependency_not_installed': '{dependency} is not installed', 'creating_temp': 'creating temporary resources', 'extracting_frames': 'extracting frames with a resolution of {resolution} and {fps} frames per second', 'extracting_frames_succeeded': 'extracting frames succeeded', @@ -104,6 +103,7 @@ LOCALES : Locales =\ 'install_dependency': 'choose the variant of {dependency} to install', 'skip_conda': 'skip the conda environment check', 'workflow_mode': 'choose the workflow mode', + 'workflow_strategy': 'process the temporary frames in memory or on disk', 'config_path': 'choose the config file to override defaults', 'temp_path': 'specify the directory for the temporary resources', 'jobs_path': 'specify the directory to store jobs', @@ -141,6 +141,7 @@ LOCALES : Locales =\ 'trim_frame_start': 'specify the starting frame of the target video', 'trim_frame_end': 'specify the ending frame of the target video', 'temp_frame_format': 'specify the temporary resources format', + 'temp_pixel_format': 'specify the temporary pixel format', 'target_frame_amount': 'specify the amount of target frames forwarded to the processor', 'output_image_quality': 'specify the image quality which translates to the image compression', 'output_image_scale': 'specify the image scale based on the target image', diff --git a/facefusion/processors/modules/age_modifier/core.py b/facefusion/processors/modules/age_modifier/core.py index c70a93a9..85e90b1b 100755 --- a/facefusion/processors/modules/age_modifier/core.py +++ b/facefusion/processors/modules/age_modifier/core.py @@ -23,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import thread_semaphore from facefusion.types import Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import match_frame_color, read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import match_frame_color, read_static_image, read_static_video_frame @lru_cache() @@ -180,7 +180,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/background_remover/core.py b/facefusion/processors/modules/background_remover/core.py index d91105d1..76e322f4 100644 --- a/facefusion/processors/modules/background_remover/core.py +++ b/facefusion/processors/modules/background_remover/core.py @@ -22,7 +22,7 @@ from facefusion.program_helper import find_argument_group from facefusion.sanitizer import sanitize_int_range from facefusion.thread_helper import thread_semaphore from facefusion.types import Args, DownloadScope, InferencePool, InferenceProvider, Mask, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_frame @lru_cache() @@ -479,7 +479,7 @@ def clear_inference_pool() -> None: inference_manager.clear_inference_pool(__name__, model_names) -def resolve_inference_providers() -> List[InferenceProvider]: +def override_inference_providers() -> List[InferenceProvider]: model_type = get_model_options().get('type') if is_macos() and has_execution_provider('coreml') or is_windows() and has_execution_provider('directml') and model_type == 'corridor_key': @@ -558,7 +558,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/deep_swapper/core.py b/facefusion/processors/modules/deep_swapper/core.py index ba66721f..291826dc 100755 --- a/facefusion/processors/modules/deep_swapper/core.py +++ b/facefusion/processors/modules/deep_swapper/core.py @@ -23,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import thread_semaphore from facefusion.types import Args, DownloadScope, Face, InferencePool, Mask, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import conditional_match_frame_color, read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import conditional_match_frame_color, read_static_image, read_static_video_frame @lru_cache() @@ -334,7 +334,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/expression_restorer/core.py b/facefusion/processors/modules/expression_restorer/core.py index 9b7e464f..56785f95 100755 --- a/facefusion/processors/modules/expression_restorer/core.py +++ b/facefusion/processors/modules/expression_restorer/core.py @@ -23,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, LivePortraitExpression, from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore, thread_semaphore from facefusion.types import Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_frame @lru_cache() @@ -167,7 +167,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/face_debugger/core.py b/facefusion/processors/modules/face_debugger/core.py index 09b5b671..62f8b3c5 100755 --- a/facefusion/processors/modules/face_debugger/core.py +++ b/facefusion/processors/modules/face_debugger/core.py @@ -19,7 +19,7 @@ from facefusion.processors.modules.face_debugger.types import FaceDebuggerInputs from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.types import Args, Face, InferencePool, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_frame def get_inference_pool() -> InferencePool: @@ -77,7 +77,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') == 'strict': diff --git a/facefusion/processors/modules/face_editor/core.py b/facefusion/processors/modules/face_editor/core.py index d20c2f81..56899bd0 100755 --- a/facefusion/processors/modules/face_editor/core.py +++ b/facefusion/processors/modules/face_editor/core.py @@ -23,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, LivePortraitExpression, from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore, thread_semaphore from facefusion.types import Args, DownloadScope, Face, FaceLandmark68, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_frame @lru_cache() @@ -302,7 +302,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/face_enhancer/core.py b/facefusion/processors/modules/face_enhancer/core.py index fbcb3552..e578c350 100755 --- a/facefusion/processors/modules/face_enhancer/core.py +++ b/facefusion/processors/modules/face_enhancer/core.py @@ -21,7 +21,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import thread_semaphore from facefusion.types import Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import blend_frame, read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import blend_frame, read_static_image, read_static_video_frame @lru_cache() @@ -358,7 +358,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/face_swapper/core.py b/facefusion/processors/modules/face_swapper/core.py index 653becf3..1b19183e 100755 --- a/facefusion/processors/modules/face_swapper/core.py +++ b/facefusion/processors/modules/face_swapper/core.py @@ -26,7 +26,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore from facefusion.types import Args, DownloadScope, Embedding, Face, InferencePool, InferenceProvider, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_images, read_static_video_chunk, read_static_video_frame, unpack_resolution +from facefusion.vision import read_static_image, read_static_images, read_static_video_frame, unpack_resolution @lru_cache() @@ -502,20 +502,19 @@ def clear_inference_pool() -> None: inference_manager.clear_inference_pool(__name__, model_names) -def resolve_inference_providers() -> List[InferenceProvider]: +def adjust_inference_providers() -> List[InferenceProvider]: model_precision = get_model_options().get('precision') - model_type = get_model_options().get('type') + workflow_mode = state_manager.get_item('workflow_mode') - if is_macos() and has_execution_provider('coreml'): - if model_type in [ 'ghost', 'uniface' ] or model_precision == 'fp16': - return\ - [ - (facefusion.choices.execution_provider_set.get('coreml'), - { - 'ModelFormat': 'MLProgram', - 'SpecializationStrategy': 'FastPrediction' - }) - ] + if is_macos() and has_execution_provider('coreml') and model_precision == 'fp16' and workflow_mode == 'image-to-video': + return\ + [ + (facefusion.choices.execution_provider_set.get('coreml'), + { + 'ModelFormat': 'MLProgram', + 'MLComputeUnits': 'CPUAndGPU' + }) + ] return [] @@ -610,7 +609,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/frame_colorizer/core.py b/facefusion/processors/modules/frame_colorizer/core.py index e903709c..50225a5e 100644 --- a/facefusion/processors/modules/frame_colorizer/core.py +++ b/facefusion/processors/modules/frame_colorizer/core.py @@ -20,7 +20,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import thread_semaphore from facefusion.types import Args, DownloadScope, InferencePool, InferenceProvider, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import blend_frame, read_static_image, read_static_video_chunk, read_static_video_frame, unpack_resolution +from facefusion.vision import blend_frame, read_static_image, read_static_video_frame, unpack_resolution @lru_cache() @@ -172,7 +172,7 @@ def clear_inference_pool() -> None: inference_manager.clear_inference_pool(__name__, model_names) -def resolve_inference_providers() -> List[InferenceProvider]: +def override_inference_providers() -> List[InferenceProvider]: if is_macos() and has_execution_provider('coreml'): return [ facefusion.choices.execution_provider_set.get('cpu') ] @@ -250,7 +250,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/frame_enhancer/core.py b/facefusion/processors/modules/frame_enhancer/core.py index 22767275..5517e2f9 100644 --- a/facefusion/processors/modules/frame_enhancer/core.py +++ b/facefusion/processors/modules/frame_enhancer/core.py @@ -20,7 +20,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore from facefusion.types import Args, DownloadScope, InferencePool, InferenceProvider, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import blend_frame, create_tile_frames, merge_tile_frames, read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import blend_frame, create_tile_frames, merge_tile_frames, read_static_image, read_static_video_frame @lru_cache() @@ -558,7 +558,7 @@ def clear_inference_pool() -> None: inference_manager.clear_inference_pool(__name__, model_names) -def resolve_inference_providers() -> List[InferenceProvider]: +def adjust_inference_providers() -> List[InferenceProvider]: model_precision = get_model_options().get('precision') if is_macos() and has_execution_provider('coreml') and model_precision == 'fp16': @@ -566,8 +566,7 @@ def resolve_inference_providers() -> List[InferenceProvider]: [ (facefusion.choices.execution_provider_set.get('coreml'), { - 'ModelFormat': 'MLProgram', - 'SpecializationStrategy': 'FastPrediction' + 'ModelFormat': 'MLProgram' }) ] @@ -637,7 +636,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() video_manager.clear_video_pool() if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: diff --git a/facefusion/processors/modules/frame_enhancer/types.py b/facefusion/processors/modules/frame_enhancer/types.py index d593b5fb..fd1ebb0b 100644 --- a/facefusion/processors/modules/frame_enhancer/types.py +++ b/facefusion/processors/modules/frame_enhancer/types.py @@ -9,4 +9,4 @@ FrameEnhancerInputs = TypedDict('FrameEnhancerInputs', 'temp_vision_mask' : Mask }) -FrameEnhancerModel = Literal['clear_reality_x4', 'face_dat_x4', 'lsdir_x4', 'nomos8k_sc_x4', 'real_esrgan_x2', 'real_esrgan_x2_fp16', 'real_esrgan_x4', 'real_esrgan_x4_fp16', 'real_esrgan_x8', 'real_esrgan_x8_fp16', 'real_hatgan_x4', 'real_web_photo_x4', 'realistic_rescaler_x4', 'remacri_x4', 'siax_x4', 'span_kendata_x4', 'swin2_sr_x4', 'tghq_face_x8', 'ultra_sharp_x4', 'ultra_sharp_2_x4'] +FrameEnhancerModel = Literal['clear_reality_x4', 'face_dat_x4', 'nomos8k_sc_x4', 'real_esrgan_x2', 'real_esrgan_x2_fp16', 'real_esrgan_x4', 'real_esrgan_x4_fp16', 'real_esrgan_x8', 'real_esrgan_x8_fp16', 'real_hatgan_x4', 'real_web_photo_x4', 'realistic_rescaler_x4', 'remacri_x4', 'siax_x4', 'span_kendata_x4', 'swin2_sr_x4', 'tghq_face_x8', 'ultra_sharp_x4', 'ultra_sharp_2_x4'] diff --git a/facefusion/processors/modules/lip_syncer/core.py b/facefusion/processors/modules/lip_syncer/core.py index 227f4b71..53d268e4 100755 --- a/facefusion/processors/modules/lip_syncer/core.py +++ b/facefusion/processors/modules/lip_syncer/core.py @@ -23,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore from facefusion.types import Args, AudioFrame, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_frame @lru_cache() @@ -185,7 +185,6 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() - read_static_video_chunk.cache_clear() read_static_voice.cache_clear() video_manager.clear_video_pool() diff --git a/facefusion/program.py b/facefusion/program.py index 0eab8db4..2ebd67e2 100755 --- a/facefusion/program.py +++ b/facefusion/program.py @@ -41,11 +41,11 @@ def create_config_path_program() -> ArgumentParser: def create_workflow_program() -> ArgumentParser: program = ArgumentParser(add_help = False) - group_paths = program.add_argument_group('paths') + group_workflow = program.add_argument_group('workflow') capability_store.register_capability_set( [ - group_paths.add_argument( + group_workflow.add_argument( '--workflow-mode', help = translator.get('help.workflow_mode'), default = config.get_str_value('workflow', 'workflow_mode', 'auto'), @@ -54,6 +54,17 @@ def create_workflow_program() -> ArgumentParser: ], scopes = [ 'api', 'cli' ] ) + capability_store.register_capability_set( + [ + group_workflow.add_argument( + '--workflow-strategy', + help = translator.get('help.workflow_strategy'), + default = config.get_str_value('workflow', 'workflow_strategy', 'memory'), + choices = facefusion.choices.workflow_strategies + ) + ], + scopes = [ 'api', 'cli' ] + ) return program @@ -592,6 +603,17 @@ def create_frame_extraction_program() -> ArgumentParser: ], scopes = [ 'api', 'cli' ] ) + capability_store.register_capability_set( + [ + group_frame_extraction.add_argument( + '--temp-pixel-format', + help = translator.get('help.temp_pixel_format'), + default = config.get_str_value('frame_extraction', 'temp_pixel_format', 'bgr24'), + choices = facefusion.choices.temp_pixel_formats + ) + ], + scopes = [ 'api', 'cli' ] + ) return program @@ -1033,6 +1055,7 @@ def collect_step_program() -> ArgumentParser: create_frame_extraction_program(), create_frame_distribution_program(), create_output_creation_program(), + create_workflow_program(), create_processors_program() ], add_help = False @@ -1043,7 +1066,6 @@ def collect_job_program() -> ArgumentParser: return ArgumentParser( parents = [ - create_workflow_program(), create_execution_program(), create_download_providers_program(), create_memory_program(), @@ -1205,7 +1227,6 @@ def create_program() -> ArgumentParser: parents = [ create_job_id_program(), - create_workflow_program(), create_config_path_program(), create_jobs_path_program(), create_source_paths_program(), @@ -1222,7 +1243,6 @@ def create_program() -> ArgumentParser: parents = [ create_job_id_program(), - create_workflow_program(), create_step_index_program(), create_config_path_program(), create_jobs_path_program(), @@ -1239,7 +1259,6 @@ def create_program() -> ArgumentParser: parents = [ create_job_id_program(), - create_workflow_program(), create_step_index_program(), create_config_path_program(), create_jobs_path_program(), diff --git a/facefusion/streamer.py b/facefusion/streamer.py index c663d9f8..ff58c03c 100644 --- a/facefusion/streamer.py +++ b/facefusion/streamer.py @@ -2,10 +2,9 @@ import os import subprocess from collections import deque from concurrent.futures import ThreadPoolExecutor -from typing import Deque, Iterator +from typing import Deque, Iterator, List import cv2 -import numpy from tqdm import tqdm from facefusion import ffmpeg_builder, logger, state_manager, translator @@ -15,11 +14,12 @@ from facefusion.ffmpeg import open_ffmpeg from facefusion.filesystem import is_directory from facefusion.processors.core import get_processors_modules from facefusion.types import Buffer, Fps, StreamMode, VisionFrame -from facefusion.vision import extract_vision_mask, read_static_images +from facefusion.vision import extract_vision_mask, is_vision_frame, read_static_images def multi_process_capture(camera_capture : cv2.VideoCapture, camera_fps : Fps) -> Iterator[VisionFrame]: capture_deque : Deque[VisionFrame] = deque() + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) with tqdm(desc = translator.get('streaming'), unit = 'frame', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: with ThreadPoolExecutor(max_workers = state_manager.get_item('execution_thread_count')) as executor: @@ -30,8 +30,8 @@ def multi_process_capture(camera_capture : cv2.VideoCapture, camera_fps : Fps) - if analyse_stream(capture_vision_frame, camera_fps): camera_capture.release() - if numpy.any(capture_vision_frame): - future = executor.submit(process_stream_frame, capture_vision_frame) + if is_vision_frame(capture_vision_frame): + future = executor.submit(process_stream_frame, source_vision_frames, capture_vision_frame) futures.append(future) for future_done in [ future for future in futures if future.done() ]: @@ -44,8 +44,7 @@ def multi_process_capture(camera_capture : cv2.VideoCapture, camera_fps : Fps) - yield capture_deque.popleft() -def process_stream_frame(target_vision_frame : VisionFrame) -> VisionFrame: - source_vision_frames = read_static_images(state_manager.get_item('source_paths')) +def process_stream_frame(source_vision_frames : List[VisionFrame], target_vision_frame : VisionFrame) -> VisionFrame: source_audio_frame = create_empty_audio_frame() source_voice_frame = create_empty_audio_frame() temp_vision_frame = target_vision_frame.copy() diff --git a/facefusion/temp_helper.py b/facefusion/temp_helper.py index a5dd6838..44bb042f 100644 --- a/facefusion/temp_helper.py +++ b/facefusion/temp_helper.py @@ -2,6 +2,7 @@ import os from typing import List from facefusion.filesystem import create_directory, get_file_extension, get_file_name, move_file, remove_directory, resolve_file_pattern +from facefusion.types import FrameSet def get_temp_file_path(temp_path : str, output_path : str) -> str: @@ -20,6 +21,17 @@ def resolve_temp_frame_paths(temp_path : str, output_path : str, temp_frame_form return resolve_file_pattern(temp_frames_pattern) +def resolve_temp_frame_set(temp_path : str, output_path : str, temp_frame_format : str) -> FrameSet: + temp_frames_pattern = get_temp_frames_pattern(temp_path, output_path, temp_frame_format, '*') + temp_frame_set = {} + + for temp_frame_path in resolve_file_pattern(temp_frames_pattern): + frame_number = int(get_file_name(temp_frame_path)) + temp_frame_set[frame_number] = temp_frame_path + + return temp_frame_set + + def get_temp_frames_pattern(temp_path : str, output_path : str, temp_frame_format : str, temp_frame_prefix : str) -> str: temp_directory_path = get_temp_directory_path(temp_path, output_path) return os.path.join(temp_directory_path, temp_frame_prefix + '.' + temp_frame_format) diff --git a/facefusion/types.py b/facefusion/types.py index aa3f1c58..a47ded9c 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -1,4 +1,5 @@ import ctypes +import subprocess from collections import namedtuple from datetime import datetime from threading import Lock @@ -66,22 +67,19 @@ Locales : TypeAlias = Dict[Language, Dict[str, Any]] LocalePoolSet : TypeAlias = Dict[str, Locales] WorkflowMode = Literal['auto', 'audio-to-image:frames', 'audio-to-image:video', 'image-to-image', 'image-to-video', 'image-to-video:frames'] +WorkflowStrategy = Literal['disk', 'memory'] -VideoCaptureSet : TypeAlias = Dict[str, cv2.VideoCapture] -VideoWriterSet : TypeAlias = Dict[str, cv2.VideoWriter] CameraCaptureSet : TypeAlias = Dict[str, cv2.VideoCapture] -VideoPoolSet = TypedDict('VideoPoolSet', -{ - 'capture' : VideoCaptureSet, - 'writer' : VideoWriterSet -}) CameraPoolSet = TypedDict('CameraPoolSet', { 'capture' : CameraCaptureSet }) ColorMode = Literal['rgb', 'rgba'] +ColorSpace = Literal['bt601', 'bt709', 'bt2020'] +ColorTransfer : TypeAlias = str VisionFrame : TypeAlias = NDArray[Any] +VisionFrameSet : TypeAlias = Dict[int, VisionFrame] Mask : TypeAlias = NDArray[Any] Points : TypeAlias = NDArray[Any] Distance : TypeAlias = NDArray[Any] @@ -197,6 +195,7 @@ AudioFormat = Literal['flac', 'm4a', 'mp3', 'ogg', 'opus', 'wav'] ImageFormat = Literal['bmp', 'jpeg', 'png', 'tiff', 'webp'] VideoFormat = Literal['avi', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mxf', 'webm', 'wmv'] TempFrameFormat = Literal['bmp', 'jpeg', 'png', 'tiff'] +TempPixelFormat = Literal['bgr24', 'bgra'] FrameSet : TypeAlias = Dict[int, str] @@ -235,8 +234,38 @@ VideoMetadata = TypedDict('VideoMetadata', 'frame_total' : int, 'fps' : Fps, 'resolution' : Resolution, - 'bit_rate' : BitRate + 'bit_rate' : BitRate, + 'color_transfer' : ColorTransfer }) +VideoReaderMetadata : TypeAlias = VideoMetadata +VideoWriterMetadata = TypedDict('VideoWriterMetadata', +{ + 'fps' : Fps, + 'resolution' : Resolution +}) +VideoReader = TypedDict('VideoReader', +{ + 'id' : str, + 'file_path' : str, + 'process' : subprocess.Popen[bytes], + 'metadata' : VideoReaderMetadata, + 'frame_number' : int +}) +VideoReaderSet : TypeAlias = Dict[str, VideoReader] +VideoWriter = TypedDict('VideoWriter', +{ + 'id' : str, + 'file_path' : str, + 'process' : subprocess.Popen[bytes], + 'metadata' : VideoWriterMetadata +}) +VideoWriterSet : TypeAlias = Dict[str, VideoWriter] +VideoPoolSet = TypedDict('VideoPoolSet', +{ + 'reader' : VideoReaderSet, + 'writer' : VideoWriterSet +}) +FrameStoreSet : TypeAlias = Dict[str, VisionFrameSet] AudioAsset = TypedDict('AudioAsset', { 'id' : AssetId, @@ -475,7 +504,6 @@ StateValue : TypeAlias = Any StateKey = Literal\ [ 'command', - 'workflow_mode', 'config_path', 'temp_path', 'jobs_path', @@ -518,6 +546,7 @@ StateKey = Literal\ 'trim_frame_start', 'trim_frame_end', 'temp_frame_format', + 'temp_pixel_format', 'keep_temp', 'target_frame_amount', 'output_image_quality', @@ -531,6 +560,8 @@ StateKey = Literal\ 'output_video_quality', 'output_video_scale', 'output_video_fps', + 'workflow_mode', + 'workflow_strategy', 'processors', 'execution_device_ids', 'execution_providers', @@ -549,7 +580,6 @@ StateKey = Literal\ State = TypedDict('State', { 'command' : str, - 'workflow_mode' : WorkflowMode, 'config_path' : str, 'temp_path' : str, 'jobs_path' : str, @@ -592,6 +622,7 @@ State = TypedDict('State', 'trim_frame_start' : int, 'trim_frame_end' : int, 'temp_frame_format' : TempFrameFormat, + 'temp_pixel_format' : TempPixelFormat, 'keep_temp' : bool, 'target_frame_amount' : int, 'output_image_quality' : int, @@ -605,6 +636,8 @@ State = TypedDict('State', 'output_video_scale' : Scale, 'output_video_fps' : float, 'output_audio_fps' : float, + 'workflow_mode' : WorkflowMode, + 'workflow_strategy' : WorkflowStrategy, 'processors' : List[str], 'execution_device_ids' : List[int], 'execution_providers' : List[ExecutionProvider], diff --git a/facefusion/video_manager.py b/facefusion/video_manager.py index b63203cc..798b3b2c 100644 --- a/facefusion/video_manager.py +++ b/facefusion/video_manager.py @@ -1,46 +1,155 @@ -import cv2 +import hashlib +import uuid +from io import BufferedReader +from typing import Optional, cast -from facefusion.types import VideoPoolSet +import numpy + +from facefusion import ffmpeg, ffprobe, frame_store, vision +from facefusion.common_helper import get_first, get_last +from facefusion.types import Fps, Resolution, VideoPoolSet, VideoReader, VideoWriter, VisionFrame, VisionFrameSet VIDEO_POOL_SET : VideoPoolSet =\ { - 'capture': {}, + 'reader': {}, 'writer': {} } -def get_video_capture(video_path : str) -> cv2.VideoCapture: - if video_path not in VIDEO_POOL_SET.get('capture'): - video_capture = cv2.VideoCapture(video_path) +def get_reader(video_path : str, context : str) -> VideoReader: + reader_id = hashlib.sha1((video_path + '_' + context).encode()).hexdigest() - if video_capture.isOpened(): - VIDEO_POOL_SET['capture'][video_path] = video_capture + if reader_id not in VIDEO_POOL_SET.get('reader'): + video_metadata = ffprobe.extract_static_video_metadata(video_path) - return VIDEO_POOL_SET.get('capture').get(video_path) + VIDEO_POOL_SET['reader'][reader_id] =\ + { + 'id': reader_id, + 'file_path': video_path, + 'process': ffmpeg.create_video_reader(video_path, 0, video_metadata), + 'metadata': video_metadata, + 'frame_number': 0 + } + + return VIDEO_POOL_SET.get('reader').get(reader_id) -def conditional_set_video_frame_position(video_capture : cv2.VideoCapture, frame_position : int) -> bool: - if not video_capture.get(cv2.CAP_PROP_POS_FRAMES) == frame_position: - return video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_position) - return True +def conditional_seek_video_reader(video_reader : VideoReader, frame_number : int = 0) -> None: + frame_total = video_reader.get('metadata').get('frame_total') + frame_number = min(frame_total - 1, frame_number) + skip_total = frame_number - video_reader.get('frame_number') + skip_margin = 128 + + if 0 < skip_total <= skip_margin: + drain_video_reader(video_reader, skip_total) + + if not video_reader.get('frame_number') == frame_number: + seek_video_reader(video_reader, frame_number) -def get_video_writer(video_path : str) -> cv2.VideoWriter: +def seek_video_reader(video_reader : VideoReader, frame_number : int = 0) -> None: + close_video_reader(video_reader) + + video_reader['process'] = ffmpeg.create_video_reader(video_reader.get('file_path'), frame_number, video_reader.get('metadata')) + video_reader['frame_number'] = frame_number + + +def drain_video_reader(video_reader : VideoReader, skip_total : int) -> None: + width, height = video_reader.get('metadata').get('resolution') + channel_total = 3 + frame_size = width * height * channel_total + + for _ in range(skip_total): + video_reader.get('process').stdout.read(frame_size) + + video_reader['frame_number'] = video_reader.get('frame_number') + skip_total + + +def read_video_frame(video_reader : VideoReader) -> Optional[VisionFrame]: + width, height = video_reader.get('metadata').get('resolution') + channel_total = 3 + video_stream = cast(BufferedReader, video_reader.get('process').stdout) + vision_frame = numpy.empty(width * height * channel_total, numpy.uint8) + + if video_stream.readinto(vision_frame) == vision_frame.size: + video_reader['frame_number'] = video_reader.get('frame_number') + 1 + return vision_frame.reshape(height, width, channel_total) + + return None + + +def read_video_frames(video_reader : VideoReader, frame_start : int, frame_end : int) -> VisionFrameSet: + reader_id = video_reader.get('id') + frame_set = frame_store.get_frame_store(reader_id) + keep_margin = 4 + frame_gaps = [] + + for frame_number in range(frame_start, frame_end + 1): + if frame_number not in frame_set: + frame_gaps.append(frame_number) + + if frame_gaps: + collect_video_frames(video_reader, get_first(frame_gaps), get_last(frame_gaps)) + + frame_store.reduce_frames(reader_id, frame_start - keep_margin, frame_end + keep_margin) + return frame_store.select_frame_set(reader_id, frame_start, frame_end) + + +def collect_video_frames(video_reader : VideoReader, frame_start : int, frame_end : int) -> None: + reader_id = video_reader.get('id') + skip_total = frame_start - video_reader.get('frame_number') + skip_margin = 16 + + if skip_total < 0 or skip_total > skip_margin: + seek_video_reader(video_reader, frame_start) + + for frame_number in range(video_reader.get('frame_number'), frame_end + 1): + vision_frame = read_video_frame(video_reader) + + if vision.is_vision_frame(vision_frame): + frame_store.set_frame(reader_id, frame_number, vision_frame) + + +def close_video_reader(video_reader : VideoReader) -> None: + video_reader.get('process').kill() + video_reader.get('process').wait() + + +def get_writer(video_path : str, temp_video_fps : Fps, temp_video_resolution : Resolution, output_video_resolution : Resolution, output_video_fps : Fps) -> VideoWriter: if video_path not in VIDEO_POOL_SET.get('writer'): - video_writer = cv2.VideoWriter() - - if video_writer.isOpened(): - VIDEO_POOL_SET['writer'][video_path] = video_writer + VIDEO_POOL_SET['writer'][video_path] =\ + { + 'id': uuid.uuid4().hex, + 'file_path': video_path, + 'process': ffmpeg.create_video_writer(video_path, temp_video_fps, temp_video_resolution, output_video_resolution, output_video_fps), + 'metadata': + { + 'fps': output_video_fps, + 'resolution': output_video_resolution + } + } return VIDEO_POOL_SET.get('writer').get(video_path) +def write_video_frame(video_writer : VideoWriter, vision_frame : VisionFrame) -> None: + video_writer.get('process').stdin.write(vision_frame.data) + + +def close_video_writer(video_writer : VideoWriter) -> bool: + video_writer.get('process').stdin.close() + video_writer.get('process').wait() + + return video_writer.get('process').returncode == 0 + + def clear_video_pool() -> None: - for video_capture in VIDEO_POOL_SET.get('capture').values(): - video_capture.release() + for video_reader in VIDEO_POOL_SET.get('reader').values(): + close_video_reader(video_reader) + frame_store.clear_frames(video_reader.get('id')) for video_writer in VIDEO_POOL_SET.get('writer').values(): - video_writer.release() + close_video_writer(video_writer) - VIDEO_POOL_SET['capture'].clear() + VIDEO_POOL_SET['reader'].clear() VIDEO_POOL_SET['writer'].clear() diff --git a/facefusion/vision.py b/facefusion/vision.py index b696fb82..24cce8a6 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -1,17 +1,17 @@ import math from functools import lru_cache -from typing import Dict, List, Optional, Tuple +from typing import List, Optional, Tuple import cv2 import numpy from cv2.typing import Size +from facefusion import ffprobe, video_manager from facefusion.common_helper import is_windows from facefusion.filesystem import get_file_extension, is_image, is_video from facefusion.media_helper import restrict_trim_frame from facefusion.thread_helper import thread_lock, thread_semaphore from facefusion.types import ColorMode, Duration, Fps, Mask, Orientation, Resolution, Scale, VisionFrame -from facefusion.video_manager import conditional_set_video_frame_position, get_video_capture def read_static_images(image_paths : List[str], color_mode : ColorMode = 'rgb') -> List[VisionFrame]: @@ -78,61 +78,30 @@ def read_static_video_frame(video_path : str, frame_number : int = 0) -> Optiona def read_video_frame(video_path : str, frame_number : int = 0) -> Optional[VisionFrame]: if is_video(video_path): - video_capture = get_video_capture(video_path) + video_reader = video_manager.get_reader(video_path, 'read_video_frame') - if video_capture and video_capture.isOpened(): - video_frame_total = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) - video_frame_position = min(video_frame_total, frame_number) - - with thread_semaphore(): - conditional_set_video_frame_position(video_capture, video_frame_position) - has_vision_frame, vision_frame = video_capture.read() - - if has_vision_frame: - return vision_frame + with thread_semaphore(): + video_manager.conditional_seek_video_reader(video_reader, frame_number) + return video_manager.read_video_frame(video_reader) return None -@lru_cache(maxsize = 2) -def read_static_video_chunk(video_path : str, chunk_number : int, chunk_size : int) -> Dict[int, VisionFrame]: - return read_video_chunk(video_path, chunk_number, chunk_size) - - -def read_video_chunk(video_path : str, chunk_number : int, chunk_size : int) -> Dict[int, VisionFrame]: - video_frame_chunk = {} - - if is_video(video_path) and chunk_number > -1: - video_capture = get_video_capture(video_path) - - if video_capture and video_capture.isOpened(): - video_frame_total = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) - video_frame_position = chunk_number * chunk_size - - with thread_semaphore(): - conditional_set_video_frame_position(video_capture, video_frame_position) - - for frame_number in range(video_frame_position, min(video_frame_position + chunk_size, video_frame_total)): - has_vision_frame, vision_frame = video_capture.read() - - if has_vision_frame: - video_frame_chunk[frame_number] = vision_frame - - return video_frame_chunk - - def select_video_frames(video_path : str, frame_number : int = 0, frame_offset : int = 2) -> List[VisionFrame]: vision_frames = [] - chunk_size = frame_offset * 2 + 1 + frame_start = frame_number - frame_offset + frame_end = frame_number + frame_offset if is_video(video_path): with thread_lock(): - for current_number in range(frame_number - frame_offset, frame_number + frame_offset + 1): - video_frame_chunk = read_static_video_chunk(video_path, current_number // chunk_size, chunk_size) + video_reader = video_manager.get_reader(video_path, 'select_video_frames') + frame_set = video_manager.read_video_frames(video_reader, max(frame_start, 0), frame_end) + + for frame_number in range(frame_start, frame_end + 1): vision_frame = create_empty_vision_frame() - if current_number in video_frame_chunk: - vision_frame = video_frame_chunk.get(current_number) + if frame_number in frame_set: + vision_frame = frame_set.get(frame_number) vision_frames.append(vision_frame) @@ -141,12 +110,7 @@ def select_video_frames(video_path : str, frame_number : int = 0, frame_offset : def count_video_frame_total(video_path : str) -> int: if is_video(video_path): - video_capture = get_video_capture(video_path) - - if video_capture and video_capture.isOpened(): - with thread_semaphore(): - video_frame_total = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) - return video_frame_total + return ffprobe.extract_static_video_metadata(video_path).get('frame_total') return 0 @@ -162,12 +126,7 @@ def predict_video_frame_total(video_path : str, fps : Fps, trim_frame_start : in def detect_video_fps(video_path : str) -> Optional[Fps]: if is_video(video_path): - video_capture = get_video_capture(video_path) - - if video_capture and video_capture.isOpened(): - with thread_semaphore(): - video_fps = video_capture.get(cv2.CAP_PROP_FPS) - return video_fps + return ffprobe.extract_static_video_metadata(video_path).get('fps') return None @@ -196,13 +155,7 @@ def restrict_trim_video_frame(video_path : str, trim_frame_start : Optional[int] def detect_video_resolution(video_path : str) -> Optional[Resolution]: if is_video(video_path): - video_capture = get_video_capture(video_path) - - if video_capture and video_capture.isOpened(): - with thread_semaphore(): - width = video_capture.get(cv2.CAP_PROP_FRAME_WIDTH) - height = video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT) - return int(width), int(height) + return ffprobe.extract_static_video_metadata(video_path).get('resolution') return None @@ -340,6 +293,10 @@ def create_empty_vision_frame() -> VisionFrame: return numpy.zeros((1, 1, 3)).astype(numpy.uint8) +def is_vision_frame(vision_frame : VisionFrame) -> bool: + return numpy.ndim(vision_frame) == 3 + + def create_tile_frames(vision_frame : VisionFrame, size : Size) -> Tuple[List[VisionFrame], int, int]: tile_width = size[0] - 2 * size[2] pad_size_top = size[1] + size[2] diff --git a/facefusion/workflows/core.py b/facefusion/workflows/core.py index 3172ec40..1e1e0db2 100644 --- a/facefusion/workflows/core.py +++ b/facefusion/workflows/core.py @@ -1,5 +1,6 @@ -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import List +from collections import deque +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Deque, List import numpy from tqdm import tqdm @@ -9,9 +10,9 @@ from facefusion.audio import create_empty_audio_frame, get_audio_frame, get_voic from facefusion.common_helper import get_first from facefusion.filesystem import filter_audio_paths from facefusion.processors.core import get_processors_modules -from facefusion.temp_helper import clear_temp_directory, create_temp_directory, resolve_temp_frame_paths +from facefusion.temp_helper import clear_temp_directory, create_temp_directory, resolve_temp_frame_set from facefusion.types import AudioFrame, ErrorCode, VisionFrame -from facefusion.vision import conditional_merge_vision_mask, extract_vision_mask, read_static_image, read_static_images, read_static_video_frame, restrict_video_fps, select_video_frames, write_image +from facefusion.vision import conditional_merge_vision_mask, extract_vision_mask, read_static_image, read_static_images, read_static_video_frame, restrict_trim_video_frame, restrict_video_fps, select_video_frames, write_image def is_process_stopping() -> bool: @@ -39,7 +40,9 @@ def conditional_get_source_audio_frame(frame_number : int) -> AudioFrame: fps = state_manager.get_item('output_audio_fps') if state_manager.get_item('workflow_mode') == 'image-to-video': + trim_frame_start, _ = restrict_trim_video_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + frame_number = frame_number - trim_frame_start source_audio_frame = get_audio_frame(source_audio_path, fps, frame_number) @@ -55,7 +58,9 @@ def conditional_get_source_voice_frame(frame_number : int) -> AudioFrame: temp_fps = state_manager.get_item('output_audio_fps') if state_manager.get_item('workflow_mode') == 'image-to-video': + trim_frame_start, _ = restrict_trim_video_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) temp_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + frame_number = frame_number - trim_frame_start source_voice_frame = get_voice_frame(source_audio_path, temp_fps, frame_number) @@ -102,26 +107,53 @@ def process_temp_frame(temp_frame_path : str, frame_number : int) -> bool: return write_image(temp_frame_path, temp_vision_frame) -def process_frames() -> ErrorCode: - temp_frame_paths = resolve_temp_frame_paths(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) +def process_temp_vision_frame(target_vision_frames : List[VisionFrame], temp_vision_frame : VisionFrame, frame_number : int) -> VisionFrame: + reference_vision_frame = conditional_get_reference_vision_frame() + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) + source_audio_frame = conditional_get_source_audio_frame(frame_number) + source_voice_frame = conditional_get_source_voice_frame(frame_number) + temp_vision_mask = extract_vision_mask(temp_vision_frame) - if temp_frame_paths: - with tqdm(total = len(temp_frame_paths), desc = translator.get('processing'), unit = 'frame', ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: + for processor_module in get_processors_modules(state_manager.get_item('processors')): + temp_vision_frame, temp_vision_mask = processor_module.process_frame( + { + 'reference_vision_frame': reference_vision_frame, + 'source_vision_frames': source_vision_frames, + 'source_audio_frame': source_audio_frame, + 'source_voice_frame': source_voice_frame, + 'target_vision_frames': target_vision_frames, + 'temp_vision_frame': temp_vision_frame[:, :, :3], + 'temp_vision_mask': temp_vision_mask + }) + + return conditional_merge_vision_mask(temp_vision_frame, temp_vision_mask) + + +def process_frames() -> ErrorCode: + temp_frame_set = resolve_temp_frame_set(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) + + if temp_frame_set: + with tqdm(total = len(temp_frame_set), 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 = [] + futures : Deque[Future[bool]] = deque() - for frame_number, temp_frame_path in enumerate(temp_frame_paths): + for frame_number, temp_frame_path in temp_frame_set.items(): future = executor.submit(process_temp_frame, temp_frame_path, frame_number) futures.append(future) - for future in as_completed(futures): - if is_process_stopping(): - for __future__ in futures: - __future__.cancel() + while futures: + future = futures.popleft() - if not future.cancelled(): + if is_process_stopping(): + + for pending_future in futures: + pending_future.cancel() + + futures.clear() + + else: future.result() progress.update() diff --git a/facefusion/workflows/image_to_video.py b/facefusion/workflows/image_to_video.py index 359be3a8..011f8bdc 100644 --- a/facefusion/workflows/image_to_video.py +++ b/facefusion/workflows/image_to_video.py @@ -1,9 +1,9 @@ from functools import partial -from facefusion import process_manager +from facefusion import process_manager, state_manager from facefusion.types import ErrorCode from facefusion.workflows.core import clear, process_frames, setup -from facefusion.workflows.to_video import analyse_video, create_temp_frames, finalize_video, merge_frames, restore_audio +from facefusion.workflows.to_video import analyse_video, create_temp_frames, finalize_video, merge_frames, process_memory_frames, restore_audio def process(start_time : float) -> ErrorCode: @@ -11,14 +11,26 @@ def process(start_time : float) -> ErrorCode: [ analyse_video, clear, - setup, - create_temp_frames, - process_frames, - merge_frames, + setup + ] + + if state_manager.get_item('workflow_strategy') == 'disk': + tasks.extend( + [ + create_temp_frames, + process_frames, + merge_frames + ]) + + if state_manager.get_item('workflow_strategy') == 'memory': + tasks.append(process_memory_frames) + + tasks.extend( + [ restore_audio, partial(finalize_video, start_time), clear - ] + ]) process_manager.start() diff --git a/facefusion/workflows/to_image.py b/facefusion/workflows/to_image.py index 119cac26..8cfa7ef6 100644 --- a/facefusion/workflows/to_image.py +++ b/facefusion/workflows/to_image.py @@ -1,5 +1,6 @@ from facefusion import content_analyser, ffmpeg, logger, process_manager, state_manager, translator from facefusion.filesystem import is_image +from facefusion.processors.core import get_processors_modules from facefusion.temp_helper import get_temp_file_path from facefusion.time_helper import calculate_end_time from facefusion.types import ErrorCode @@ -31,6 +32,9 @@ def process_image() -> ErrorCode: temp_image_path = get_temp_file_path(state_manager.get_temp_path(), state_manager.get_item('output_path')) process_temp_frame(temp_image_path, 0) + for processor_module in get_processors_modules(state_manager.get_item('processors')): + processor_module.post_process() + if is_process_stopping(): return 4 return 0 diff --git a/facefusion/workflows/to_video.py b/facefusion/workflows/to_video.py index 7e2212bc..81ccd93d 100644 --- a/facefusion/workflows/to_video.py +++ b/facefusion/workflows/to_video.py @@ -1,12 +1,21 @@ -from facefusion import content_analyser, ffmpeg, logger, state_manager, translator, video_manager -from facefusion.common_helper import get_first +from collections import deque +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Deque + +import cv2 +import numpy +from tqdm import tqdm + +from facefusion import content_analyser, ffmpeg, logger, process_manager, state_manager, translator, video_manager +from facefusion.common_helper import get_first, get_middle from facefusion.filesystem import filter_audio_paths, is_video from facefusion.media_helper import restrict_trim_frame +from facefusion.processors.core import get_processors_modules from facefusion.temp_helper import move_temp_file, resolve_temp_frame_paths from facefusion.time_helper import calculate_end_time -from facefusion.types import ErrorCode, Fps, Resolution -from facefusion.vision import detect_image_resolution, detect_video_resolution, pack_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution -from facefusion.workflows.core import is_process_stopping +from facefusion.types import ErrorCode, Fps, Resolution, VisionFrame +from facefusion.vision import detect_image_resolution, detect_video_resolution, pack_resolution, read_static_video_frame, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames +from facefusion.workflows.core import is_process_stopping, process_temp_vision_frame def analyse_video() -> ErrorCode: @@ -34,6 +43,75 @@ def create_temp_frames() -> ErrorCode: return 0 +def process_memory_frame(frame_number : int, temp_video_resolution : Resolution) -> VisionFrame: + 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_vision_frame = target_vision_frame.copy() + + if not (target_vision_frame.shape[1], target_vision_frame.shape[0]) == temp_video_resolution: + temp_vision_frame = cv2.resize(target_vision_frame, temp_video_resolution) + + temp_vision_frame = process_temp_vision_frame(target_vision_frames, temp_vision_frame, frame_number) + + if state_manager.get_item('temp_pixel_format') == 'bgra': + temp_vision_frame = cv2.cvtColor(temp_vision_frame, cv2.COLOR_BGR2BGRA) + + if state_manager.get_item('temp_pixel_format') == 'bgr24': + temp_vision_frame = temp_vision_frame[:, :, :3] + + return numpy.ascontiguousarray(temp_vision_frame) + + +def process_memory_frames() -> ErrorCode: + trim_frame_start, trim_frame_end = restrict_trim_video_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')) + temp_frame_range = range(trim_frame_start, trim_frame_end) + + if temp_frame_range: + video_writer = video_manager.get_writer(state_manager.get_item('output_path'), temp_video_fps, temp_video_resolution, output_video_resolution, state_manager.get_item('output_video_fps')) + + with tqdm(total = len(temp_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')) + + read_static_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) + + with ThreadPoolExecutor(max_workers = state_manager.get_item('execution_thread_count')) as executor: + futures : Deque[Future[VisionFrame]] = deque() + + for frame_number in temp_frame_range: + future = executor.submit(process_memory_frame, frame_number, temp_video_resolution) + futures.append(future) + + while futures: + future = futures.popleft() + + if is_process_stopping(): + + for pending_future in futures: + pending_future.cancel() + + futures.clear() + + else: + video_manager.write_video_frame(video_writer, future.result()) + progress.update() + + if not video_manager.close_video_writer(video_writer): + process_manager.stop() + + for processor_module in get_processors_modules(state_manager.get_item('processors')): + processor_module.post_process() + + if is_process_stopping(): + return 4 + else: + logger.error(translator.get('temp_frames_not_found'), __name__) + return 1 + return 0 + + def merge_frames() -> ErrorCode: temp_frame_paths = resolve_temp_frame_paths(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) trim_frame_start, trim_frame_end = restrict_trim_frame(len(temp_frame_paths), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) diff --git a/requirements.txt b/requirements.txt index 7e67d4df..e72c0933 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,11 @@ -numpy==2.4.2 -onnx==1.20.1 -onnxruntime==1.26.0 -opencv-python==4.13.0.92 +numpy==2.4.6 +onnx==1.22.0 +onnxruntime==1.28.0 +opencv-python-headless==5.0.0.93 psutil==7.2.2 python-multipart==0.0.28 -tqdm==4.67.3 -scipy==1.16.3 +tqdm==4.70.0 +scipy==1.18.0 starlette==0.52.1 uvicorn==0.41.0 websockets==16.0 diff --git a/tests/test_api_stream_video.py b/tests/test_api_stream_video.py index a1e22244..50b49110 100644 --- a/tests/test_api_stream_video.py +++ b/tests/test_api_stream_video.py @@ -79,7 +79,7 @@ def test_run_video_encode_loop(video_codec : VideoCodec, payload_type : int) -> video_queue : Queue[Tuple[Time, Future[BufferPack]]] = Queue(maxsize = 30) with ThreadPoolExecutor(max_workers = 1) as executor: - video_queue.put((0.1, executor.submit(process_video_frame, video_frame))) + video_queue.put((0.1, executor.submit(process_video_frame, [], video_frame))) with patch('facefusion.apis.stream_video.rtc.send_video') as send_video_mock: encode_loop_thread = threading.Thread(target = run_video_encode_loop, args = (rtc_peer, video_queue), daemon = True) @@ -267,7 +267,7 @@ def test_handle_video_frame(video_codec : VideoCodec) -> None: with ThreadPoolExecutor(max_workers = 1) as executor: with patch('facefusion.apis.stream_video.decode_video_frame', return_value = video_frame): with patch('facefusion.apis.stream_video.process_video_frame', return_value = BufferPack(buffer = video_frame.tobytes(), resolution = (426, 226))): - handle_video_frame(video_codec, video_decoder, video_queue, executor, bytes(), 0) + handle_video_frame([], video_codec, video_decoder, video_queue, executor, bytes(), 0) _, video_future = video_queue.get_nowait() video_buffer = video_future.result().get('buffer') diff --git a/tests/test_face_tracker.py b/tests/test_face_tracker.py index 0c655e6d..134d2874 100644 --- a/tests/test_face_tracker.py +++ b/tests/test_face_tracker.py @@ -7,7 +7,7 @@ from facefusion.download import conditional_download from facefusion.face_creator import get_many_faces, get_one_face from facefusion.face_store import clear_faces from facefusion.face_tracker import create_face_tracks, select_face_track, track_faces -from facefusion.vision import read_static_video_chunk, read_static_video_frame +from facefusion.vision import read_static_video_frame, select_video_frames from .assert_helper import get_test_example_file, get_test_examples_directory @@ -47,8 +47,7 @@ def before_each() -> None: def test_track_faces() -> None: target_path = get_test_example_file('target-240p.mp4') - video_frame_chunk = read_static_video_chunk(target_path, 0, 7) - target_vision_frames = [ video_frame_chunk.get(frame_number) for frame_number in sorted(video_frame_chunk) ] + target_vision_frames = select_video_frames(target_path, 3, 3) empty_vision_frame = numpy.zeros_like(get_first(target_vision_frames)) target_vision_frames[2] = empty_vision_frame @@ -58,7 +57,7 @@ def test_track_faces() -> None: assert len(track_faces(target_vision_frames, 0.3)) == 1 - target_vision_frames = [ video_frame_chunk.get(frame_number) for frame_number in sorted(video_frame_chunk)[:5] ] + target_vision_frames = select_video_frames(target_path, 3, 3)[:5] target_vision_frames[0] = empty_vision_frame target_vision_frames[1] = empty_vision_frame target_vision_frames[2] = empty_vision_frame diff --git a/tests/test_ffmpeg.py b/tests/test_ffmpeg.py index 54bea288..c114ae45 100644 --- a/tests/test_ffmpeg.py +++ b/tests/test_ffmpeg.py @@ -7,7 +7,7 @@ import facefusion.ffmpeg from facefusion import ffmpeg, ffmpeg_builder, process_manager, state_manager from facefusion.download import conditional_download from facefusion.ffmpeg import concat_video, extract_frames, merge_video, read_audio_buffer, replace_audio, restore_audio, sanitize_audio, sanitize_image, sanitize_video, spawn_frames -from facefusion.ffprobe import probe_entries +from facefusion.ffprobe import probe_audio_entries, probe_video_entries from facefusion.filesystem import copy_file, is_image from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_file_path, resolve_temp_frame_paths from facefusion.types import EncoderSet @@ -253,10 +253,10 @@ def test_sanitize_audio() -> None: ] assert sanitize_audio(file_content, output_paths[0], 'strict') is True - assert probe_entries(output_paths[0], [ 'codec_name' ]).get('codec_name') == 'mp3' + assert probe_audio_entries(output_paths[0], [ 'codec_name' ]).get('codec_name') == 'mp3' assert sanitize_audio(file_content, output_paths[1], 'moderate') is True - assert probe_entries(output_paths[1], [ 'codec_name' ]).get('codec_name') == 'pcm_s16le' + assert probe_audio_entries(output_paths[1], [ 'codec_name' ]).get('codec_name') == 'pcm_s16le' def test_sanitize_image() -> None: @@ -278,9 +278,9 @@ def test_sanitize_video() -> None: ] assert sanitize_video(file_content, output_paths[0], 'strict') is True - assert probe_entries(output_paths[0], [ 'codec_name' ]).get('codec_name') == 'h264' + assert probe_video_entries(output_paths[0], [ 'codec_name' ]).get('codec_name') == 'h264' assert sanitize_video(file_content, output_paths[1], 'moderate') is True - assert probe_entries(output_paths[1], [ 'codec_name' ]).get('codec_name') == 'hevc' + assert probe_video_entries(output_paths[1], [ 'codec_name' ]).get('codec_name') == 'hevc' diff --git a/tests/test_ffprobe.py b/tests/test_ffprobe.py index 8240c3f8..7f3d1517 100644 --- a/tests/test_ffprobe.py +++ b/tests/test_ffprobe.py @@ -26,13 +26,14 @@ def before_all() -> None: ffmpeg_builder.set_output(get_test_example_file('source-48000khz-2ch.wav')) ) ) - ffmpeg.run_ffmpeg( - ffmpeg_builder.chain( - ffmpeg_builder.set_input(get_test_example_file('target-240p.mp4')), - ffmpeg_builder.set_video_duration(1), - ffmpeg_builder.set_output(get_test_example_file('target-240p-1s.mov')) + for video_format in [ 'mkv', 'mov' ]: + ffmpeg.run_ffmpeg( + ffmpeg_builder.chain( + ffmpeg_builder.set_input(get_test_example_file('target-240p.mp4')), + ffmpeg_builder.set_video_duration(1), + ffmpeg_builder.set_output(get_test_example_file('target-240p-1s.' + video_format)) + ) ) - ) def test_extract_audio_metadata() -> None: @@ -40,7 +41,7 @@ def test_extract_audio_metadata() -> None: assert audio_metadata.get('sample_rate') == 44100 assert audio_metadata.get('channel_total') == 1 - assert audio_metadata.get('frame_total') == 167039 + assert audio_metadata.get('frame_total') == 167040 assert audio_metadata.get('bit_rate') == 128000 audio_metadata = extract_audio_metadata(get_test_example_file('source-48000khz-2ch.wav')) @@ -48,7 +49,7 @@ def test_extract_audio_metadata() -> None: assert audio_metadata.get('sample_rate') == 48000 assert audio_metadata.get('channel_total') == 2 assert audio_metadata.get('frame_total') == 91200 - assert audio_metadata.get('bit_rate') == 1536000 + assert audio_metadata.get('bit_rate') == 1536328 def test_extract_video_metadata() -> None: @@ -57,7 +58,15 @@ def test_extract_video_metadata() -> None: assert video_metadata.get('fps') == 25.0 assert video_metadata.get('duration') == 10.8 assert video_metadata.get('resolution') == (426, 226) - assert video_metadata.get('bit_rate') == 138754 + assert video_metadata.get('bit_rate') == 141981 + assert video_metadata.get('color_transfer') == 'smpte170m' + + video_metadata = extract_video_metadata(get_test_example_file('target-240p-1s.mkv')) + + assert video_metadata.get('fps') == 25.0 + assert video_metadata.get('duration') == 1.0 + assert video_metadata.get('frame_total') == 25 + assert video_metadata.get('resolution') == (426, 226) video_metadata = extract_video_metadata(get_test_example_file('target-240p-1s.mov')) diff --git a/tests/test_ffprobe_builder.py b/tests/test_ffprobe_builder.py index 88aaf641..050b9097 100644 --- a/tests/test_ffprobe_builder.py +++ b/tests/test_ffprobe_builder.py @@ -1,7 +1,7 @@ from shutil import which from facefusion import ffprobe_builder -from facefusion.ffprobe_builder import chain, format_to_key_value, run, set_input, show_entries +from facefusion.ffprobe_builder import chain, format_to_key_value, run, select_stream, set_input, show_stream_entries def test_run() -> None: @@ -10,15 +10,21 @@ def test_run() -> None: def test_chain() -> None: assert chain( - ffprobe_builder.show_entries([ 'sample_rate' ]), + ffprobe_builder.select_stream('a:0'), + ffprobe_builder.show_stream_entries([ 'sample_rate' ]), ffprobe_builder.format_to_key_value(), ffprobe_builder.set_input('audio.mp3') - ) == [ '-show_entries', 'stream=sample_rate', '-of', 'default=noprint_wrappers=1', '-i', 'audio.mp3' ] + ) == [ '-select_streams', 'a:0', '-show_entries', 'stream=sample_rate', '-of', 'default=noprint_wrappers=1', '-i', 'audio.mp3' ] -def test_show_entries() -> None: - assert show_entries([ 'duration' ]) == [ '-show_entries', 'stream=duration' ] - assert show_entries([ 'duration', 'sample_rate']) == [ '-show_entries', 'stream=duration,sample_rate' ] +def test_select_stream() -> None: + assert select_stream('a:0') == [ '-select_streams', 'a:0' ] + assert select_stream('v:0') == [ '-select_streams', 'v:0' ] + + +def test_show_stream_entries() -> None: + assert show_stream_entries([ 'duration' ]) == [ '-show_entries', 'stream=duration' ] + assert show_stream_entries([ 'duration', 'sample_rate' ]) == [ '-show_entries', 'stream=duration,sample_rate' ] def test_format_to_key_value() -> None: diff --git a/tests/test_frame_store.py b/tests/test_frame_store.py new file mode 100644 index 00000000..9f8ae199 --- /dev/null +++ b/tests/test_frame_store.py @@ -0,0 +1,73 @@ +import pytest + +from facefusion import process_manager +from facefusion.download import conditional_download +from facefusion.frame_store import clear_frames, get_frame_store, reduce_frames, select_frame_set, set_frame +from facefusion.vision import read_video_frame +from .assert_helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + process_manager.start() + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_frames('reader-1') + clear_frames('reader-2') + + +def test_get_frame_store() -> None: + frame_store = get_frame_store('reader-1') + + assert frame_store == {} + assert get_frame_store('reader-1') is frame_store + + +def test_set_frame() -> None: + target_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) + + set_frame('reader-1', 5, target_frame) + + assert get_frame_store('reader-1').get(5) is target_frame + + +def test_select_frame_set() -> None: + first_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) + fifth_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 5) + + set_frame('reader-1', 2, first_frame) + set_frame('reader-1', 5, fifth_frame) + + assert sorted(select_frame_set('reader-1', 0, 4)) == [ 2 ] + assert select_frame_set('reader-1', 0, 4).get(2) is first_frame + assert sorted(select_frame_set('reader-1', 2, 6)) == [ 2, 5 ] + assert select_frame_set('reader-1', 2, 6).get(5) is fifth_frame + assert select_frame_set('reader-1', 8, 12) == {} + + +def test_reduce_frames() -> None: + target_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) + + for frame_number in range(0, 10): + set_frame('reader-1', frame_number, target_frame) + + reduce_frames('reader-1', 4, 6) + + assert sorted(get_frame_store('reader-1')) == [ 4, 5, 6 ] + + +def test_clear_frames() -> None: + target_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) + + set_frame('reader-1', 0, target_frame) + set_frame('reader-2', 0, target_frame) + clear_frames('reader-1') + + assert get_frame_store('reader-1') == {} + assert get_frame_store('reader-2').get(0) is target_frame diff --git a/tests/test_video_manager.py b/tests/test_video_manager.py new file mode 100644 index 00000000..2c28757d --- /dev/null +++ b/tests/test_video_manager.py @@ -0,0 +1,203 @@ +import tempfile + +import numpy +import pytest + +from facefusion import ffmpeg, ffmpeg_builder, process_manager, state_manager +from facefusion.common_helper import is_linux, is_macos, is_windows +from facefusion.download import conditional_download +from facefusion.ffprobe import extract_video_metadata +from facefusion.frame_store import get_frame_store +from facefusion.temp_helper import create_temp_directory, get_temp_file_path +from facefusion.video_manager import clear_video_pool, close_video_reader, close_video_writer, collect_video_frames, conditional_seek_video_reader, drain_video_reader, get_reader, get_writer, read_video_frame, read_video_frames, seek_video_reader, write_video_frame +from .assert_helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + process_manager.start() + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + + for video_fps in [ 25, 30 ]: + ffmpeg.run_ffmpeg( + ffmpeg_builder.chain( + ffmpeg_builder.set_input(get_test_example_file('target-240p.mp4')), + ffmpeg_builder.set_video_fps(video_fps), + ffmpeg_builder.set_output(get_test_example_file('target-240p-' + str(video_fps) + 'fps.mp4')) + ) + ) + + state_manager.init_item('temp_path', tempfile.gettempdir()) + state_manager.init_item('temp_frame_format', 'png') + state_manager.init_item('temp_pixel_format', 'bgr24') + state_manager.init_item('output_video_encoder', 'libx264') + state_manager.init_item('output_video_quality', 80) + state_manager.init_item('output_video_preset', 'veryfast') + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_video_pool() + + +def test_get_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + video_metadata = video_reader.get('metadata') + + assert video_metadata.get('resolution') == (426, 226) + assert video_metadata.get('fps') == 25.0 + assert video_metadata.get('frame_total') == 270 + + assert get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') is video_reader + assert not get_reader(get_test_example_file('target-240p-25fps.mp4'), 'select_video_frames').get('id') == video_reader.get('id') + + +def test_conditional_seek_video_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + video_frames = {} + + for frame_number in range(30): + video_frames[frame_number] = read_video_frame(video_reader) + + for frame_number in [ 5, 17, 29 ]: + conditional_seek_video_reader(video_reader, frame_number) + + assert numpy.array_equal(read_video_frame(video_reader), video_frames.get(frame_number)) is True + + +def test_seek_video_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + video_frames = {} + + for frame_number in range(30): + video_frames[frame_number] = read_video_frame(video_reader) + + for frame_number in [ 5, 17, 29 ]: + seek_video_reader(video_reader, frame_number) + + assert numpy.array_equal(read_video_frame(video_reader), video_frames.get(frame_number)) is True + + +def test_drain_video_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + + drain_video_reader(video_reader, 10) + + assert video_reader.get('frame_number') == 10 + + vision_frame = read_video_frame(video_reader) + seek_video_reader(video_reader, 10) + + assert numpy.array_equal(vision_frame, read_video_frame(video_reader)) is True + + +def test_read_video_frame() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + + assert read_video_frame(video_reader).shape == (226, 426, 3) + assert video_reader.get('frame_number') == 1 + + seek_video_reader(video_reader, 269) + + assert read_video_frame(video_reader).shape == (226, 426, 3) + assert read_video_frame(video_reader) is None + + +def test_read_video_frames() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + + assert sorted(read_video_frames(video_reader, 0, 4)) == [ 0, 1, 2, 3, 4 ] + + frame_number = video_reader.get('frame_number') + + assert sorted(read_video_frames(video_reader, 1, 3)) == [ 1, 2, 3 ] + assert video_reader.get('frame_number') == frame_number + + read_video_frames(video_reader, 21, 25) + + assert min(get_frame_store(video_reader.get('id'))) == 17 + assert max(get_frame_store(video_reader.get('id'))) == 25 + assert sorted(read_video_frames(video_reader, 268, 275)) == [ 268, 269 ] + + +def test_collect_video_frames() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'select_video_frames') + + collect_video_frames(video_reader, 20, 24) + + assert sorted(get_frame_store(video_reader.get('id'))) == [ 20, 21, 22, 23, 24 ] + assert video_reader.get('frame_number') == 25 + + +def test_close_video_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'select_video_frames') + + read_video_frames(video_reader, 0, 4) + close_video_reader(video_reader) + + if is_windows(): + assert video_reader.get('process').returncode == 1 + + if is_linux() or is_macos(): + assert video_reader.get('process').returncode == -9 + + +def test_get_writer() -> None: + target_path = get_test_example_file('target-240p-25fps.mp4') + create_temp_directory(state_manager.get_temp_path(), target_path) + video_writer = get_writer(target_path, 25.0, (426, 226), (426, 226), 25.0) + + assert get_writer(target_path, 25.0, (426, 226), (426, 226), 25.0) is video_writer + + +def test_write_video_frame() -> None: + target_path = get_test_example_file('target-240p-25fps.mp4') + create_temp_directory(state_manager.get_temp_path(), target_path) + video_reader = get_reader(target_path, 'read_video_frame') + video_writer = get_writer(target_path, 25.0, (426, 226), (426, 226), 25.0) + + for frame_number in range(25): + write_video_frame(video_writer, read_video_frame(video_reader)) + + assert close_video_writer(video_writer) is True + + video_metadata = extract_video_metadata(get_temp_file_path(state_manager.get_temp_path(), target_path)) + + assert video_metadata.get('duration') == 1.0 + assert video_metadata.get('frame_total') == 25 + assert video_metadata.get('fps') == 25.0 + assert video_metadata.get('resolution') == (426, 226) + assert video_metadata.get('color_transfer') == 'bt709' + + +def test_close_video_writer() -> None: + target_path = get_test_example_file('target-240p-30fps.mp4') + create_temp_directory(state_manager.get_temp_path(), target_path) + video_reader = get_reader(target_path, 'read_video_frame') + video_writer = get_writer(target_path, 30.0, (426, 226), (426, 226), 30.0) + + write_video_frame(video_writer, read_video_frame(video_reader)) + + assert close_video_writer(video_writer) is True + + +def test_clear_video_pool() -> None: + target_path = get_test_example_file('target-240p-25fps.mp4') + create_temp_directory(state_manager.get_temp_path(), target_path) + video_reader = get_reader(target_path, 'select_video_frames') + video_writer = get_writer(target_path, 25.0, (426, 226), (426, 226), 25.0) + + read_video_frames(video_reader, 0, 4) + write_video_frame(video_writer, read_video_frame(video_reader)) + clear_video_pool() + + if is_windows(): + assert video_reader.get('process').returncode == 1 + + if is_linux() or is_macos(): + assert video_reader.get('process').returncode == -9 + + assert video_writer.get('process').returncode == 0 diff --git a/tests/test_vision.py b/tests/test_vision.py index 04a8071b..a499c7b3 100644 --- a/tests/test_vision.py +++ b/tests/test_vision.py @@ -4,7 +4,7 @@ import pytest from facefusion import ffmpeg, ffmpeg_builder, process_manager from facefusion.download import conditional_download -from facefusion.vision import calculate_histogram_difference, count_video_frame_total, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution, match_frame_color, normalize_resolution, pack_resolution, predict_video_frame_total, read_image, read_video_chunk, read_video_frame, restrict_image_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames, unpack_resolution, write_image +from facefusion.vision import calculate_histogram_difference, count_video_frame_total, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution, match_frame_color, normalize_resolution, pack_resolution, predict_video_frame_total, read_image, read_video_frame, restrict_image_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames, unpack_resolution, write_image from .assert_helper import get_test_example_file, get_test_examples_directory, get_test_output_path, prepare_test_output_directory @@ -136,11 +136,6 @@ def test_read_video_frame() -> None: assert read_video_frame('invalid') is None -def test_read_video_chunk() -> None: - assert len(read_video_chunk(get_test_example_file('target-240p-25fps.mp4'), 1, 40)) == 40 - assert read_video_chunk('invalid', 1, 40) == {} - - def test_select_video_frames() -> None: assert len(select_video_frames(get_test_example_file('target-240p-25fps.mp4'), 50, 5)) == 11 assert len(select_video_frames(get_test_example_file('target-240p-25fps.mp4'), 1, 5)) == 11