diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 0000000000000000000000000000000000000000..4cd58731e3b8b11a77cfca951f8d3eea5771ef84 --- /dev/null +++ b/.coveragerc @@ -0,0 +1,2 @@ +[run] +patch = subprocess diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000000000000000000000000000000000000..fdb96ba2d200d5839f4acb490fc3413ea548966d --- /dev/null +++ b/.editorconfig @@ -0,0 +1,8 @@ +root = true + +[*] +end_of_line = lf +insert_final_newline = true +indent_size = 4 +indent_style = tab +trim_trailing_whitespace = true diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000000000000000000000000000000000000..7bd108f24ef44590eaeb5ce0ba4f40a746e50c3b --- /dev/null +++ b/.flake8 @@ -0,0 +1,6 @@ +[flake8] +select = E22, E23, E24, E27, E3, E4, E7, F, I1, I2 +per-file-ignores = facefusion.py:E402, install.py:E402 +plugins = flake8-import-order +application_import_names = facefusion +import-order-style = pycharm diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..2448af21d6983b6c4e360a4d6f7b1bd9c3a4cab7 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +.github/preview.png filter=lfs diff=lfs merge=lfs -text diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 0000000000000000000000000000000000000000..933813985ce29ad6e071385db6693d15eb27fadb --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1 @@ +custom: [ buymeacoffee.com/facefusion, ko-fi.com/facefusion ] diff --git a/.github/preview.png b/.github/preview.png new file mode 100644 index 0000000000000000000000000000000000000000..77877a060ec08c6b3b0f41bf194180f71a0e9a58 --- /dev/null +++ b/.github/preview.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1e4ce6832887293dd83fc472877564215bb6d0e148168f8cd6f38b42348945fe +size 1354191 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000000000000000000000000000000000000..d7e4af80ac37c128d246483d804a2c076d49d8fa --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,58 @@ +name: ci + +on: [ push, pull_request ] + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: pip install flake8 + - run: pip install flake8-import-order + - run: pip install mypy + - run: flake8 facefusion.py install.py + - run: flake8 facefusion tests + - run: mypy facefusion.py install.py + - run: mypy facefusion tests + test: + strategy: + matrix: + os: [ macos-latest, ubuntu-latest, windows-latest ] + runs-on: ${{ matrix.os }} + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Set up FFmpeg + uses: AnimMouse/setup-ffmpeg@v1 + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: python install.py --onnxruntime default --skip-conda + - run: pip install pytest + - run: pytest + report: + needs: test + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Set up FFmpeg + uses: FedericoCarboni/setup-ffmpeg@v3 + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: python install.py --onnxruntime default --skip-conda + - run: pip install coveralls + - run: pip install pytest + - run: pip install pytest-cov + - run: pytest tests --cov facefusion + - run: coveralls --service github + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..199ace197d71a90c6723c6bb51e7497c6f39de37 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +__pycache__ +.assets +.claude +.caches +.idea +.jobs +.vscode diff --git a/LICENSE.md b/LICENSE.md new file mode 100644 index 0000000000000000000000000000000000000000..2adc121491ddd49fabbcd6e0d1ba1ec74853509a --- /dev/null +++ b/LICENSE.md @@ -0,0 +1,3 @@ +OpenRAIL-AS license + +Copyright (c) 2025 Henry Ruhs diff --git a/README.md b/README.md index f787b9dd17b3c512e821db46e67ba10af6acdafa..ad0fc08bb0364c2f521217f34d7bc6710961abc7 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,26 @@ ---- -title: Facefusion -emoji: 📊 -colorFrom: yellow -colorTo: purple -sdk: gradio -sdk_version: 6.2.0 -app_file: app.py -pinned: false ---- - -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +--- +title: FaceFusion +emoji: 🌖 +colorFrom: indigo +colorTo: pink +sdk: gradio +sdk_version: 5.44.1 +app_file: app.py +pinned: false +license: mit +--- + +# FaceFusion + +> Industry leading face manipulation platform. + +## Deployment on Hugging Face Spaces +This Space runs FaceFusion using the **Gradio SDK**. + +**Hardware Recommendation:** +- **CPU Basic**: Slow, not recommended for video. +- **GPU (Zero/Pro)**: Recommended for reasonable speed. + +## Configuration +The `packages.txt` file installs `ffmpeg` which is required for media processing. +`app.py` is configured to bypass CLI arguments and start the UI directly. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..d4340148d1d434ef0ab9588a31774acc90f0a122 --- /dev/null +++ b/app.py @@ -0,0 +1,74 @@ +import os +import sys +import shutil +import gradio as gr +from typing import List +from types import ModuleType + +# Force current directory into python path +sys.path.append(os.path.abspath(".")) + +# Mock or Setup FaceFusion imports +# We need to manually initialize state_manager because we are skipping the CLI argument parsing +from facefusion import state_manager, metadata, config +import facefusion.uis.core as ui_core +import facefusion.uis.layouts.default as default_layout +from facefusion import logger + +# Set default state for FaceFusion +# These values mimic what 'python facefusion.py run' would produce with default args +state_manager.init_item('command', 'run') +state_manager.init_item('ui_layouts', ['default']) +state_manager.init_item('system_memory_limit', 0) +state_manager.init_item('face_selector_mode', 'reference') +state_manager.init_item('face_selector_order', 'large-small') +state_manager.init_item('face_selector_gender', None) +state_manager.init_item('face_selector_race', None) +state_manager.init_item('face_detector_model', 'yolo_face') +state_manager.init_item('face_detector_size', '640x640') +state_manager.init_item('face_detector_score', 0.5) +state_manager.init_item('face_landmarker_model', '2dfan4') +state_manager.init_item('face_landmarker_score', 0.5) +state_manager.init_item('face_mask_types', ['box']) +state_manager.init_item('face_mask_areas', ['skin', 'l-eyebrow', 'r-eyebrow', 'l-eye', 'r-eye', 'eye-g', 'l-ear', 'r-ear', 'ear-r', 'nose', 'mouth', 'u-lip', 'l-lip', 'neck', 'neck-l', 'cloth']) # Default all +state_manager.init_item('face_mask_regions', ['skin', 'l-eyebrow', 'r-eyebrow', 'l-eye', 'r-eye', 'eye-g', 'l-ear', 'r-ear', 'ear-r', 'nose', 'mouth', 'u-lip', 'l-lip', 'neck', 'neck-l', 'cloth']) +state_manager.init_item('face_mask_blur', 0.3) +state_manager.init_item('face_mask_padding', (0, 0, 0, 0)) +state_manager.init_item('face_occluder_model', 'xseg_1') +state_manager.init_item('face_parser_model', 'bisenet_resnet_34') +state_manager.init_item('processors', ['face_swapper']) +state_manager.init_item('execution_providers', ['cpu']) # Change to 'cuda' if GPU space +state_manager.init_item('execution_thread_count', 4) +state_manager.init_item('execution_device_ids', [0]) +state_manager.init_item('temp_path', os.path.join(os.getcwd(), '.tmp')) +state_manager.init_item('output_image_quality', 80) +state_manager.init_item('output_video_quality', 80) +state_manager.init_item('output_video_fps', 30) +state_manager.init_item('output_video_encoder', 'libx264') +state_manager.init_item('output_audio_encoder', 'aac') +state_manager.init_item('log_level', 'info') + +# Helper to load layout +def get_theme() -> gr.Theme: + return ui_core.get_theme() + +def get_css() -> str: + return ui_core.get_css() + +# Initialize UI +ui_core.init() + +# Launch logic adapted from facefusion.uis.core.launch +# We construct the Block here and expose it as 'demo' for Hugging Face +with gr.Blocks(theme=get_theme(), css=get_css(), title=metadata.get('name') + ' ' + metadata.get('version')) as demo: + # Render layout + # Since we only use 'default', we import it directly or load it + default_layout.render() + default_layout.listen() + + # Run setup + default_layout.run(demo) + +# Launch +if __name__ == "__main__": + demo.launch() diff --git a/facefusion.ico b/facefusion.ico new file mode 100644 index 0000000000000000000000000000000000000000..7dd4da75931ebe108efbb84b28d1585d04495533 Binary files /dev/null and b/facefusion.ico differ diff --git a/facefusion.ini b/facefusion.ini new file mode 100644 index 0000000000000000000000000000000000000000..ef2c5209738892657fec151eb0f4fc9686b69b11 --- /dev/null +++ b/facefusion.ini @@ -0,0 +1,131 @@ +[paths] +temp_path = +jobs_path = +source_paths = +target_path = +output_path = + +[patterns] +source_pattern = +target_pattern = +output_pattern = + +[face_detector] +face_detector_model = +face_detector_size = +face_detector_margin = +face_detector_angles = +face_detector_score = + +[face_landmarker] +face_landmarker_model = +face_landmarker_score = + +[face_selector] +face_selector_mode = +face_selector_order = +face_selector_age_start = +face_selector_age_end = +face_selector_gender = +face_selector_race = +reference_face_position = +reference_face_distance = +reference_frame_number = + +[face_masker] +face_occluder_model = +face_parser_model = +face_mask_types = +face_mask_areas = +face_mask_regions = +face_mask_blur = +face_mask_padding = + +[voice_extractor] +voice_extractor_model = + +[frame_extraction] +trim_frame_start = +trim_frame_end = +temp_frame_format = +keep_temp = + +[output_creation] +output_image_quality = +output_image_scale = +output_audio_encoder = +output_audio_quality = +output_audio_volume = +output_video_encoder = +output_video_preset = +output_video_quality = +output_video_scale = +output_video_fps = + +[processors] +processors = +age_modifier_model = +age_modifier_direction = +background_remover_model = +background_remover_color = +deep_swapper_model = +deep_swapper_morph = +expression_restorer_model = +expression_restorer_factor = +expression_restorer_areas = +face_debugger_items = +face_editor_model = +face_editor_eyebrow_direction = +face_editor_eye_gaze_horizontal = +face_editor_eye_gaze_vertical = +face_editor_eye_open_ratio = +face_editor_lip_open_ratio = +face_editor_mouth_grim = +face_editor_mouth_pout = +face_editor_mouth_purse = +face_editor_mouth_smile = +face_editor_mouth_position_horizontal = +face_editor_mouth_position_vertical = +face_editor_head_pitch = +face_editor_head_yaw = +face_editor_head_roll = +face_enhancer_model = +face_enhancer_blend = +face_enhancer_weight = +face_swapper_model = +face_swapper_pixel_boost = +face_swapper_weight = +frame_colorizer_model = +frame_colorizer_size = +frame_colorizer_blend = +frame_enhancer_model = +frame_enhancer_blend = +lip_syncer_model = +lip_syncer_weight = + +[uis] +open_browser = +ui_layouts = +ui_workflow = + +[download] +download_providers = +download_scope = + +[benchmark] +benchmark_mode = +benchmark_resolutions = +benchmark_cycle_count = + +[execution] +execution_device_ids = +execution_providers = +execution_thread_count = + +[memory] +video_memory_strategy = +system_memory_limit = + +[misc] +log_level = +halt_on_error = diff --git a/facefusion.py b/facefusion.py new file mode 100644 index 0000000000000000000000000000000000000000..cdf4723548957587611d8ca16411cf0a5cb8c50c --- /dev/null +++ b/facefusion.py @@ -0,0 +1,10 @@ +#!/usr/bin/env python3 + +import os + +os.environ['OMP_NUM_THREADS'] = '1' + +from facefusion import core + +if __name__ == '__main__': + core.cli() diff --git a/facefusion/__init__.py b/facefusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/app_context.py b/facefusion/app_context.py new file mode 100644 index 0000000000000000000000000000000000000000..ee33b992913a845d691136b6eeef065c7a2578bf --- /dev/null +++ b/facefusion/app_context.py @@ -0,0 +1,16 @@ +import os +import sys + +from facefusion.types import AppContext + + +def detect_app_context() -> AppContext: + frame = sys._getframe(1) + + while frame: + if os.path.join('facefusion', 'jobs') in frame.f_code.co_filename: + return 'cli' + if os.path.join('facefusion', 'uis') in frame.f_code.co_filename: + return 'ui' + frame = frame.f_back + return 'cli' diff --git a/facefusion/args.py b/facefusion/args.py new file mode 100644 index 0000000000000000000000000000000000000000..88361f70ff5b6b7a60f83d2547cf48fecbfccd33 --- /dev/null +++ b/facefusion/args.py @@ -0,0 +1,131 @@ +from facefusion import state_manager +from facefusion.filesystem import get_file_name, is_video, resolve_file_paths +from facefusion.jobs import job_store +from facefusion.normalizer import normalize_fps, normalize_space +from facefusion.processors.core import get_processors_modules +from facefusion.types import ApplyStateItem, Args +from facefusion.vision import detect_video_fps + + +def reduce_step_args(args : Args) -> Args: + step_args =\ + { + key: args[key] for key in args if key in job_store.get_step_keys() + } + return step_args + + +def reduce_job_args(args : Args) -> Args: + job_args =\ + { + key: args[key] for key in args if key in job_store.get_job_keys() + } + return job_args + + +def collect_step_args() -> Args: + step_args =\ + { + key: state_manager.get_item(key) for key in job_store.get_step_keys() #type:ignore[arg-type] + } + return step_args + + +def collect_job_args() -> Args: + job_args =\ + { + key: state_manager.get_item(key) for key in job_store.get_job_keys() #type:ignore[arg-type] + } + return job_args + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + # general + apply_state_item('command', args.get('command')) + # paths + 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')) + apply_state_item('target_path', args.get('target_path')) + apply_state_item('output_path', args.get('output_path')) + # patterns + apply_state_item('source_pattern', args.get('source_pattern')) + apply_state_item('target_pattern', args.get('target_pattern')) + apply_state_item('output_pattern', args.get('output_pattern')) + # face detector + apply_state_item('face_detector_model', args.get('face_detector_model')) + apply_state_item('face_detector_size', args.get('face_detector_size')) + apply_state_item('face_detector_margin', normalize_space(args.get('face_detector_margin'))) + apply_state_item('face_detector_angles', args.get('face_detector_angles')) + apply_state_item('face_detector_score', args.get('face_detector_score')) + # face landmarker + apply_state_item('face_landmarker_model', args.get('face_landmarker_model')) + apply_state_item('face_landmarker_score', args.get('face_landmarker_score')) + # face selector + apply_state_item('face_selector_mode', args.get('face_selector_mode')) + apply_state_item('face_selector_order', args.get('face_selector_order')) + apply_state_item('face_selector_age_start', args.get('face_selector_age_start')) + apply_state_item('face_selector_age_end', args.get('face_selector_age_end')) + apply_state_item('face_selector_gender', args.get('face_selector_gender')) + apply_state_item('face_selector_race', args.get('face_selector_race')) + apply_state_item('reference_face_position', args.get('reference_face_position')) + apply_state_item('reference_face_distance', args.get('reference_face_distance')) + apply_state_item('reference_frame_number', args.get('reference_frame_number')) + # face masker + apply_state_item('face_occluder_model', args.get('face_occluder_model')) + apply_state_item('face_parser_model', args.get('face_parser_model')) + apply_state_item('face_mask_types', args.get('face_mask_types')) + apply_state_item('face_mask_areas', args.get('face_mask_areas')) + apply_state_item('face_mask_regions', args.get('face_mask_regions')) + apply_state_item('face_mask_blur', args.get('face_mask_blur')) + apply_state_item('face_mask_padding', normalize_space(args.get('face_mask_padding'))) + # voice extractor + apply_state_item('voice_extractor_model', args.get('voice_extractor_model')) + # frame extraction + 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('keep_temp', args.get('keep_temp')) + # output creation + apply_state_item('output_image_quality', args.get('output_image_quality')) + apply_state_item('output_image_scale', args.get('output_image_scale')) + apply_state_item('output_audio_encoder', args.get('output_audio_encoder')) + apply_state_item('output_audio_quality', args.get('output_audio_quality')) + apply_state_item('output_audio_volume', args.get('output_audio_volume')) + apply_state_item('output_video_encoder', args.get('output_video_encoder')) + apply_state_item('output_video_preset', args.get('output_video_preset')) + apply_state_item('output_video_quality', args.get('output_video_quality')) + apply_state_item('output_video_scale', args.get('output_video_scale')) + if args.get('output_video_fps') or is_video(args.get('target_path')): + 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) + # processors + available_processors = [ get_file_name(file_path) for file_path in resolve_file_paths('facefusion/processors/modules') ] + apply_state_item('processors', args.get('processors')) + for processor_module in get_processors_modules(available_processors): + processor_module.apply_args(args, apply_state_item) + # uis + apply_state_item('open_browser', args.get('open_browser')) + apply_state_item('ui_layouts', args.get('ui_layouts')) + apply_state_item('ui_workflow', args.get('ui_workflow')) + # execution + apply_state_item('execution_device_ids', args.get('execution_device_ids')) + apply_state_item('execution_providers', args.get('execution_providers')) + apply_state_item('execution_thread_count', args.get('execution_thread_count')) + # download + apply_state_item('download_providers', args.get('download_providers')) + apply_state_item('download_scope', args.get('download_scope')) + # benchmark + apply_state_item('benchmark_mode', args.get('benchmark_mode')) + apply_state_item('benchmark_resolutions', args.get('benchmark_resolutions')) + apply_state_item('benchmark_cycle_count', args.get('benchmark_cycle_count')) + # memory + apply_state_item('video_memory_strategy', args.get('video_memory_strategy')) + apply_state_item('system_memory_limit', args.get('system_memory_limit')) + # misc + apply_state_item('log_level', args.get('log_level')) + apply_state_item('halt_on_error', args.get('halt_on_error')) + # jobs + apply_state_item('job_id', args.get('job_id')) + apply_state_item('job_status', args.get('job_status')) + apply_state_item('step_index', args.get('step_index')) diff --git a/facefusion/audio.py b/facefusion/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..4e7b13deb0454dda877ebbc54141e1693e15a285 --- /dev/null +++ b/facefusion/audio.py @@ -0,0 +1,143 @@ +from functools import lru_cache +from typing import Any, List, Optional + +import numpy +import scipy +from numpy.typing import NDArray + +from facefusion.ffmpeg import read_audio_buffer +from facefusion.filesystem import is_audio +from facefusion.types import Audio, AudioFrame, Fps, Mel, MelFilterBank, Spectrogram +from facefusion.voice_extractor import batch_extract_voice + + +@lru_cache(maxsize = 64) +def read_static_audio(audio_path : str, fps : Fps) -> Optional[List[AudioFrame]]: + return read_audio(audio_path, fps) + + +def read_audio(audio_path : str, fps : Fps) -> Optional[List[AudioFrame]]: + audio_sample_rate = 48000 + audio_sample_size = 16 + audio_channel_total = 2 + + if is_audio(audio_path): + audio_buffer = read_audio_buffer(audio_path, audio_sample_rate, audio_sample_size, audio_channel_total) + audio = numpy.frombuffer(audio_buffer, dtype = numpy.int16).reshape(-1, 2) + audio = prepare_audio(audio) + spectrogram = create_spectrogram(audio) + audio_frames = extract_audio_frames(spectrogram, fps) + return audio_frames + return None + + +@lru_cache(maxsize = 64) +def read_static_voice(audio_path : str, fps : Fps) -> Optional[List[AudioFrame]]: + return read_voice(audio_path, fps) + + +def read_voice(audio_path : str, fps : Fps) -> Optional[List[AudioFrame]]: + voice_sample_rate = 48000 + voice_sample_size = 16 + voice_channel_total = 2 + voice_chunk_size = 240 * 1024 + voice_step_size = 180 * 1024 + + if is_audio(audio_path): + audio_buffer = read_audio_buffer(audio_path, voice_sample_rate, voice_sample_size, voice_channel_total) + audio = numpy.frombuffer(audio_buffer, dtype = numpy.int16).reshape(-1, 2) + audio = batch_extract_voice(audio, voice_chunk_size, voice_step_size) + audio = prepare_voice(audio) + spectrogram = create_spectrogram(audio) + audio_frames = extract_audio_frames(spectrogram, fps) + return audio_frames + return None + + +def get_audio_frame(audio_path : str, fps : Fps, frame_number : int = 0) -> Optional[AudioFrame]: + if is_audio(audio_path): + audio_frames = read_static_audio(audio_path, fps) + if frame_number in range(len(audio_frames)): + return audio_frames[frame_number] + return None + + +def extract_audio_frames(spectrogram : Spectrogram, fps : Fps) -> List[AudioFrame]: + audio_frames = [] + mel_filter_total = 80 + audio_step_size = 16 + indices = numpy.arange(0, spectrogram.shape[1], mel_filter_total / fps).astype(numpy.int16) + indices = indices[indices >= audio_step_size] + + for index in indices: + start = max(0, index - audio_step_size) + audio_frames.append(spectrogram[:, start:index]) + + return audio_frames + + +def get_voice_frame(audio_path : str, fps : Fps, frame_number : int = 0) -> Optional[AudioFrame]: + if is_audio(audio_path): + voice_frames = read_static_voice(audio_path, fps) + if frame_number in range(len(voice_frames)): + return voice_frames[frame_number] + return None + + +def create_empty_audio_frame() -> AudioFrame: + mel_filter_total = 80 + audio_step_size = 16 + audio_frame = numpy.zeros((mel_filter_total, audio_step_size)).astype(numpy.int16) + return audio_frame + + +def prepare_audio(audio : Audio) -> Audio: + if audio.ndim > 1: + audio = numpy.mean(audio, axis = 1) + audio = audio / numpy.max(numpy.abs(audio), axis = 0) + audio = scipy.signal.lfilter([ 1.0, -0.97 ], [ 1.0 ], audio) + return audio + + +def prepare_voice(audio : Audio) -> Audio: + audio_sample_rate = 48000 + audio_resample_rate = 16000 + audio_resample_factor = round(len(audio) * audio_resample_rate / audio_sample_rate) + audio = scipy.signal.resample(audio, audio_resample_factor) + audio = prepare_audio(audio) + return audio + + +def convert_hertz_to_mel(hertz : float) -> float: + return 2595 * numpy.log10(1 + hertz / 700) + + +def convert_mel_to_hertz(mel : Mel) -> NDArray[Any]: + return 700 * (10 ** (mel / 2595) - 1) + + +def create_mel_filter_bank() -> MelFilterBank: + audio_sample_rate = 16000 + audio_frequency_min = 55.0 + audio_frequency_max = 7600.0 + mel_filter_total = 80 + mel_bin_total = 800 + mel_filter_bank = numpy.zeros((mel_filter_total, mel_bin_total // 2 + 1)) + mel_frequency_range = numpy.linspace(convert_hertz_to_mel(audio_frequency_min), convert_hertz_to_mel(audio_frequency_max), mel_filter_total + 2) + indices = numpy.floor((mel_bin_total + 1) * convert_mel_to_hertz(mel_frequency_range) / audio_sample_rate).astype(numpy.int16) + + for index in range(mel_filter_total): + start = indices[index] + end = indices[index + 1] + mel_filter_bank[index, start:end] = scipy.signal.windows.triang(end - start) + + return mel_filter_bank + + +def create_spectrogram(audio : Audio) -> Spectrogram: + mel_bin_total = 800 + mel_bin_overlap = 600 + mel_filter_bank = create_mel_filter_bank() + spectrogram = scipy.signal.stft(audio, nperseg = mel_bin_total, nfft = mel_bin_total, noverlap = mel_bin_overlap)[2] + spectrogram = numpy.dot(mel_filter_bank, numpy.abs(spectrogram)) + return spectrogram diff --git a/facefusion/benchmarker.py b/facefusion/benchmarker.py new file mode 100644 index 0000000000000000000000000000000000000000..b2b3bbf5556d6694665f03e1358c204b7a84060e --- /dev/null +++ b/facefusion/benchmarker.py @@ -0,0 +1,111 @@ +import hashlib +import os +import statistics +import tempfile +from time import perf_counter +from typing import Iterator, List + +import facefusion.choices +from facefusion import content_analyser, core, state_manager +from facefusion.cli_helper import render_table +from facefusion.download import conditional_download, resolve_download_url +from facefusion.face_store import clear_static_faces +from facefusion.filesystem import get_file_extension +from facefusion.types import BenchmarkCycleSet +from facefusion.vision import count_video_frame_total, detect_video_fps + + +def pre_check() -> bool: + conditional_download('.assets/examples', + [ + resolve_download_url('examples-3.0.0', 'source.jpg'), + resolve_download_url('examples-3.0.0', 'source.mp3'), + resolve_download_url('examples-3.0.0', 'target-240p.mp4'), + resolve_download_url('examples-3.0.0', 'target-360p.mp4'), + resolve_download_url('examples-3.0.0', 'target-540p.mp4'), + resolve_download_url('examples-3.0.0', 'target-720p.mp4'), + resolve_download_url('examples-3.0.0', 'target-1080p.mp4'), + resolve_download_url('examples-3.0.0', 'target-1440p.mp4'), + resolve_download_url('examples-3.0.0', 'target-2160p.mp4') + ]) + return True + + +def run() -> Iterator[List[BenchmarkCycleSet]]: + benchmark_resolutions = state_manager.get_item('benchmark_resolutions') + benchmark_cycle_count = state_manager.get_item('benchmark_cycle_count') + + state_manager.init_item('source_paths', [ '.assets/examples/source.jpg', '.assets/examples/source.mp3' ]) + state_manager.init_item('face_landmarker_score', 0) + state_manager.init_item('temp_frame_format', 'bmp') + state_manager.init_item('output_audio_volume', 0) + state_manager.init_item('output_video_preset', 'ultrafast') + state_manager.init_item('video_memory_strategy', 'tolerant') + + benchmarks = [] + target_paths = [ facefusion.choices.benchmark_set.get(benchmark_resolution) for benchmark_resolution in benchmark_resolutions if benchmark_resolution in facefusion.choices.benchmark_set ] + + for target_path in target_paths: + state_manager.init_item('target_path', target_path) + state_manager.init_item('output_path', suggest_output_path(state_manager.get_item('target_path'))) + benchmarks.append(cycle(benchmark_cycle_count)) + yield benchmarks + + +def cycle(cycle_count : int) -> BenchmarkCycleSet: + process_times = [] + video_frame_total = count_video_frame_total(state_manager.get_item('target_path')) + state_manager.init_item('output_video_fps', detect_video_fps(state_manager.get_item('target_path'))) + + if state_manager.get_item('benchmark_mode') == 'warm': + core.conditional_process() + + for index in range(cycle_count): + if state_manager.get_item('benchmark_mode') == 'cold': + content_analyser.analyse_image.cache_clear() + content_analyser.analyse_video.cache_clear() + clear_static_faces() + + start_time = perf_counter() + core.conditional_process() + end_time = perf_counter() + process_times.append(end_time - start_time) + + average_run = round(statistics.mean(process_times), 2) + fastest_run = round(min(process_times), 2) + slowest_run = round(max(process_times), 2) + relative_fps = round(video_frame_total * cycle_count / sum(process_times), 2) + + return\ + { + 'target_path': state_manager.get_item('target_path'), + 'cycle_count': cycle_count, + 'average_run': average_run, + 'fastest_run': fastest_run, + 'slowest_run': slowest_run, + 'relative_fps': relative_fps + } + + +def suggest_output_path(target_path : str) -> str: + target_file_extension = get_file_extension(target_path) + return os.path.join(tempfile.gettempdir(), hashlib.sha1().hexdigest()[:8] + target_file_extension) + + +def render() -> None: + benchmarks = [] + headers =\ + [ + 'target_path', + 'cycle_count', + 'average_run', + 'fastest_run', + 'slowest_run', + 'relative_fps' + ] + + for benchmark in run(): + benchmarks = benchmark + + contents = [ list(benchmark_set.values()) for benchmark_set in benchmarks ] + render_table(headers, contents) diff --git a/facefusion/camera_manager.py b/facefusion/camera_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..903b511533a2421894c5831a0d745c3e33260bee --- /dev/null +++ b/facefusion/camera_manager.py @@ -0,0 +1,53 @@ +from typing import List + +import cv2 + +from facefusion.types import CameraPoolSet + +CAMERA_POOL_SET : CameraPoolSet =\ +{ + 'capture': {} +} + + +def get_local_camera_capture(camera_id : int) -> cv2.VideoCapture: + camera_key = str(camera_id) + + if camera_key not in CAMERA_POOL_SET.get('capture'): + camera_capture = cv2.VideoCapture(camera_id) + + if camera_capture.isOpened(): + CAMERA_POOL_SET['capture'][camera_key] = camera_capture + + return CAMERA_POOL_SET.get('capture').get(camera_key) + + +def get_remote_camera_capture(camera_url : str) -> cv2.VideoCapture: + if camera_url not in CAMERA_POOL_SET.get('capture'): + camera_capture = cv2.VideoCapture(camera_url) + + if camera_capture.isOpened(): + CAMERA_POOL_SET['capture'][camera_url] = camera_capture + + return CAMERA_POOL_SET.get('capture').get(camera_url) + + +def clear_camera_pool() -> None: + for camera_capture in CAMERA_POOL_SET.get('capture').values(): + camera_capture.release() + + CAMERA_POOL_SET['capture'].clear() + + +def detect_local_camera_ids(id_start : int, id_end : int) -> List[int]: + local_camera_ids = [] + + for camera_id in range(id_start, id_end): + cv2.setLogLevel(0) + camera_capture = get_local_camera_capture(camera_id) + cv2.setLogLevel(3) + + if camera_capture and camera_capture.isOpened(): + local_camera_ids.append(camera_id) + + return local_camera_ids diff --git a/facefusion/choices.py b/facefusion/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..71c75d5c431c26cf30346d00dd9de4ea14f1b850 --- /dev/null +++ b/facefusion/choices.py @@ -0,0 +1,169 @@ +import logging +from typing import List, Sequence + +from facefusion.common_helper import create_float_range, create_int_range +from facefusion.types import Angle, AudioEncoder, AudioFormat, AudioTypeSet, BenchmarkMode, BenchmarkResolution, BenchmarkSet, DownloadProvider, DownloadProviderSet, DownloadScope, EncoderSet, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskArea, FaceMaskAreaSet, FaceMaskRegion, FaceMaskRegionSet, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorMode, FaceSelectorOrder, Gender, ImageFormat, ImageTypeSet, JobStatus, LogLevel, LogLevelSet, Race, Score, TempFrameFormat, UiWorkflow, VideoEncoder, VideoFormat, VideoMemoryStrategy, VideoPreset, VideoTypeSet, VoiceExtractorModel + +face_detector_set : FaceDetectorSet =\ +{ + 'many': [ '640x640' ], + 'retinaface': [ '160x160', '320x320', '480x480', '512x512', '640x640' ], + 'scrfd': [ '160x160', '320x320', '480x480', '512x512', '640x640' ], + 'yolo_face': [ '640x640' ], + 'yunet': [ '640x640' ] +} +face_detector_models : List[FaceDetectorModel] = list(face_detector_set.keys()) +face_landmarker_models : List[FaceLandmarkerModel] = [ 'many', '2dfan4', 'peppa_wutz' ] +face_selector_modes : List[FaceSelectorMode] = [ 'many', 'one', 'reference' ] +face_selector_orders : List[FaceSelectorOrder] = [ 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best' ] +face_selector_genders : List[Gender] = [ 'female', 'male' ] +face_selector_races : List[Race] = [ 'white', 'black', 'latino', 'asian', 'indian', 'arabic' ] +face_occluder_models : List[FaceOccluderModel] = [ 'many', 'xseg_1', 'xseg_2', 'xseg_3' ] +face_parser_models : List[FaceParserModel] = [ 'bisenet_resnet_18', 'bisenet_resnet_34' ] +face_mask_types : List[FaceMaskType] = [ 'box', 'occlusion', 'area', 'region' ] +face_mask_area_set : FaceMaskAreaSet =\ +{ + 'upper-face': [ 0, 1, 2, 31, 32, 33, 34, 35, 14, 15, 16, 26, 25, 24, 23, 22, 21, 20, 19, 18, 17 ], + 'lower-face': [ 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 35, 34, 33, 32, 31 ], + 'mouth': [ 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67 ] +} +face_mask_region_set : FaceMaskRegionSet =\ +{ + 'skin': 1, + 'left-eyebrow': 2, + 'right-eyebrow': 3, + 'left-eye': 4, + 'right-eye': 5, + 'glasses': 6, + 'nose': 10, + 'mouth': 11, + 'upper-lip': 12, + 'lower-lip': 13 +} +face_mask_areas : List[FaceMaskArea] = list(face_mask_area_set.keys()) +face_mask_regions : List[FaceMaskRegion] = list(face_mask_region_set.keys()) + +voice_extractor_models : List[VoiceExtractorModel] = [ 'kim_vocal_1', 'kim_vocal_2', 'uvr_mdxnet' ] + +audio_type_set : AudioTypeSet =\ +{ + 'flac': 'audio/flac', + 'm4a': 'audio/mp4', + 'mp3': 'audio/mpeg', + 'ogg': 'audio/ogg', + 'opus': 'audio/opus', + 'wav': 'audio/x-wav' +} +image_type_set : ImageTypeSet =\ +{ + 'bmp': 'image/bmp', + 'jpeg': 'image/jpeg', + 'png': 'image/png', + 'tiff': 'image/tiff', + 'webp': 'image/webp' +} +video_type_set : VideoTypeSet =\ +{ + 'avi': 'video/x-msvideo', + 'm4v': 'video/mp4', + 'mkv': 'video/x-matroska', + 'mp4': 'video/mp4', + 'mpeg': 'video/mpeg', + 'mov': 'video/quicktime', + 'mxf': 'application/mxf', + 'webm': 'video/webm', + 'wmv': 'video/x-ms-wmv' +} +audio_formats : List[AudioFormat] = list(audio_type_set.keys()) +image_formats : List[ImageFormat] = list(image_type_set.keys()) +video_formats : List[VideoFormat] = list(video_type_set.keys()) +temp_frame_formats : List[TempFrameFormat] = [ 'bmp', 'jpeg', 'png', 'tiff' ] + +output_encoder_set : EncoderSet =\ +{ + 'audio': [ 'flac', 'aac', 'libmp3lame', 'libopus', 'libvorbis', 'pcm_s16le', 'pcm_s32le' ], + 'video': [ 'libx264', 'libx264rgb', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc', 'h264_amf', 'hevc_amf', 'h264_qsv', 'hevc_qsv', 'h264_videotoolbox', 'hevc_videotoolbox', 'rawvideo' ] +} +output_audio_encoders : List[AudioEncoder] = output_encoder_set.get('audio') +output_video_encoders : List[VideoEncoder] = output_encoder_set.get('video') +output_video_presets : List[VideoPreset] = [ 'ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow' ] + +benchmark_modes : List[BenchmarkMode] = [ 'warm', 'cold' ] +benchmark_set : BenchmarkSet =\ +{ + '240p': '.assets/examples/target-240p.mp4', + '360p': '.assets/examples/target-360p.mp4', + '540p': '.assets/examples/target-540p.mp4', + '720p': '.assets/examples/target-720p.mp4', + '1080p': '.assets/examples/target-1080p.mp4', + '1440p': '.assets/examples/target-1440p.mp4', + '2160p': '.assets/examples/target-2160p.mp4' +} +benchmark_resolutions : List[BenchmarkResolution] = list(benchmark_set.keys()) + +execution_provider_set : ExecutionProviderSet =\ +{ + 'cuda': 'CUDAExecutionProvider', + 'tensorrt': 'TensorrtExecutionProvider', + 'directml': 'DmlExecutionProvider', + 'rocm': 'ROCMExecutionProvider', + 'migraphx': 'MIGraphXExecutionProvider', + 'openvino': 'OpenVINOExecutionProvider', + 'coreml': 'CoreMLExecutionProvider', + 'cpu': 'CPUExecutionProvider' +} +execution_providers : List[ExecutionProvider] = list(execution_provider_set.keys()) +download_provider_set : DownloadProviderSet =\ +{ + 'github': + { + 'urls': + [ + 'https://github.com' + ], + 'path': '/facefusion/facefusion-assets/releases/download/{base_name}/{file_name}' + }, + 'huggingface': + { + 'urls': + [ + 'https://huggingface.co', + 'https://hf-mirror.com' + ], + 'path': '/facefusion/{base_name}/resolve/main/{file_name}' + } +} +download_providers : List[DownloadProvider] = list(download_provider_set.keys()) +download_scopes : List[DownloadScope] = [ 'lite', 'full' ] + +video_memory_strategies : List[VideoMemoryStrategy] = [ 'strict', 'moderate', 'tolerant' ] + +log_level_set : LogLevelSet =\ +{ + 'error': logging.ERROR, + 'warn': logging.WARNING, + 'info': logging.INFO, + 'debug': logging.DEBUG +} +log_levels : List[LogLevel] = list(log_level_set.keys()) + +ui_workflows : List[UiWorkflow] = [ 'instant_runner', 'job_runner', 'job_manager' ] +job_statuses : List[JobStatus] = [ 'drafted', 'queued', 'completed', 'failed' ] + +benchmark_cycle_count_range : Sequence[int] = create_int_range(1, 10, 1) +execution_thread_count_range : Sequence[int] = create_int_range(1, 32, 1) +system_memory_limit_range : Sequence[int] = create_int_range(0, 128, 4) +face_detector_margin_range : Sequence[int] = create_int_range(0, 100, 1) +face_detector_angles : Sequence[Angle] = create_int_range(0, 270, 90) +face_detector_score_range : Sequence[Score] = create_float_range(0.0, 1.0, 0.05) +face_landmarker_score_range : Sequence[Score] = create_float_range(0.0, 1.0, 0.05) +face_mask_blur_range : Sequence[float] = create_float_range(0.0, 1.0, 0.05) +face_mask_padding_range : Sequence[int] = create_int_range(0, 100, 1) +face_selector_age_range : Sequence[int] = create_int_range(0, 100, 1) +reference_face_distance_range : Sequence[float] = create_float_range(0.0, 1.0, 0.05) +output_image_quality_range : Sequence[int] = create_int_range(0, 100, 1) +output_image_scale_range : Sequence[float] = create_float_range(0.25, 8.0, 0.25) +output_audio_quality_range : Sequence[int] = create_int_range(0, 100, 1) +output_audio_volume_range : Sequence[int] = create_int_range(0, 100, 1) +output_video_quality_range : Sequence[int] = create_int_range(0, 100, 1) +output_video_scale_range : Sequence[float] = create_float_range(0.25, 8.0, 0.25) diff --git a/facefusion/cli_helper.py b/facefusion/cli_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..01418dca1132fba54c37470798d7026d98a8ee77 --- /dev/null +++ b/facefusion/cli_helper.py @@ -0,0 +1,35 @@ +from typing import List, Tuple + +from facefusion.logger import get_package_logger +from facefusion.types import TableContent, TableHeader + + +def render_table(headers : List[TableHeader], contents : List[List[TableContent]]) -> None: + package_logger = get_package_logger() + table_column, table_separator = create_table_parts(headers, contents) + + package_logger.critical(table_separator) + package_logger.critical(table_column.format(*headers)) + package_logger.critical(table_separator) + + for content in contents: + content = [ str(value) for value in content ] + package_logger.critical(table_column.format(*content)) + + package_logger.critical(table_separator) + + +def create_table_parts(headers : List[TableHeader], contents : List[List[TableContent]]) -> Tuple[str, str]: + column_parts = [] + separator_parts = [] + widths = [ len(header) for header in headers ] + + for content in contents: + for index, value in enumerate(content): + widths[index] = max(widths[index], len(str(value))) + + for width in widths: + column_parts.append('{:<' + str(width) + '}') + separator_parts.append('-' * width) + + return '| ' + ' | '.join(column_parts) + ' |', '+-' + '-+-'.join(separator_parts) + '-+' diff --git a/facefusion/common_helper.py b/facefusion/common_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..5527c86cf877e8f51bbae45f08992c5c2eb5b07b --- /dev/null +++ b/facefusion/common_helper.py @@ -0,0 +1,84 @@ +import platform +from typing import Any, Iterable, Optional, Reversible, Sequence + + +def is_linux() -> bool: + return platform.system().lower() == 'linux' + + +def is_macos() -> bool: + return platform.system().lower() == 'darwin' + + +def is_windows() -> bool: + return platform.system().lower() == 'windows' + + +def create_int_metavar(int_range : Sequence[int]) -> str: + return '[' + str(int_range[0]) + '..' + str(int_range[-1]) + ':' + str(calculate_int_step(int_range)) + ']' + + +def create_float_metavar(float_range : Sequence[float]) -> str: + return '[' + str(float_range[0]) + '..' + str(float_range[-1]) + ':' + str(calculate_float_step(float_range)) + ']' + + +def create_int_range(start : int, end : int, step : int) -> Sequence[int]: + int_range = [] + current = start + + while current <= end: + int_range.append(current) + current += step + return int_range + + +def create_float_range(start : float, end : float, step : float) -> Sequence[float]: + float_range = [] + current = start + + while current <= end: + float_range.append(round(current, 2)) + current = round(current + step, 2) + return float_range + + +def calculate_int_step(int_range : Sequence[int]) -> int: + return int_range[1] - int_range[0] + + +def calculate_float_step(float_range : Sequence[float]) -> float: + return round(float_range[1] - float_range[0], 2) + + +def cast_int(value : Any) -> Optional[int]: + try: + return int(value) + except (ValueError, TypeError): + return None + + +def cast_float(value : Any) -> Optional[float]: + try: + return float(value) + except (ValueError, TypeError): + return None + + +def cast_bool(value : Any) -> Optional[bool]: + if value == 'True': + return True + if value == 'False': + return False + return None + + +def get_first(__list__ : Any) -> Any: + if isinstance(__list__, Iterable): + return next(iter(__list__), None) + return None + + +def get_last(__list__ : Any) -> Any: + if isinstance(__list__, Reversible): + return next(reversed(__list__), None) + return None diff --git a/facefusion/config.py b/facefusion/config.py new file mode 100644 index 0000000000000000000000000000000000000000..132b383a159cdfd5f386bb4e28dea26ef95db98f --- /dev/null +++ b/facefusion/config.py @@ -0,0 +1,74 @@ +from configparser import ConfigParser +from typing import List, Optional + +from facefusion import state_manager +from facefusion.common_helper import cast_bool, cast_float, cast_int + +CONFIG_PARSER = None + + +def get_config_parser() -> ConfigParser: + global CONFIG_PARSER + + if CONFIG_PARSER is None: + CONFIG_PARSER = ConfigParser() + CONFIG_PARSER.read(state_manager.get_item('config_path'), encoding = 'utf-8') + return CONFIG_PARSER + + +def clear_config_parser() -> None: + global CONFIG_PARSER + + CONFIG_PARSER = None + + +def get_str_value(section : str, option : str, fallback : Optional[str] = None) -> Optional[str]: + config_parser = get_config_parser() + + if config_parser.has_option(section, option) and config_parser.get(section, option).strip(): + return config_parser.get(section, option) + return fallback + + +def get_int_value(section : str, option : str, fallback : Optional[str] = None) -> Optional[int]: + config_parser = get_config_parser() + + if config_parser.has_option(section, option) and config_parser.get(section, option).strip(): + return config_parser.getint(section, option) + return cast_int(fallback) + + +def get_float_value(section : str, option : str, fallback : Optional[str] = None) -> Optional[float]: + config_parser = get_config_parser() + + if config_parser.has_option(section, option) and config_parser.get(section, option).strip(): + return config_parser.getfloat(section, option) + return cast_float(fallback) + + +def get_bool_value(section : str, option : str, fallback : Optional[str] = None) -> Optional[bool]: + config_parser = get_config_parser() + + if config_parser.has_option(section, option) and config_parser.get(section, option).strip(): + return config_parser.getboolean(section, option) + return cast_bool(fallback) + + +def get_str_list(section : str, option : str, fallback : Optional[str] = None) -> Optional[List[str]]: + config_parser = get_config_parser() + + if config_parser.has_option(section, option) and config_parser.get(section, option).strip(): + return config_parser.get(section, option).split() + if fallback: + return fallback.split() + return None + + +def get_int_list(section : str, option : str, fallback : Optional[str] = None) -> Optional[List[int]]: + config_parser = get_config_parser() + + if config_parser.has_option(section, option) and config_parser.get(section, option).strip(): + return list(map(int, config_parser.get(section, option).split())) + if fallback: + return list(map(int, fallback.split())) + return None diff --git a/facefusion/content_analyser.py b/facefusion/content_analyser.py new file mode 100644 index 0000000000000000000000000000000000000000..0140989679e336cfe6f2502c5f9ee8272cf92a5b --- /dev/null +++ b/facefusion/content_analyser.py @@ -0,0 +1,247 @@ +from functools import lru_cache +from typing import List, Tuple + +import numpy +from tqdm import tqdm + +from facefusion import inference_manager, state_manager, translator +from facefusion.common_helper import is_macos +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.execution import has_execution_provider +from facefusion.filesystem import resolve_relative_path +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.types import Detection, DownloadScope, DownloadSet, ExecutionProvider, Fps, InferencePool, ModelSet, VisionFrame +from facefusion.vision import detect_video_fps, fit_contain_frame, read_image, read_video_frame + +STREAM_COUNTER = 0 + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'nsfw_1': + { + '__metadata__': + { + 'vendor': 'EraX', + 'license': 'Apache-2.0', + 'year': 2024 + }, + 'hashes': + { + 'content_analyser': + { + 'url': resolve_download_url('models-3.3.0', 'nsfw_1.hash'), + 'path': resolve_relative_path('../.assets/models/nsfw_1.hash') + } + }, + 'sources': + { + 'content_analyser': + { + 'url': resolve_download_url('models-3.3.0', 'nsfw_1.onnx'), + 'path': resolve_relative_path('../.assets/models/nsfw_1.onnx') + } + }, + 'size': (640, 640), + 'mean': (0.0, 0.0, 0.0), + 'standard_deviation': (1.0, 1.0, 1.0) + }, + 'nsfw_2': + { + '__metadata__': + { + 'vendor': 'Marqo', + 'license': 'Apache-2.0', + 'year': 2024 + }, + 'hashes': + { + 'content_analyser': + { + 'url': resolve_download_url('models-3.3.0', 'nsfw_2.hash'), + 'path': resolve_relative_path('../.assets/models/nsfw_2.hash') + } + }, + 'sources': + { + 'content_analyser': + { + 'url': resolve_download_url('models-3.3.0', 'nsfw_2.onnx'), + 'path': resolve_relative_path('../.assets/models/nsfw_2.onnx') + } + }, + 'size': (384, 384), + 'mean': (0.5, 0.5, 0.5), + 'standard_deviation': (0.5, 0.5, 0.5) + }, + 'nsfw_3': + { + '__metadata__': + { + 'vendor': 'Freepik', + 'license': 'MIT', + 'year': 2025 + }, + 'hashes': + { + 'content_analyser': + { + 'url': resolve_download_url('models-3.3.0', 'nsfw_3.hash'), + 'path': resolve_relative_path('../.assets/models/nsfw_3.hash') + } + }, + 'sources': + { + 'content_analyser': + { + 'url': resolve_download_url('models-3.3.0', 'nsfw_3.onnx'), + 'path': resolve_relative_path('../.assets/models/nsfw_3.onnx') + } + }, + 'size': (448, 448), + 'mean': (0.48145466, 0.4578275, 0.40821073), + 'standard_deviation': (0.26862954, 0.26130258, 0.27577711) + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ] + _, model_source_set = collect_model_downloads() + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def resolve_execution_providers() -> List[ExecutionProvider]: + if is_macos() and has_execution_provider('coreml'): + return [ 'cpu' ] + return state_manager.get_item('execution_providers') + + +def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_set = create_static_model_set('full') + model_hash_set = {} + model_source_set = {} + + for content_analyser_model in [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ]: + model_hash_set[content_analyser_model] = model_set.get(content_analyser_model).get('hashes').get('content_analyser') + model_source_set[content_analyser_model] = model_set.get(content_analyser_model).get('sources').get('content_analyser') + + return model_hash_set, model_source_set + + +def pre_check() -> bool: + model_hash_set, model_source_set = collect_model_downloads() + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def analyse_stream(vision_frame : VisionFrame, video_fps : Fps) -> bool: + global STREAM_COUNTER + + STREAM_COUNTER = STREAM_COUNTER + 1 + if STREAM_COUNTER % int(video_fps) == 0: + return analyse_frame(vision_frame) + return False + + +def analyse_frame(vision_frame : VisionFrame) -> bool: + return detect_nsfw(vision_frame) + + +@lru_cache() +def analyse_image(image_path : str) -> bool: + vision_frame = read_image(image_path) + return analyse_frame(vision_frame) + + +@lru_cache() +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) + rate = 0.0 + total = 0 + counter = 0 + + 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) + total += 1 + + if analyse_frame(vision_frame): + counter += 1 + + if counter > 0 and total > 0: + rate = counter / total * 100 + + progress.set_postfix(rate = rate) + progress.update() + + return bool(rate > 10.0) + + +def detect_nsfw(vision_frame : VisionFrame) -> bool: + is_nsfw_1 = detect_with_nsfw_1(vision_frame) + is_nsfw_2 = detect_with_nsfw_2(vision_frame) + is_nsfw_3 = detect_with_nsfw_3(vision_frame) + + return is_nsfw_1 and is_nsfw_2 or is_nsfw_1 and is_nsfw_3 or is_nsfw_2 and is_nsfw_3 + + +def detect_with_nsfw_1(vision_frame : VisionFrame) -> bool: + detect_vision_frame = prepare_detect_frame(vision_frame, 'nsfw_1') + detection = forward_nsfw(detect_vision_frame, 'nsfw_1') + detection_score = numpy.max(numpy.amax(detection[:, 4:], axis = 1)) + return bool(detection_score > 0.2) + + +def detect_with_nsfw_2(vision_frame : VisionFrame) -> bool: + detect_vision_frame = prepare_detect_frame(vision_frame, 'nsfw_2') + detection = forward_nsfw(detect_vision_frame, 'nsfw_2') + detection_score = detection[0] - detection[1] + return bool(detection_score > 0.25) + + +def detect_with_nsfw_3(vision_frame : VisionFrame) -> bool: + detect_vision_frame = prepare_detect_frame(vision_frame, 'nsfw_3') + detection = forward_nsfw(detect_vision_frame, 'nsfw_3') + detection_score = (detection[2] + detection[3]) - (detection[0] + detection[1]) + return bool(detection_score > 10.5) + + +def forward_nsfw(vision_frame : VisionFrame, model_name : str) -> Detection: + content_analyser = get_inference_pool().get(model_name) + + with conditional_thread_semaphore(): + detection = content_analyser.run(None, + { + 'input': vision_frame + })[0] + + if model_name in [ 'nsfw_2', 'nsfw_3' ]: + return detection[0] + + return detection + + +def prepare_detect_frame(temp_vision_frame : VisionFrame, model_name : str) -> VisionFrame: + model_set = create_static_model_set('full').get(model_name) + model_size = model_set.get('size') + model_mean = model_set.get('mean') + model_standard_deviation = model_set.get('standard_deviation') + + detect_vision_frame = fit_contain_frame(temp_vision_frame, model_size) + detect_vision_frame = detect_vision_frame[:, :, ::-1] / 255.0 + detect_vision_frame -= model_mean + detect_vision_frame /= model_standard_deviation + detect_vision_frame = numpy.expand_dims(detect_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) + return detect_vision_frame diff --git a/facefusion/core.py b/facefusion/core.py new file mode 100644 index 0000000000000000000000000000000000000000..79a7dc6142a096857d119a162f39ada24bfe2da4 --- /dev/null +++ b/facefusion/core.py @@ -0,0 +1,350 @@ +import inspect +import itertools +import shutil +import signal +import sys +from time import time + +from facefusion import benchmarker, cli_helper, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, hash_helper, logger, state_manager, translator, voice_extractor +from facefusion.args import apply_args, collect_job_args, reduce_job_args, reduce_step_args +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.exit_helper import hard_exit, signal_exit +from facefusion.filesystem import get_file_extension, get_file_name, is_image, is_video, resolve_file_paths, resolve_file_pattern +from facefusion.jobs import job_helper, job_manager, job_runner +from facefusion.jobs.job_list import compose_job_list +from facefusion.memory import limit_system_memory +from facefusion.processors.core import get_processors_modules +from facefusion.program import create_program +from facefusion.program_helper import validate_args +from facefusion.types import Args, ErrorCode +from facefusion.workflows import image_to_image, image_to_video + + +def cli() -> None: + if pre_check(): + signal.signal(signal.SIGINT, signal_exit) + program = create_program() + + if validate_args(program): + args = vars(program.parse_args()) + apply_args(args, state_manager.init_item) + + if state_manager.get_item('command'): + logger.init(state_manager.get_item('log_level')) + route(args) + else: + program.print_help() + else: + hard_exit(2) + else: + hard_exit(2) + + +def route(args : Args) -> None: + system_memory_limit = state_manager.get_item('system_memory_limit') + + if system_memory_limit and system_memory_limit > 0: + limit_system_memory(system_memory_limit) + + if state_manager.get_item('command') == 'force-download': + error_code = force_download() + hard_exit(error_code) + + if state_manager.get_item('command') == 'benchmark': + if not common_pre_check() or not processors_pre_check() or not benchmarker.pre_check(): + hard_exit(2) + benchmarker.render() + + if state_manager.get_item('command') in [ 'job-list', 'job-create', 'job-submit', 'job-submit-all', 'job-delete', 'job-delete-all', 'job-add-step', 'job-remix-step', 'job-insert-step', 'job-remove-step' ]: + if not job_manager.init_jobs(state_manager.get_item('jobs_path')): + hard_exit(1) + error_code = route_job_manager(args) + hard_exit(error_code) + + if state_manager.get_item('command') == 'run': + import facefusion.uis.core as ui + + if not common_pre_check() or not processors_pre_check(): + hard_exit(2) + for ui_layout in ui.get_ui_layouts_modules(state_manager.get_item('ui_layouts')): + if not ui_layout.pre_check(): + hard_exit(2) + ui.init() + ui.launch() + + if state_manager.get_item('command') == 'headless-run': + if not job_manager.init_jobs(state_manager.get_item('jobs_path')): + hard_exit(1) + error_code = process_headless(args) + hard_exit(error_code) + + if state_manager.get_item('command') == 'batch-run': + if not job_manager.init_jobs(state_manager.get_item('jobs_path')): + hard_exit(1) + error_code = process_batch(args) + hard_exit(error_code) + + if state_manager.get_item('command') in [ 'job-run', 'job-run-all', 'job-retry', 'job-retry-all' ]: + if not job_manager.init_jobs(state_manager.get_item('jobs_path')): + hard_exit(1) + error_code = route_job_runner() + hard_exit(error_code) + + +def pre_check() -> bool: + if sys.version_info < (3, 10): + 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 + return True + + +def common_pre_check() -> bool: + common_modules =\ + [ + content_analyser, + face_classifier, + face_detector, + face_landmarker, + face_masker, + face_recognizer, + voice_extractor + ] + + content_analyser_content = inspect.getsource(content_analyser).encode() + content_analyser_hash = hash_helper.create_hash(content_analyser_content) + + return all(module.pre_check() for module in common_modules) and content_analyser_hash == 'b14e7b92' + + +def processors_pre_check() -> bool: + for processor_module in get_processors_modules(state_manager.get_item('processors')): + if not processor_module.pre_check(): + return False + return True + + +def force_download() -> ErrorCode: + common_modules =\ + [ + content_analyser, + face_classifier, + face_detector, + face_landmarker, + face_masker, + face_recognizer, + voice_extractor + ] + available_processors = [ get_file_name(file_path) for file_path in resolve_file_paths('facefusion/processors/modules') ] + processor_modules = get_processors_modules(available_processors) + + for module in common_modules + processor_modules: + if hasattr(module, 'create_static_model_set'): + for model in module.create_static_model_set(state_manager.get_item('download_scope')).values(): + model_hash_set = model.get('hashes') + model_source_set = model.get('sources') + + if model_hash_set and model_source_set: + if not conditional_download_hashes(model_hash_set) or not conditional_download_sources(model_source_set): + return 1 + + return 0 + + +def route_job_manager(args : Args) -> ErrorCode: + if state_manager.get_item('command') == 'job-list': + job_headers, job_contents = compose_job_list(state_manager.get_item('job_status')) + + if job_contents: + cli_helper.render_table(job_headers, job_contents) + return 0 + return 1 + + if state_manager.get_item('command') == 'job-create': + if job_manager.create_job(state_manager.get_item('job_id')): + logger.info(translator.get('job_created').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.error(translator.get('job_not_created').format(job_id = state_manager.get_item('job_id')), __name__) + return 1 + + if state_manager.get_item('command') == 'job-submit': + if job_manager.submit_job(state_manager.get_item('job_id')): + logger.info(translator.get('job_submitted').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.error(translator.get('job_not_submitted').format(job_id = state_manager.get_item('job_id')), __name__) + return 1 + + if state_manager.get_item('command') == 'job-submit-all': + if job_manager.submit_jobs(state_manager.get_item('halt_on_error')): + logger.info(translator.get('job_all_submitted'), __name__) + return 0 + logger.error(translator.get('job_all_not_submitted'), __name__) + return 1 + + if state_manager.get_item('command') == 'job-delete': + if job_manager.delete_job(state_manager.get_item('job_id')): + logger.info(translator.get('job_deleted').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.error(translator.get('job_not_deleted').format(job_id = state_manager.get_item('job_id')), __name__) + return 1 + + if state_manager.get_item('command') == 'job-delete-all': + if job_manager.delete_jobs(state_manager.get_item('halt_on_error')): + logger.info(translator.get('job_all_deleted'), __name__) + return 0 + logger.error(translator.get('job_all_not_deleted'), __name__) + return 1 + + if state_manager.get_item('command') == 'job-add-step': + step_args = reduce_step_args(args) + + if job_manager.add_step(state_manager.get_item('job_id'), step_args): + logger.info(translator.get('job_step_added').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.error(translator.get('job_step_not_added').format(job_id = state_manager.get_item('job_id')), __name__) + return 1 + + if state_manager.get_item('command') == 'job-remix-step': + step_args = reduce_step_args(args) + + if job_manager.remix_step(state_manager.get_item('job_id'), state_manager.get_item('step_index'), step_args): + logger.info(translator.get('job_remix_step_added').format(job_id = state_manager.get_item('job_id'), step_index = state_manager.get_item('step_index')), __name__) + return 0 + logger.error(translator.get('job_remix_step_not_added').format(job_id = state_manager.get_item('job_id'), step_index = state_manager.get_item('step_index')), __name__) + return 1 + + if state_manager.get_item('command') == 'job-insert-step': + step_args = reduce_step_args(args) + + if job_manager.insert_step(state_manager.get_item('job_id'), state_manager.get_item('step_index'), step_args): + logger.info(translator.get('job_step_inserted').format(job_id = state_manager.get_item('job_id'), step_index = state_manager.get_item('step_index')), __name__) + return 0 + logger.error(translator.get('job_step_not_inserted').format(job_id = state_manager.get_item('job_id'), step_index = state_manager.get_item('step_index')), __name__) + return 1 + + if state_manager.get_item('command') == 'job-remove-step': + if job_manager.remove_step(state_manager.get_item('job_id'), state_manager.get_item('step_index')): + logger.info(translator.get('job_step_removed').format(job_id = state_manager.get_item('job_id'), step_index = state_manager.get_item('step_index')), __name__) + return 0 + logger.error(translator.get('job_step_not_removed').format(job_id = state_manager.get_item('job_id'), step_index = state_manager.get_item('step_index')), __name__) + return 1 + return 1 + + +def route_job_runner() -> ErrorCode: + if state_manager.get_item('command') == 'job-run': + logger.info(translator.get('running_job').format(job_id = state_manager.get_item('job_id')), __name__) + if job_runner.run_job(state_manager.get_item('job_id'), process_step): + logger.info(translator.get('processing_job_succeeded').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.info(translator.get('processing_job_failed').format(job_id = state_manager.get_item('job_id')), __name__) + return 1 + + if state_manager.get_item('command') == 'job-run-all': + logger.info(translator.get('running_jobs'), __name__) + if job_runner.run_jobs(process_step, state_manager.get_item('halt_on_error')): + logger.info(translator.get('processing_jobs_succeeded'), __name__) + return 0 + logger.info(translator.get('processing_jobs_failed'), __name__) + return 1 + + if state_manager.get_item('command') == 'job-retry': + logger.info(translator.get('retrying_job').format(job_id = state_manager.get_item('job_id')), __name__) + if job_runner.retry_job(state_manager.get_item('job_id'), process_step): + logger.info(translator.get('processing_job_succeeded').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.info(translator.get('processing_job_failed').format(job_id = state_manager.get_item('job_id')), __name__) + return 1 + + if state_manager.get_item('command') == 'job-retry-all': + logger.info(translator.get('retrying_jobs'), __name__) + if job_runner.retry_jobs(process_step, state_manager.get_item('halt_on_error')): + logger.info(translator.get('processing_jobs_succeeded'), __name__) + return 0 + logger.info(translator.get('processing_jobs_failed'), __name__) + return 1 + return 2 + + +def process_headless(args : Args) -> ErrorCode: + job_id = job_helper.suggest_job_id('headless') + step_args = reduce_step_args(args) + + if job_manager.create_job(job_id) and job_manager.add_step(job_id, step_args) and job_manager.submit_job(job_id) and job_runner.run_job(job_id, process_step): + return 0 + return 1 + + +def process_batch(args : Args) -> ErrorCode: + job_id = job_helper.suggest_job_id('batch') + step_args = reduce_step_args(args) + job_args = reduce_job_args(args) + source_paths = resolve_file_pattern(job_args.get('source_pattern')) + target_paths = resolve_file_pattern(job_args.get('target_pattern')) + + if job_manager.create_job(job_id): + if source_paths and target_paths: + for index, (source_path, target_path) in enumerate(itertools.product(source_paths, target_paths)): + step_args['source_paths'] = [ source_path ] + step_args['target_path'] = target_path + + try: + step_args['output_path'] = job_args.get('output_pattern').format(index = index, source_name = get_file_name(source_path), target_name = get_file_name(target_path), target_extension = get_file_extension(target_path)) + except KeyError: + return 1 + + if not job_manager.add_step(job_id, step_args): + return 1 + if job_manager.submit_job(job_id) and job_runner.run_job(job_id, process_step): + return 0 + + if not source_paths and target_paths: + for index, target_path in enumerate(target_paths): + step_args['target_path'] = target_path + + try: + step_args['output_path'] = job_args.get('output_pattern').format(index = index, target_name = get_file_name(target_path), target_extension = get_file_extension(target_path)) + except KeyError: + return 1 + + if not job_manager.add_step(job_id, step_args): + return 1 + if job_manager.submit_job(job_id) and job_runner.run_job(job_id, process_step): + return 0 + return 1 + + +def process_step(job_id : str, step_index : int, step_args : Args) -> bool: + step_total = job_manager.count_step_total(job_id) + step_args.update(collect_job_args()) + apply_args(step_args, state_manager.set_item) + + logger.info(translator.get('processing_step').format(step_current = step_index + 1, step_total = step_total), __name__) + if common_pre_check() and processors_pre_check(): + error_code = conditional_process() + return error_code == 0 + return False + + +def conditional_process() -> ErrorCode: + start_time = time() + + for processor_module in get_processors_modules(state_manager.get_item('processors')): + if not processor_module.pre_process('output'): + return 2 + + if is_image(state_manager.get_item('target_path')): + return image_to_image.process(start_time) + if is_video(state_manager.get_item('target_path')): + return image_to_video.process(start_time) + + return 0 + + diff --git a/facefusion/curl_builder.py b/facefusion/curl_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..3aa9643702cde0777dccb6fbdd31d42693ada6cd --- /dev/null +++ b/facefusion/curl_builder.py @@ -0,0 +1,32 @@ +import itertools +import shutil +from typing import List + +from facefusion import metadata +from facefusion.types import Command + + +def run(commands : List[Command]) -> List[Command]: + user_agent = metadata.get('name') + '/' + metadata.get('version') + + return [ shutil.which('curl'), '--user-agent', user_agent, '--insecure', '--location', '--silent' ] + commands + + +def chain(*commands : List[Command]) -> List[Command]: + return list(itertools.chain(*commands)) + + +def ping(url : str) -> List[Command]: + return [ '-I', url ] + + +def download(url : str, download_file_path : str) -> List[Command]: + return [ '--create-dirs', '--continue-at', '-', '--output', download_file_path, url ] + + +def set_timeout(timeout : int) -> List[Command]: + return [ '--connect-timeout', str(timeout) ] + + +def set_retry(retry : int) -> List[Command]: + return [ '--retry', str(retry) ] diff --git a/facefusion/download.py b/facefusion/download.py new file mode 100644 index 0000000000000000000000000000000000000000..abd3fc07f1a6083655550ed7d37bccdbb11e123d --- /dev/null +++ b/facefusion/download.py @@ -0,0 +1,175 @@ +import os +import subprocess +from functools import lru_cache +from typing import List, Optional, Tuple +from urllib.parse import urlparse + +from tqdm import tqdm + +import facefusion.choices +from facefusion import curl_builder, logger, process_manager, state_manager, translator +from facefusion.filesystem import get_file_name, get_file_size, is_file, remove_file +from facefusion.hash_helper import validate_hash +from facefusion.types import Command, DownloadProvider, DownloadSet + + +def open_curl(commands : List[Command]) -> subprocess.Popen[bytes]: + commands = curl_builder.run(commands) + return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.PIPE) + + +def conditional_download(download_directory_path : str, urls : List[str]) -> None: + for url in urls: + download_file_name = os.path.basename(urlparse(url).path) + download_file_path = os.path.join(download_directory_path, download_file_name) + initial_size = get_file_size(download_file_path) + download_size = get_static_download_size(url) + + if initial_size < download_size: + with tqdm(total = download_size, initial = initial_size, desc = translator.get('downloading'), unit = 'B', unit_scale = True, unit_divisor = 1024, ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: + commands = curl_builder.chain( + curl_builder.download(url, download_file_path), + curl_builder.set_timeout(5), + curl_builder.set_retry(5) + ) + open_curl(commands) + current_size = initial_size + progress.set_postfix(download_providers = state_manager.get_item('download_providers'), file_name = download_file_name) + + while current_size < download_size: + if is_file(download_file_path): + current_size = get_file_size(download_file_path) + progress.update(current_size - progress.n) + + +@lru_cache(maxsize = 64) +def get_static_download_size(url : str) -> int: + commands = curl_builder.chain( + curl_builder.ping(url), + curl_builder.set_timeout(5) + ) + process = open_curl(commands) + lines = reversed(process.stdout.readlines()) + + for line in lines: + __line__ = line.decode().lower() + if 'content-length:' in __line__: + _, content_length = __line__.split('content-length:') + return int(content_length) + + return 0 + + +@lru_cache(maxsize = 64) +def ping_static_url(url : str) -> bool: + commands = curl_builder.chain( + curl_builder.ping(url), + curl_builder.set_timeout(5) + ) + process = open_curl(commands) + process.communicate() + return process.returncode == 0 + + +def conditional_download_hashes(hash_set : DownloadSet) -> bool: + hash_paths = [ hash_set.get(hash_key).get('path') for hash_key in hash_set.keys() ] + + process_manager.check() + _, invalid_hash_paths = validate_hash_paths(hash_paths) + if invalid_hash_paths: + for index in hash_set: + if hash_set.get(index).get('path') in invalid_hash_paths: + invalid_hash_url = hash_set.get(index).get('url') + if invalid_hash_url: + download_directory_path = os.path.dirname(hash_set.get(index).get('path')) + conditional_download(download_directory_path, [ invalid_hash_url ]) + + valid_hash_paths, invalid_hash_paths = validate_hash_paths(hash_paths) + + for valid_hash_path in valid_hash_paths: + valid_hash_file_name = get_file_name(valid_hash_path) + logger.debug(translator.get('validating_hash_succeeded').format(hash_file_name = valid_hash_file_name), __name__) + for invalid_hash_path in invalid_hash_paths: + invalid_hash_file_name = get_file_name(invalid_hash_path) + logger.error(translator.get('validating_hash_failed').format(hash_file_name = invalid_hash_file_name), __name__) + + if not invalid_hash_paths: + process_manager.end() + return not invalid_hash_paths + + +def conditional_download_sources(source_set : DownloadSet) -> bool: + source_paths = [ source_set.get(source_key).get('path') for source_key in source_set.keys() ] + + process_manager.check() + _, invalid_source_paths = validate_source_paths(source_paths) + if invalid_source_paths: + for index in source_set: + if source_set.get(index).get('path') in invalid_source_paths: + invalid_source_url = source_set.get(index).get('url') + if invalid_source_url: + download_directory_path = os.path.dirname(source_set.get(index).get('path')) + conditional_download(download_directory_path, [ invalid_source_url ]) + + valid_source_paths, invalid_source_paths = validate_source_paths(source_paths) + + for valid_source_path in valid_source_paths: + valid_source_file_name = get_file_name(valid_source_path) + logger.debug(translator.get('validating_source_succeeded').format(source_file_name = valid_source_file_name), __name__) + for invalid_source_path in invalid_source_paths: + invalid_source_file_name = get_file_name(invalid_source_path) + logger.error(translator.get('validating_source_failed').format(source_file_name = invalid_source_file_name), __name__) + + if remove_file(invalid_source_path): + logger.error(translator.get('deleting_corrupt_source').format(source_file_name = invalid_source_file_name), __name__) + + if not invalid_source_paths: + process_manager.end() + return not invalid_source_paths + + +def validate_hash_paths(hash_paths : List[str]) -> Tuple[List[str], List[str]]: + valid_hash_paths = [] + invalid_hash_paths = [] + + for hash_path in hash_paths: + if is_file(hash_path): + valid_hash_paths.append(hash_path) + else: + invalid_hash_paths.append(hash_path) + + return valid_hash_paths, invalid_hash_paths + + +def validate_source_paths(source_paths : List[str]) -> Tuple[List[str], List[str]]: + valid_source_paths = [] + invalid_source_paths = [] + + for source_path in source_paths: + if validate_hash(source_path): + valid_source_paths.append(source_path) + else: + invalid_source_paths.append(source_path) + + return valid_source_paths, invalid_source_paths + + +def resolve_download_url(base_name : str, file_name : str) -> Optional[str]: + download_providers = state_manager.get_item('download_providers') + + for download_provider in download_providers: + download_url = resolve_download_url_by_provider(download_provider, base_name, file_name) + if download_url: + return download_url + + return None + + +def resolve_download_url_by_provider(download_provider : DownloadProvider, base_name : str, file_name : str) -> Optional[str]: + download_provider_value = facefusion.choices.download_provider_set.get(download_provider) + + for download_provider_url in download_provider_value.get('urls'): + if ping_static_url(download_provider_url): + return download_provider_url + download_provider_value.get('path').format(base_name = base_name, file_name = file_name) + + return None diff --git a/facefusion/execution.py b/facefusion/execution.py new file mode 100644 index 0000000000000000000000000000000000000000..71e3981b566187ce109be81e869e45ca576604f3 --- /dev/null +++ b/facefusion/execution.py @@ -0,0 +1,160 @@ +import shutil +import subprocess +import xml.etree.ElementTree as ElementTree +from functools import lru_cache +from typing import List, Optional + +from onnxruntime import get_available_providers, set_default_logger_severity + +import facefusion.choices +from facefusion.types import ExecutionDevice, ExecutionProvider, InferenceSessionProvider, ValueAndUnit + +set_default_logger_severity(3) + + +def has_execution_provider(execution_provider : ExecutionProvider) -> bool: + return execution_provider in get_available_execution_providers() + + +def get_available_execution_providers() -> List[ExecutionProvider]: + inference_session_providers = get_available_providers() + available_execution_providers : List[ExecutionProvider] = [] + + for execution_provider, execution_provider_value in facefusion.choices.execution_provider_set.items(): + if execution_provider_value in inference_session_providers: + index = facefusion.choices.execution_providers.index(execution_provider) + available_execution_providers.insert(index, execution_provider) + + return available_execution_providers + + +def create_inference_session_providers(execution_device_id : int, execution_providers : List[ExecutionProvider]) -> List[InferenceSessionProvider]: + inference_session_providers : List[InferenceSessionProvider] = [] + + for execution_provider in execution_providers: + if execution_provider == 'cuda': + inference_session_providers.append((facefusion.choices.execution_provider_set.get(execution_provider), + { + 'device_id': execution_device_id, + 'cudnn_conv_algo_search': resolve_cudnn_conv_algo_search() + })) + if execution_provider == 'tensorrt': + inference_session_providers.append((facefusion.choices.execution_provider_set.get(execution_provider), + { + 'device_id': execution_device_id, + 'trt_engine_cache_enable': True, + 'trt_engine_cache_path': '.caches', + 'trt_timing_cache_enable': True, + 'trt_timing_cache_path': '.caches', + 'trt_builder_optimization_level': 5 + })) + if execution_provider in [ 'directml', 'rocm' ]: + inference_session_providers.append((facefusion.choices.execution_provider_set.get(execution_provider), + { + 'device_id': execution_device_id + })) + if execution_provider == 'migraphx': + inference_session_providers.append((facefusion.choices.execution_provider_set.get(execution_provider), + { + 'device_id': execution_device_id, + 'migraphx_model_cache_dir': '.caches' + })) + if execution_provider == 'openvino': + inference_session_providers.append((facefusion.choices.execution_provider_set.get(execution_provider), + { + 'device_type': resolve_openvino_device_type(execution_device_id), + 'precision': 'FP32' + })) + if execution_provider == 'coreml': + inference_session_providers.append((facefusion.choices.execution_provider_set.get(execution_provider), + { + 'SpecializationStrategy': 'FastPrediction', + 'ModelCacheDirectory': '.caches' + })) + + if 'cpu' in execution_providers: + inference_session_providers.append(facefusion.choices.execution_provider_set.get('cpu')) + + return inference_session_providers + + +def resolve_cudnn_conv_algo_search() -> str: + execution_devices = detect_static_execution_devices() + product_names = ('GeForce GTX 1630', 'GeForce GTX 1650', 'GeForce GTX 1660') + + for execution_device in execution_devices: + if execution_device.get('product').get('name').startswith(product_names): + return 'DEFAULT' + + return 'EXHAUSTIVE' + + +def resolve_openvino_device_type(execution_device_id : int) -> str: + if execution_device_id == 0: + return 'GPU' + return 'GPU.' + str(execution_device_id) + + +def run_nvidia_smi() -> subprocess.Popen[bytes]: + commands = [ shutil.which('nvidia-smi'), '--query', '--xml-format' ] + return subprocess.Popen(commands, stdout = subprocess.PIPE) + + +@lru_cache() +def detect_static_execution_devices() -> List[ExecutionDevice]: + return detect_execution_devices() + + +def detect_execution_devices() -> List[ExecutionDevice]: + execution_devices : List[ExecutionDevice] = [] + + try: + output, _ = run_nvidia_smi().communicate() + root_element = ElementTree.fromstring(output) + except Exception: + root_element = ElementTree.Element('xml') + + for gpu_element in root_element.findall('gpu'): + execution_devices.append( + { + 'driver_version': root_element.findtext('driver_version'), + 'framework': + { + 'name': 'CUDA', + 'version': root_element.findtext('cuda_version') + }, + 'product': + { + 'vendor': 'NVIDIA', + 'name': gpu_element.findtext('product_name').replace('NVIDIA', '').strip() + }, + 'video_memory': + { + 'total': create_value_and_unit(gpu_element.findtext('fb_memory_usage/total')), + 'free': create_value_and_unit(gpu_element.findtext('fb_memory_usage/free')) + }, + 'temperature': + { + 'gpu': create_value_and_unit(gpu_element.findtext('temperature/gpu_temp')), + 'memory': create_value_and_unit(gpu_element.findtext('temperature/memory_temp')) + }, + 'utilization': + { + 'gpu': create_value_and_unit(gpu_element.findtext('utilization/gpu_util')), + 'memory': create_value_and_unit(gpu_element.findtext('utilization/memory_util')) + } + }) + + return execution_devices + + +def create_value_and_unit(text : str) -> Optional[ValueAndUnit]: + if ' ' in text: + value, unit = text.split() + + return\ + { + 'value': int(value), + 'unit': str(unit) + } + return None diff --git a/facefusion/exit_helper.py b/facefusion/exit_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..19033de3992e211a647ca568ec329201f380fdaa --- /dev/null +++ b/facefusion/exit_helper.py @@ -0,0 +1,34 @@ +import os +import signal +import sys +from time import sleep +from types import FrameType + +from facefusion import process_manager, state_manager +from facefusion.temp_helper import clear_temp_directory +from facefusion.types import ErrorCode + + +def fatal_exit(error_code : ErrorCode) -> None: + os._exit(error_code) + + +def hard_exit(error_code : ErrorCode) -> None: + sys.exit(error_code) + + +def signal_exit(signum : int, frame : FrameType) -> None: + graceful_exit(0) + + +def graceful_exit(error_code : ErrorCode) -> None: + signal.signal(signal.SIGINT, signal.SIG_IGN) + process_manager.stop() + + while process_manager.is_processing(): + sleep(0.5) + + if state_manager.get_item('target_path'): + clear_temp_directory(state_manager.get_item('target_path')) + + hard_exit(error_code) diff --git a/facefusion/face_analyser.py b/facefusion/face_analyser.py new file mode 100644 index 0000000000000000000000000000000000000000..c49c8a4c2a2c2728cf76974286eeb55b72239b61 --- /dev/null +++ b/facefusion/face_analyser.py @@ -0,0 +1,143 @@ +from typing import List, Optional + +import numpy + +from facefusion import state_manager +from facefusion.common_helper import get_first +from facefusion.face_classifier import classify_face +from facefusion.face_detector import detect_faces, detect_faces_by_angle +from facefusion.face_helper import apply_nms, convert_to_face_landmark_5, estimate_face_angle, get_nms_threshold +from facefusion.face_landmarker import detect_face_landmark, estimate_face_landmark_68_5 +from facefusion.face_recognizer import calculate_face_embedding +from facefusion.face_store import get_static_faces, set_static_faces +from facefusion.types import BoundingBox, Face, FaceLandmark5, FaceLandmarkSet, FaceScoreSet, Score, VisionFrame + + +def create_faces(vision_frame : VisionFrame, bounding_boxes : List[BoundingBox], face_scores : List[Score], face_landmarks_5 : List[FaceLandmark5]) -> List[Face]: + faces = [] + nms_threshold = get_nms_threshold(state_manager.get_item('face_detector_model'), state_manager.get_item('face_detector_angles')) + keep_indices = apply_nms(bounding_boxes, face_scores, state_manager.get_item('face_detector_score'), nms_threshold) + + for index in keep_indices: + bounding_box = bounding_boxes[index] + face_score = face_scores[index] + face_landmark_5 = face_landmarks_5[index] + face_landmark_5_68 = face_landmark_5 + face_landmark_68_5 = estimate_face_landmark_68_5(face_landmark_5_68) + face_landmark_68 = face_landmark_68_5 + face_landmark_score_68 = 0.0 + face_angle = estimate_face_angle(face_landmark_68_5) + + if state_manager.get_item('face_landmarker_score') > 0: + face_landmark_68, face_landmark_score_68 = detect_face_landmark(vision_frame, bounding_box, face_angle) + if face_landmark_score_68 > state_manager.get_item('face_landmarker_score'): + face_landmark_5_68 = convert_to_face_landmark_5(face_landmark_68) + + face_landmark_set : FaceLandmarkSet =\ + { + '5': face_landmark_5, + '5/68': face_landmark_5_68, + '68': face_landmark_68, + '68/5': face_landmark_68_5 + } + face_score_set : FaceScoreSet =\ + { + 'detector': face_score, + 'landmarker': face_landmark_score_68 + } + face_embedding, face_embedding_norm = calculate_face_embedding(vision_frame, face_landmark_set.get('5/68')) + gender, age, race = classify_face(vision_frame, face_landmark_set.get('5/68')) + faces.append(Face( + bounding_box = bounding_box, + score_set = face_score_set, + landmark_set = face_landmark_set, + angle = face_angle, + embedding = face_embedding, + embedding_norm = face_embedding_norm, + gender = gender, + age = age, + race = race + )) + return faces + + +def get_one_face(faces : List[Face], position : int = 0) -> Optional[Face]: + if faces: + position = min(position, len(faces) - 1) + return faces[position] + return None + + +def get_average_face(faces : List[Face]) -> Optional[Face]: + face_embeddings = [] + face_embeddings_norm = [] + + if faces: + first_face = get_first(faces) + + for face in faces: + face_embeddings.append(face.embedding) + face_embeddings_norm.append(face.embedding_norm) + + return Face( + bounding_box = first_face.bounding_box, + score_set = first_face.score_set, + landmark_set = first_face.landmark_set, + angle = first_face.angle, + embedding = numpy.mean(face_embeddings, axis = 0), + embedding_norm = numpy.mean(face_embeddings_norm, axis = 0), + gender = first_face.gender, + age = first_face.age, + race = first_face.race + ) + return None + + +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): + static_faces = get_static_faces(vision_frame) + if static_faces: + many_faces.extend(static_faces) + else: + all_bounding_boxes = [] + all_face_scores = [] + all_face_landmarks_5 = [] + + for face_detector_angle in state_manager.get_item('face_detector_angles'): + if face_detector_angle == 0: + bounding_boxes, face_scores, face_landmarks_5 = detect_faces(vision_frame) + else: + bounding_boxes, face_scores, face_landmarks_5 = detect_faces_by_angle(vision_frame, face_detector_angle) + all_bounding_boxes.extend(bounding_boxes) + all_face_scores.extend(face_scores) + all_face_landmarks_5.extend(face_landmarks_5) + + if all_bounding_boxes and all_face_scores and all_face_landmarks_5 and state_manager.get_item('face_detector_score') > 0: + faces = create_faces(vision_frame, all_bounding_boxes, all_face_scores, all_face_landmarks_5) + + if faces: + many_faces.extend(faces) + set_static_faces(vision_frame, faces) + return many_faces + + +def scale_face(target_face : Face, target_vision_frame : VisionFrame, temp_vision_frame : VisionFrame) -> Face: + scale_x = temp_vision_frame.shape[1] / target_vision_frame.shape[1] + scale_y = temp_vision_frame.shape[0] / target_vision_frame.shape[0] + + bounding_box = target_face.bounding_box * [ scale_x, scale_y, scale_x, scale_y ] + landmark_set =\ + { + '5': target_face.landmark_set.get('5') * numpy.array([ scale_x, scale_y ]), + '5/68': target_face.landmark_set.get('5/68') * numpy.array([ scale_x, scale_y ]), + '68': target_face.landmark_set.get('68') * numpy.array([ scale_x, scale_y ]), + '68/5': target_face.landmark_set.get('68/5') * numpy.array([ scale_x, scale_y ]) + } + + return target_face._replace( + bounding_box = bounding_box, + landmark_set = landmark_set + ) diff --git a/facefusion/face_classifier.py b/facefusion/face_classifier.py new file mode 100644 index 0000000000000000000000000000000000000000..ccac76c8e8da97ef2ee78d543f697d7d83d86478 --- /dev/null +++ b/facefusion/face_classifier.py @@ -0,0 +1,140 @@ +from functools import lru_cache +from typing import List, Tuple + +import numpy + +from facefusion import inference_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.face_helper import warp_face_by_face_landmark_5 +from facefusion.filesystem import resolve_relative_path +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.types import Age, DownloadScope, FaceLandmark5, Gender, InferencePool, ModelOptions, ModelSet, Race, VisionFrame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'fairface': + { + '__metadata__': + { + 'vendor': 'dchen236', + 'license': 'Non-Commercial', + 'year': 2021 + }, + 'hashes': + { + 'face_classifier': + { + 'url': resolve_download_url('models-3.0.0', 'fairface.hash'), + 'path': resolve_relative_path('../.assets/models/fairface.hash') + } + }, + 'sources': + { + 'face_classifier': + { + 'url': resolve_download_url('models-3.0.0', 'fairface.onnx'), + 'path': resolve_relative_path('../.assets/models/fairface.onnx') + } + }, + 'template': 'arcface_112_v2', + 'size': (224, 224), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ 'fairface' ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ 'fairface' ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + return create_static_model_set('full').get('fairface') + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def classify_face(temp_vision_frame : VisionFrame, face_landmark_5 : FaceLandmark5) -> Tuple[Gender, Age, Race]: + model_template = get_model_options().get('template') + model_size = get_model_options().get('size') + model_mean = get_model_options().get('mean') + model_standard_deviation = get_model_options().get('standard_deviation') + crop_vision_frame, _ = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5, model_template, model_size) + crop_vision_frame = crop_vision_frame.astype(numpy.float32)[:, :, ::-1] / 255.0 + crop_vision_frame -= model_mean + crop_vision_frame /= model_standard_deviation + crop_vision_frame = crop_vision_frame.transpose(2, 0, 1) + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0) + gender_id, age_id, race_id = forward(crop_vision_frame) + gender = categorize_gender(gender_id[0]) + age = categorize_age(age_id[0]) + race = categorize_race(race_id[0]) + return gender, age, race + + +def forward(crop_vision_frame : VisionFrame) -> Tuple[List[int], List[int], List[int]]: + face_classifier = get_inference_pool().get('face_classifier') + + with conditional_thread_semaphore(): + race_id, gender_id, age_id = face_classifier.run(None, + { + 'input': crop_vision_frame + }) + + return gender_id, age_id, race_id + + +def categorize_gender(gender_id : int) -> Gender: + if gender_id == 1: + return 'female' + return 'male' + + +def categorize_age(age_id : int) -> Age: + if age_id == 0: + return range(0, 2) + if age_id == 1: + return range(3, 9) + if age_id == 2: + return range(10, 19) + if age_id == 3: + return range(20, 29) + if age_id == 4: + return range(30, 39) + if age_id == 5: + return range(40, 49) + if age_id == 6: + return range(50, 59) + if age_id == 7: + return range(60, 69) + return range(70, 100) + + +def categorize_race(race_id : int) -> Race: + if race_id == 1: + return 'black' + if race_id == 2: + return 'latino' + if race_id == 3 or race_id == 4: + return 'asian' + if race_id == 5: + return 'indian' + if race_id == 6: + return 'arabic' + return 'white' diff --git a/facefusion/face_detector.py b/facefusion/face_detector.py new file mode 100644 index 0000000000000000000000000000000000000000..6ff3f8f3d643590e65382726f33966dfa5374b8b --- /dev/null +++ b/facefusion/face_detector.py @@ -0,0 +1,458 @@ +from functools import lru_cache +from typing import List, Sequence, Tuple + +import cv2 +import numpy + +from facefusion import inference_manager, state_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.face_helper import create_rotation_matrix_and_size, create_static_anchors, distance_to_bounding_box, distance_to_face_landmark_5, normalize_bounding_box, transform_bounding_box, transform_points +from facefusion.filesystem import resolve_relative_path +from facefusion.thread_helper import thread_semaphore +from facefusion.types import Angle, BoundingBox, Detection, DownloadScope, DownloadSet, FaceLandmark5, InferencePool, Margin, ModelSet, Score, VisionFrame +from facefusion.vision import restrict_frame, unpack_resolution + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'retinaface': + { + '__metadata__': + { + 'vendor': 'InsightFace', + 'license': 'Non-Commercial', + 'year': 2020 + }, + 'hashes': + { + 'retinaface': + { + 'url': resolve_download_url('models-3.0.0', 'retinaface_10g.hash'), + 'path': resolve_relative_path('../.assets/models/retinaface_10g.hash') + } + }, + 'sources': + { + 'retinaface': + { + 'url': resolve_download_url('models-3.0.0', 'retinaface_10g.onnx'), + 'path': resolve_relative_path('../.assets/models/retinaface_10g.onnx') + } + } + }, + 'scrfd': + { + '__metadata__': + { + 'vendor': 'InsightFace', + 'license': 'Non-Commercial', + 'year': 2021 + }, + 'hashes': + { + 'scrfd': + { + 'url': resolve_download_url('models-3.0.0', 'scrfd_2.5g.hash'), + 'path': resolve_relative_path('../.assets/models/scrfd_2.5g.hash') + } + }, + 'sources': + { + 'scrfd': + { + 'url': resolve_download_url('models-3.0.0', 'scrfd_2.5g.onnx'), + 'path': resolve_relative_path('../.assets/models/scrfd_2.5g.onnx') + } + } + }, + 'yolo_face': + { + '__metadata__': + { + 'vendor': 'derronqi', + 'license': 'GPL-3.0', + 'year': 2022 + }, + 'hashes': + { + 'yolo_face': + { + 'url': resolve_download_url('models-3.0.0', 'yoloface_8n.hash'), + 'path': resolve_relative_path('../.assets/models/yoloface_8n.hash') + } + }, + 'sources': + { + 'yolo_face': + { + 'url': resolve_download_url('models-3.0.0', 'yoloface_8n.onnx'), + 'path': resolve_relative_path('../.assets/models/yoloface_8n.onnx') + } + } + }, + 'yunet': + { + '__metadata__': + { + 'vendor': 'OpenCV', + 'license': 'MIT', + 'year': 2023 + }, + 'hashes': + { + 'yunet': + { + 'url': resolve_download_url('models-3.4.0', 'yunet_2023_mar.hash'), + 'path': resolve_relative_path('../.assets/models/yunet_2023_mar.hash') + } + }, + 'sources': + { + 'yunet': + { + 'url': resolve_download_url('models-3.4.0', 'yunet_2023_mar.onnx'), + 'path': resolve_relative_path('../.assets/models/yunet_2023_mar.onnx') + } + } + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('face_detector_model') ] + _, model_source_set = collect_model_downloads() + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('face_detector_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_set = create_static_model_set('full') + model_hash_set = {} + model_source_set = {} + + for face_detector_model in [ 'retinaface', 'scrfd', 'yolo_face', 'yunet' ]: + if state_manager.get_item('face_detector_model') in [ 'many', face_detector_model ]: + model_hash_set[face_detector_model] = model_set.get(face_detector_model).get('hashes').get(face_detector_model) + model_source_set[face_detector_model] = model_set.get(face_detector_model).get('sources').get(face_detector_model) + + return model_hash_set, model_source_set + + +def pre_check() -> bool: + model_hash_set, model_source_set = collect_model_downloads() + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def detect_faces(vision_frame : VisionFrame) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + margin_top, margin_right, margin_bottom, margin_left = prepare_margin(vision_frame) + margin_vision_frame = numpy.pad(vision_frame, ((margin_top, margin_bottom), (margin_left, margin_right), (0, 0))) + all_bounding_boxes : List[BoundingBox] = [] + all_face_scores : List[Score] = [] + all_face_landmarks_5 : List[FaceLandmark5] = [] + + if state_manager.get_item('face_detector_model') in [ 'many', 'retinaface' ]: + bounding_boxes, face_scores, face_landmarks_5 = detect_with_retinaface(margin_vision_frame, state_manager.get_item('face_detector_size')) + all_bounding_boxes.extend(bounding_boxes) + all_face_scores.extend(face_scores) + all_face_landmarks_5.extend(face_landmarks_5) + + if state_manager.get_item('face_detector_model') in [ 'many', 'scrfd' ]: + bounding_boxes, face_scores, face_landmarks_5 = detect_with_scrfd(margin_vision_frame, state_manager.get_item('face_detector_size')) + all_bounding_boxes.extend(bounding_boxes) + all_face_scores.extend(face_scores) + all_face_landmarks_5.extend(face_landmarks_5) + + if state_manager.get_item('face_detector_model') in [ 'many', 'yolo_face' ]: + bounding_boxes, face_scores, face_landmarks_5 = detect_with_yolo_face(margin_vision_frame, state_manager.get_item('face_detector_size')) + all_bounding_boxes.extend(bounding_boxes) + all_face_scores.extend(face_scores) + all_face_landmarks_5.extend(face_landmarks_5) + + if state_manager.get_item('face_detector_model') == 'yunet': + bounding_boxes, face_scores, face_landmarks_5 = detect_with_yunet(margin_vision_frame, state_manager.get_item('face_detector_size')) + all_bounding_boxes.extend(bounding_boxes) + all_face_scores.extend(face_scores) + all_face_landmarks_5.extend(face_landmarks_5) + + all_bounding_boxes = [ normalize_bounding_box(all_bounding_box) - numpy.array([ margin_left, margin_top, margin_left, margin_top ]) for all_bounding_box in all_bounding_boxes ] + all_face_landmarks_5 = [ all_face_landmark_5 - numpy.array([ margin_left, margin_top ]) for all_face_landmark_5 in all_face_landmarks_5 ] + return all_bounding_boxes, all_face_scores, all_face_landmarks_5 + + +def prepare_margin(vision_frame : VisionFrame) -> Margin: + margin_top = int(vision_frame.shape[0] * numpy.interp(state_manager.get_item('face_detector_margin')[0], [ 0, 100 ], [ 0, 0.5 ])) + margin_right = int(vision_frame.shape[1] * numpy.interp(state_manager.get_item('face_detector_margin')[1], [ 0, 100 ], [ 0, 0.5 ])) + margin_bottom = int(vision_frame.shape[0] * numpy.interp(state_manager.get_item('face_detector_margin')[2], [ 0, 100 ], [ 0, 0.5 ])) + margin_left = int(vision_frame.shape[1] * numpy.interp(state_manager.get_item('face_detector_margin')[3], [ 0, 100 ], [ 0, 0.5 ])) + return margin_top, margin_right, margin_bottom, margin_left + + +def detect_faces_by_angle(vision_frame : VisionFrame, face_angle : Angle) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + rotation_matrix, rotation_size = create_rotation_matrix_and_size(face_angle, vision_frame.shape[:2][::-1]) + rotation_vision_frame = cv2.warpAffine(vision_frame, rotation_matrix, rotation_size) + rotation_inverse_matrix = cv2.invertAffineTransform(rotation_matrix) + bounding_boxes, face_scores, face_landmarks_5 = detect_faces(rotation_vision_frame) + bounding_boxes = [ transform_bounding_box(bounding_box, rotation_inverse_matrix) for bounding_box in bounding_boxes ] + face_landmarks_5 = [ transform_points(face_landmark_5, rotation_inverse_matrix) for face_landmark_5 in face_landmarks_5 ] + return bounding_boxes, face_scores, face_landmarks_5 + + +def detect_with_retinaface(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + bounding_boxes = [] + face_scores = [] + face_landmarks_5 = [] + feature_strides = [ 8, 16, 32 ] + feature_map_channel = 3 + anchor_total = 2 + face_detector_score = state_manager.get_item('face_detector_score') + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = restrict_frame(vision_frame, (face_detector_width, face_detector_height)) + ratio_height = vision_frame.shape[0] / temp_vision_frame.shape[0] + ratio_width = vision_frame.shape[1] / temp_vision_frame.shape[1] + detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) + detect_vision_frame = normalize_detect_frame(detect_vision_frame, [ -1, 1 ]) + detection = forward_with_retinaface(detect_vision_frame) + + for index, feature_stride in enumerate(feature_strides): + face_scores_raw = detection[index] + keep_indices = numpy.where(face_scores_raw >= face_detector_score)[0] + + if numpy.any(keep_indices): + stride_height = face_detector_height // feature_stride + stride_width = face_detector_width // feature_stride + anchors = create_static_anchors(feature_stride, anchor_total, stride_height, stride_width) + bounding_boxes_raw = detection[index + feature_map_channel] * feature_stride + face_landmarks_5_raw = detection[index + feature_map_channel * 2] * feature_stride + + for bounding_box_raw in distance_to_bounding_box(anchors, bounding_boxes_raw)[keep_indices]: + bounding_boxes.append(numpy.array( + [ + bounding_box_raw[0] * ratio_width, + bounding_box_raw[1] * ratio_height, + bounding_box_raw[2] * ratio_width, + bounding_box_raw[3] * ratio_height + ])) + + for face_score_raw in face_scores_raw[keep_indices]: + face_scores.append(face_score_raw[0]) + + for face_landmark_raw_5 in distance_to_face_landmark_5(anchors, face_landmarks_5_raw)[keep_indices]: + face_landmarks_5.append(face_landmark_raw_5 * [ ratio_width, ratio_height ]) + + return bounding_boxes, face_scores, face_landmarks_5 + + +def detect_with_scrfd(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + bounding_boxes = [] + face_scores = [] + face_landmarks_5 = [] + feature_strides = [ 8, 16, 32 ] + feature_map_channel = 3 + anchor_total = 2 + face_detector_score = state_manager.get_item('face_detector_score') + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = restrict_frame(vision_frame, (face_detector_width, face_detector_height)) + ratio_height = vision_frame.shape[0] / temp_vision_frame.shape[0] + ratio_width = vision_frame.shape[1] / temp_vision_frame.shape[1] + detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) + detect_vision_frame = normalize_detect_frame(detect_vision_frame, [ -1, 1 ]) + detection = forward_with_scrfd(detect_vision_frame) + + for index, feature_stride in enumerate(feature_strides): + face_scores_raw = detection[index] + keep_indices = numpy.where(face_scores_raw >= face_detector_score)[0] + + if numpy.any(keep_indices): + stride_height = face_detector_height // feature_stride + stride_width = face_detector_width // feature_stride + anchors = create_static_anchors(feature_stride, anchor_total, stride_height, stride_width) + bounding_boxes_raw = detection[index + feature_map_channel] * feature_stride + face_landmarks_5_raw = detection[index + feature_map_channel * 2] * feature_stride + + for bounding_box_raw in distance_to_bounding_box(anchors, bounding_boxes_raw)[keep_indices]: + bounding_boxes.append(numpy.array( + [ + bounding_box_raw[0] * ratio_width, + bounding_box_raw[1] * ratio_height, + bounding_box_raw[2] * ratio_width, + bounding_box_raw[3] * ratio_height + ])) + + for face_score_raw in face_scores_raw[keep_indices]: + face_scores.append(face_score_raw[0]) + + for face_landmark_raw_5 in distance_to_face_landmark_5(anchors, face_landmarks_5_raw)[keep_indices]: + face_landmarks_5.append(face_landmark_raw_5 * [ ratio_width, ratio_height ]) + + return bounding_boxes, face_scores, face_landmarks_5 + + +def detect_with_yolo_face(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + bounding_boxes = [] + face_scores = [] + face_landmarks_5 = [] + face_detector_score = state_manager.get_item('face_detector_score') + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = restrict_frame(vision_frame, (face_detector_width, face_detector_height)) + ratio_height = vision_frame.shape[0] / temp_vision_frame.shape[0] + ratio_width = vision_frame.shape[1] / temp_vision_frame.shape[1] + detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) + detect_vision_frame = normalize_detect_frame(detect_vision_frame, [ 0, 1 ]) + detection = forward_with_yolo_face(detect_vision_frame) + detection = numpy.squeeze(detection).T + bounding_boxes_raw, face_scores_raw, face_landmarks_5_raw = numpy.split(detection, [ 4, 5 ], axis = 1) + keep_indices = numpy.where(face_scores_raw > face_detector_score)[0] + + if numpy.any(keep_indices): + bounding_boxes_raw, face_scores_raw, face_landmarks_5_raw = bounding_boxes_raw[keep_indices], face_scores_raw[keep_indices], face_landmarks_5_raw[keep_indices] + + for bounding_box_raw in bounding_boxes_raw: + bounding_boxes.append(numpy.array( + [ + (bounding_box_raw[0] - bounding_box_raw[2] / 2) * ratio_width, + (bounding_box_raw[1] - bounding_box_raw[3] / 2) * ratio_height, + (bounding_box_raw[0] + bounding_box_raw[2] / 2) * ratio_width, + (bounding_box_raw[1] + bounding_box_raw[3] / 2) * ratio_height + ])) + + face_scores = face_scores_raw.ravel().tolist() + face_landmarks_5_raw[:, 0::3] = (face_landmarks_5_raw[:, 0::3]) * ratio_width + face_landmarks_5_raw[:, 1::3] = (face_landmarks_5_raw[:, 1::3]) * ratio_height + + for face_landmark_raw_5 in face_landmarks_5_raw: + face_landmarks_5.append(numpy.array(face_landmark_raw_5.reshape(-1, 3)[:, :2])) + + return bounding_boxes, face_scores, face_landmarks_5 + + +def detect_with_yunet(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + bounding_boxes = [] + face_scores = [] + face_landmarks_5 = [] + feature_strides = [ 8, 16, 32 ] + feature_map_channel = 3 + anchor_total = 1 + face_detector_score = state_manager.get_item('face_detector_score') + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = restrict_frame(vision_frame, (face_detector_width, face_detector_height)) + ratio_height = vision_frame.shape[0] / temp_vision_frame.shape[0] + ratio_width = vision_frame.shape[1] / temp_vision_frame.shape[1] + detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) + detect_vision_frame = normalize_detect_frame(detect_vision_frame, [ 0, 255 ]) + detection = forward_with_yunet(detect_vision_frame) + + for index, feature_stride in enumerate(feature_strides): + face_scores_raw = (detection[index] * detection[index + feature_map_channel]).reshape(-1) + keep_indices = numpy.where(face_scores_raw >= face_detector_score)[0] + + if numpy.any(keep_indices): + stride_height = face_detector_height // feature_stride + stride_width = face_detector_width // feature_stride + anchors = create_static_anchors(feature_stride, anchor_total, stride_height, stride_width) + bounding_boxes_center = detection[index + feature_map_channel * 2].squeeze(0)[:, :2] * feature_stride + anchors + bounding_boxes_size = numpy.exp(detection[index + feature_map_channel * 2].squeeze(0)[:, 2:4]) * feature_stride + face_landmarks_5_raw = detection[index + feature_map_channel * 3].squeeze(0) + + bounding_boxes_raw = numpy.stack( + [ + bounding_boxes_center[:, 0] - bounding_boxes_size[:, 0] / 2, + bounding_boxes_center[:, 1] - bounding_boxes_size[:, 1] / 2, + bounding_boxes_center[:, 0] + bounding_boxes_size[:, 0] / 2, + bounding_boxes_center[:, 1] + bounding_boxes_size[:, 1] / 2 + ], axis = -1) + + for bounding_box_raw in bounding_boxes_raw[keep_indices]: + bounding_boxes.append(numpy.array( + [ + bounding_box_raw[0] * ratio_width, + bounding_box_raw[1] * ratio_height, + bounding_box_raw[2] * ratio_width, + bounding_box_raw[3] * ratio_height + ])) + + face_scores.extend(face_scores_raw[keep_indices]) + face_landmarks_5_raw = numpy.concatenate( + [ + face_landmarks_5_raw[:, [0, 1]] * feature_stride + anchors, + face_landmarks_5_raw[:, [2, 3]] * feature_stride + anchors, + face_landmarks_5_raw[:, [4, 5]] * feature_stride + anchors, + face_landmarks_5_raw[:, [6, 7]] * feature_stride + anchors, + face_landmarks_5_raw[:, [8, 9]] * feature_stride + anchors + ], axis = -1).reshape(-1, 5, 2) + + for face_landmark_raw_5 in face_landmarks_5_raw[keep_indices]: + face_landmarks_5.append(face_landmark_raw_5 * [ ratio_width, ratio_height ]) + + return bounding_boxes, face_scores, face_landmarks_5 + + +def forward_with_retinaface(detect_vision_frame : VisionFrame) -> Detection: + face_detector = get_inference_pool().get('retinaface') + + with thread_semaphore(): + detection = face_detector.run(None, + { + 'input': detect_vision_frame + }) + + return detection + + +def forward_with_scrfd(detect_vision_frame : VisionFrame) -> Detection: + face_detector = get_inference_pool().get('scrfd') + + with thread_semaphore(): + detection = face_detector.run(None, + { + 'input': detect_vision_frame + }) + + return detection + + +def forward_with_yolo_face(detect_vision_frame : VisionFrame) -> Detection: + face_detector = get_inference_pool().get('yolo_face') + + with thread_semaphore(): + detection = face_detector.run(None, + { + 'input': detect_vision_frame + }) + + return detection + + +def forward_with_yunet(detect_vision_frame : VisionFrame) -> Detection: + face_detector = get_inference_pool().get('yunet') + + with thread_semaphore(): + detection = face_detector.run(None, + { + 'input': detect_vision_frame + }) + + return detection + + +def prepare_detect_frame(temp_vision_frame : VisionFrame, face_detector_size : str) -> VisionFrame: + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + detect_vision_frame = numpy.zeros((face_detector_height, face_detector_width, 3)) + detect_vision_frame[:temp_vision_frame.shape[0], :temp_vision_frame.shape[1], :] = temp_vision_frame + detect_vision_frame = numpy.expand_dims(detect_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) + return detect_vision_frame + + +def normalize_detect_frame(detect_vision_frame : VisionFrame, normalize_range : Sequence[int]) -> VisionFrame: + if normalize_range == [ -1, 1 ]: + return (detect_vision_frame - 127.5) / 128.0 + if normalize_range == [ 0, 1 ]: + return detect_vision_frame / 255.0 + return detect_vision_frame diff --git a/facefusion/face_helper.py b/facefusion/face_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..37feade2f37584142b18abddf6ab61a1772c0863 --- /dev/null +++ b/facefusion/face_helper.py @@ -0,0 +1,256 @@ +from functools import lru_cache +from typing import List, Sequence, Tuple + +import cv2 +import numpy +from cv2.typing import Size + +from facefusion.types import Anchors, Angle, BoundingBox, Distance, FaceDetectorModel, FaceLandmark5, FaceLandmark68, Mask, Matrix, Points, Scale, Score, Translation, VisionFrame, WarpTemplate, WarpTemplateSet + +WARP_TEMPLATE_SET : WarpTemplateSet =\ +{ + 'arcface_112_v1': numpy.array( + [ + [ 0.35473214, 0.45658929 ], + [ 0.64526786, 0.45658929 ], + [ 0.50000000, 0.61154464 ], + [ 0.37913393, 0.77687500 ], + [ 0.62086607, 0.77687500 ] + ]), + 'arcface_112_v2': numpy.array( + [ + [ 0.34191607, 0.46157411 ], + [ 0.65653393, 0.45983393 ], + [ 0.50022500, 0.64050536 ], + [ 0.37097589, 0.82469196 ], + [ 0.63151696, 0.82325089 ] + ]), + 'arcface_128': numpy.array( + [ + [ 0.36167656, 0.40387734 ], + [ 0.63696719, 0.40235469 ], + [ 0.50019687, 0.56044219 ], + [ 0.38710391, 0.72160547 ], + [ 0.61507734, 0.72034453 ] + ]), + 'dfl_whole_face': numpy.array( + [ + [ 0.35342266, 0.39285716 ], + [ 0.62797622, 0.39285716 ], + [ 0.48660713, 0.54017860 ], + [ 0.38839287, 0.68750011 ], + [ 0.59821427, 0.68750011 ] + ]), + 'ffhq_512': numpy.array( + [ + [ 0.37691676, 0.46864664 ], + [ 0.62285697, 0.46912813 ], + [ 0.50123859, 0.61331904 ], + [ 0.39308822, 0.72541100 ], + [ 0.61150205, 0.72490465 ] + ]), + 'mtcnn_512': numpy.array( + [ + [ 0.36562865, 0.46733799 ], + [ 0.63305391, 0.46585885 ], + [ 0.50019127, 0.61942959 ], + [ 0.39032951, 0.77598822 ], + [ 0.61178945, 0.77476328 ] + ]), + 'styleganex_384': numpy.array( + [ + [ 0.42353745, 0.52289879 ], + [ 0.57725008, 0.52319972 ], + [ 0.50123859, 0.61331904 ], + [ 0.43364461, 0.68337652 ], + [ 0.57015325, 0.68306005 ] + ]) +} + + +def estimate_matrix_by_face_landmark_5(face_landmark_5 : FaceLandmark5, warp_template : WarpTemplate, crop_size : Size) -> Matrix: + warp_template_norm = WARP_TEMPLATE_SET.get(warp_template) * crop_size + affine_matrix = cv2.estimateAffinePartial2D(face_landmark_5, warp_template_norm, method = cv2.RANSAC, ransacReprojThreshold = 100)[0] + return affine_matrix + + +def warp_face_by_face_landmark_5(temp_vision_frame : VisionFrame, face_landmark_5 : FaceLandmark5, warp_template : WarpTemplate, crop_size : Size) -> Tuple[VisionFrame, Matrix]: + affine_matrix = estimate_matrix_by_face_landmark_5(face_landmark_5, warp_template, crop_size) + crop_vision_frame = cv2.warpAffine(temp_vision_frame, affine_matrix, crop_size, borderMode = cv2.BORDER_REPLICATE, flags = cv2.INTER_AREA) + return crop_vision_frame, affine_matrix + + +def warp_face_by_bounding_box(temp_vision_frame : VisionFrame, bounding_box : BoundingBox, crop_size : Size) -> Tuple[VisionFrame, Matrix]: + source_points = numpy.array([ [ bounding_box[0], bounding_box[1] ], [bounding_box[2], bounding_box[1] ], [ bounding_box[0], bounding_box[3] ] ]).astype(numpy.float32) + target_points = numpy.array([ [ 0, 0 ], [ crop_size[0], 0 ], [ 0, crop_size[1] ] ]).astype(numpy.float32) + affine_matrix = cv2.getAffineTransform(source_points, target_points) + if bounding_box[2] - bounding_box[0] > crop_size[0] or bounding_box[3] - bounding_box[1] > crop_size[1]: + interpolation_method = cv2.INTER_AREA + else: + interpolation_method = cv2.INTER_LINEAR + crop_vision_frame = cv2.warpAffine(temp_vision_frame, affine_matrix, crop_size, flags = interpolation_method) + return crop_vision_frame, affine_matrix + + +def warp_face_by_translation(temp_vision_frame : VisionFrame, translation : Translation, scale : float, crop_size : Size) -> Tuple[VisionFrame, Matrix]: + affine_matrix = numpy.array([ [ scale, 0, translation[0] ], [ 0, scale, translation[1] ] ]) + crop_vision_frame = cv2.warpAffine(temp_vision_frame, affine_matrix, crop_size) + return crop_vision_frame, affine_matrix + + +def paste_back(temp_vision_frame : VisionFrame, crop_vision_frame : VisionFrame, crop_vision_mask : Mask, affine_matrix : Matrix) -> VisionFrame: + paste_bounding_box, paste_matrix = calculate_paste_area(temp_vision_frame, crop_vision_frame, affine_matrix) + x1, y1, x2, y2 = paste_bounding_box + paste_width = x2 - x1 + paste_height = y2 - y1 + inverse_vision_mask = cv2.warpAffine(crop_vision_mask, paste_matrix, (paste_width, paste_height)).clip(0, 1) + inverse_vision_mask = numpy.expand_dims(inverse_vision_mask, axis = -1) + inverse_vision_frame = cv2.warpAffine(crop_vision_frame, paste_matrix, (paste_width, paste_height), borderMode = cv2.BORDER_REPLICATE) + temp_vision_frame = temp_vision_frame.copy() + paste_vision_frame = temp_vision_frame[y1:y2, x1:x2] + paste_vision_frame = paste_vision_frame * (1 - inverse_vision_mask) + inverse_vision_frame * inverse_vision_mask + temp_vision_frame[y1:y2, x1:x2] = paste_vision_frame.astype(temp_vision_frame.dtype) + return temp_vision_frame + + +def calculate_paste_area(temp_vision_frame : VisionFrame, crop_vision_frame : VisionFrame, affine_matrix : Matrix) -> Tuple[BoundingBox, Matrix]: + temp_height, temp_width = temp_vision_frame.shape[:2] + crop_height, crop_width = crop_vision_frame.shape[:2] + inverse_matrix = cv2.invertAffineTransform(affine_matrix) + crop_points = numpy.array([ [ 0, 0 ], [ crop_width, 0 ], [ crop_width, crop_height ], [ 0, crop_height ] ]) + paste_region_points = transform_points(crop_points, inverse_matrix) + paste_region_point_min = numpy.floor(paste_region_points.min(axis = 0)).astype(int) + paste_region_point_max = numpy.ceil(paste_region_points.max(axis = 0)).astype(int) + x1, y1 = numpy.clip(paste_region_point_min, 0, [ temp_width, temp_height ]) + x2, y2 = numpy.clip(paste_region_point_max, 0, [ temp_width, temp_height ]) + paste_bounding_box = numpy.array([ x1, y1, x2, y2 ]) + paste_matrix = inverse_matrix.copy() + paste_matrix[0, 2] -= x1 + paste_matrix[1, 2] -= y1 + return paste_bounding_box, paste_matrix + + +@lru_cache() +def create_static_anchors(feature_stride : int, anchor_total : int, stride_height : int, stride_width : int) -> Anchors: + x, y = numpy.mgrid[:stride_width, :stride_height] + anchors = numpy.stack((y, x), axis = -1) + anchors = (anchors * feature_stride).reshape((-1, 2)) + anchors = numpy.stack([ anchors ] * anchor_total, axis = 1).reshape((-1, 2)) + return anchors + + +def create_rotation_matrix_and_size(angle : Angle, size : Size) -> Tuple[Matrix, Size]: + rotation_matrix = cv2.getRotationMatrix2D((size[0] / 2, size[1] / 2), angle, 1) + rotation_size = numpy.dot(numpy.abs(rotation_matrix[:, :2]), size) + rotation_matrix[:, -1] += (rotation_size - size) * 0.5 #type:ignore[misc] + rotation_size = int(rotation_size[0]), int(rotation_size[1]) + return rotation_matrix, rotation_size + + +def create_bounding_box(face_landmark_68 : FaceLandmark68) -> BoundingBox: + x1, y1 = numpy.min(face_landmark_68, axis = 0) + x2, y2 = numpy.max(face_landmark_68, axis = 0) + bounding_box = normalize_bounding_box(numpy.array([ x1, y1, x2, y2 ])) + return bounding_box + + +def normalize_bounding_box(bounding_box : BoundingBox) -> BoundingBox: + x1, y1, x2, y2 = bounding_box + x1, x2 = sorted([ x1, x2 ]) + y1, y2 = sorted([ y1, y2 ]) + return numpy.array([ x1, y1, x2, y2 ]) + + +def transform_points(points : Points, matrix : Matrix) -> Points: + points = points.reshape(-1, 1, 2) + points = cv2.transform(points, matrix) #type:ignore[assignment] + points = points.reshape(-1, 2) + return points + + +def transform_bounding_box(bounding_box : BoundingBox, matrix : Matrix) -> BoundingBox: + points = numpy.array( + [ + [ bounding_box[0], bounding_box[1] ], + [ bounding_box[2], bounding_box[1] ], + [ bounding_box[2], bounding_box[3] ], + [ bounding_box[0], bounding_box[3] ] + ]) + points = transform_points(points, matrix) + x1, y1 = numpy.min(points, axis = 0) + x2, y2 = numpy.max(points, axis = 0) + return normalize_bounding_box(numpy.array([ x1, y1, x2, y2 ])) + + +def distance_to_bounding_box(points : Points, distance : Distance) -> BoundingBox: + x1 = points[:, 0] - distance[:, 0] + y1 = points[:, 1] - distance[:, 1] + x2 = points[:, 0] + distance[:, 2] + y2 = points[:, 1] + distance[:, 3] + bounding_box = numpy.column_stack([ x1, y1, x2, y2 ]) + return bounding_box + + +def distance_to_face_landmark_5(points : Points, distance : Distance) -> FaceLandmark5: + x = points[:, 0::2] + distance[:, 0::2] + y = points[:, 1::2] + distance[:, 1::2] + face_landmark_5 = numpy.stack((x, y), axis = -1) + return face_landmark_5 + + +def scale_face_landmark_5(face_landmark_5 : FaceLandmark5, scale : Scale) -> FaceLandmark5: + face_landmark_5_scale = face_landmark_5 - face_landmark_5[2] + face_landmark_5_scale *= scale + face_landmark_5_scale += face_landmark_5[2] + return face_landmark_5_scale + + +def convert_to_face_landmark_5(face_landmark_68 : FaceLandmark68) -> FaceLandmark5: + face_landmark_5 = numpy.array( + [ + numpy.mean(face_landmark_68[36:42], axis = 0), + numpy.mean(face_landmark_68[42:48], axis = 0), + face_landmark_68[30], + face_landmark_68[48], + face_landmark_68[54] + ]) + return face_landmark_5 + + +def estimate_face_angle(face_landmark_68 : FaceLandmark68) -> Angle: + x1, y1 = face_landmark_68[0] + x2, y2 = face_landmark_68[16] + theta = numpy.arctan2(y2 - y1, x2 - x1) + theta = numpy.degrees(theta) % 360 + angles = numpy.linspace(0, 360, 5) + index = numpy.argmin(numpy.abs(angles - theta)) + face_angle = int(angles[index] % 360) + return face_angle + + +def apply_nms(bounding_boxes : List[BoundingBox], scores : List[Score], score_threshold : float, nms_threshold : float) -> Sequence[int]: + bounding_boxes_norm = [ (x1, y1, x2 - x1, y2 - y1) for (x1, y1, x2, y2) in bounding_boxes ] + keep_indices = cv2.dnn.NMSBoxes(bounding_boxes_norm, scores, score_threshold = score_threshold, nms_threshold = nms_threshold) + return keep_indices + + +def get_nms_threshold(face_detector_model : FaceDetectorModel, face_detector_angles : List[Angle]) -> float: + if face_detector_model == 'many': + return 0.1 + if len(face_detector_angles) == 2: + return 0.3 + if len(face_detector_angles) == 3: + return 0.2 + if len(face_detector_angles) == 4: + return 0.1 + return 0.4 + + +def merge_matrix(temp_matrices : List[Matrix]) -> Matrix: + matrix = numpy.vstack([temp_matrices[0], [0, 0, 1]]) + + for temp_matrix in temp_matrices[1:]: + temp_matrix = numpy.vstack([ temp_matrix, [ 0, 0, 1 ] ]) + matrix = numpy.dot(temp_matrix, matrix) + + return matrix[:2, :] diff --git a/facefusion/face_landmarker.py b/facefusion/face_landmarker.py new file mode 100644 index 0000000000000000000000000000000000000000..7c293c42dacb4179e1b85349f04fc4905465cb3f --- /dev/null +++ b/facefusion/face_landmarker.py @@ -0,0 +1,240 @@ +from functools import lru_cache +from typing import Tuple + +import cv2 +import numpy + +from facefusion import inference_manager, state_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.face_helper import create_rotation_matrix_and_size, estimate_matrix_by_face_landmark_5, transform_points, warp_face_by_translation +from facefusion.filesystem import resolve_relative_path +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.types import Angle, BoundingBox, DownloadScope, DownloadSet, FaceLandmark5, FaceLandmark68, InferencePool, ModelSet, Prediction, Score, VisionFrame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + '2dfan4': + { + '__metadata__': + { + 'vendor': 'breadbread1984', + 'license': 'MIT', + 'year': 2018 + }, + 'hashes': + { + '2dfan4': + { + 'url': resolve_download_url('models-3.0.0', '2dfan4.hash'), + 'path': resolve_relative_path('../.assets/models/2dfan4.hash') + } + }, + 'sources': + { + '2dfan4': + { + 'url': resolve_download_url('models-3.0.0', '2dfan4.onnx'), + 'path': resolve_relative_path('../.assets/models/2dfan4.onnx') + } + }, + 'size': (256, 256) + }, + 'peppa_wutz': + { + '__metadata__': + { + 'vendor': 'Unknown', + 'license': 'Apache-2.0', + 'year': 2023 + }, + 'hashes': + { + 'peppa_wutz': + { + 'url': resolve_download_url('models-3.0.0', 'peppa_wutz.hash'), + 'path': resolve_relative_path('../.assets/models/peppa_wutz.hash') + } + }, + 'sources': + { + 'peppa_wutz': + { + 'url': resolve_download_url('models-3.0.0', 'peppa_wutz.onnx'), + 'path': resolve_relative_path('../.assets/models/peppa_wutz.onnx') + } + }, + 'size': (256, 256) + }, + 'fan_68_5': + { + '__metadata__': + { + 'vendor': 'FaceFusion', + 'license': 'OpenRAIL-M', + 'year': 2024 + }, + 'hashes': + { + 'fan_68_5': + { + 'url': resolve_download_url('models-3.0.0', 'fan_68_5.hash'), + 'path': resolve_relative_path('../.assets/models/fan_68_5.hash') + } + }, + 'sources': + { + 'fan_68_5': + { + 'url': resolve_download_url('models-3.0.0', 'fan_68_5.onnx'), + 'path': resolve_relative_path('../.assets/models/fan_68_5.onnx') + } + } + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('face_landmarker_model'), 'fan_68_5' ] + _, model_source_set = collect_model_downloads() + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('face_landmarker_model'), 'fan_68_5' ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_set = create_static_model_set('full') + model_hash_set =\ + { + 'fan_68_5': model_set.get('fan_68_5').get('hashes').get('fan_68_5') + } + model_source_set =\ + { + 'fan_68_5': model_set.get('fan_68_5').get('sources').get('fan_68_5') + } + + for face_landmarker_model in [ '2dfan4', 'peppa_wutz' ]: + if state_manager.get_item('face_landmarker_model') in [ 'many', face_landmarker_model ]: + model_hash_set[face_landmarker_model] = model_set.get(face_landmarker_model).get('hashes').get(face_landmarker_model) + model_source_set[face_landmarker_model] = model_set.get(face_landmarker_model).get('sources').get(face_landmarker_model) + + return model_hash_set, model_source_set + + +def pre_check() -> bool: + model_hash_set, model_source_set = collect_model_downloads() + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def detect_face_landmark(vision_frame : VisionFrame, bounding_box : BoundingBox, face_angle : Angle) -> Tuple[FaceLandmark68, Score]: + face_landmark_2dfan4 = None + face_landmark_peppa_wutz = None + face_landmark_score_2dfan4 = 0.0 + face_landmark_score_peppa_wutz = 0.0 + + if state_manager.get_item('face_landmarker_model') in [ 'many', '2dfan4' ]: + face_landmark_2dfan4, face_landmark_score_2dfan4 = detect_with_2dfan4(vision_frame, bounding_box, face_angle) + + if state_manager.get_item('face_landmarker_model') in [ 'many', 'peppa_wutz' ]: + face_landmark_peppa_wutz, face_landmark_score_peppa_wutz = detect_with_peppa_wutz(vision_frame, bounding_box, face_angle) + + if face_landmark_score_2dfan4 > face_landmark_score_peppa_wutz - 0.2: + return face_landmark_2dfan4, face_landmark_score_2dfan4 + return face_landmark_peppa_wutz, face_landmark_score_peppa_wutz + + +def detect_with_2dfan4(temp_vision_frame: VisionFrame, bounding_box: BoundingBox, face_angle: Angle) -> Tuple[FaceLandmark68, Score]: + model_size = create_static_model_set('full').get('2dfan4').get('size') + scale = 195 / numpy.subtract(bounding_box[2:], bounding_box[:2]).max().clip(1, None) + translation = (model_size[0] - numpy.add(bounding_box[2:], bounding_box[:2]) * scale) * 0.5 + rotation_matrix, rotation_size = create_rotation_matrix_and_size(face_angle, model_size) + crop_vision_frame, affine_matrix = warp_face_by_translation(temp_vision_frame, translation, scale, model_size) + crop_vision_frame = cv2.warpAffine(crop_vision_frame, rotation_matrix, rotation_size) + crop_vision_frame = conditional_optimize_contrast(crop_vision_frame) + crop_vision_frame = crop_vision_frame.transpose(2, 0, 1).astype(numpy.float32) / 255.0 + face_landmark_68, face_heatmap = forward_with_2dfan4(crop_vision_frame) + face_landmark_68 = face_landmark_68[:, :, :2][0] / 64 * 256 + face_landmark_68 = transform_points(face_landmark_68, cv2.invertAffineTransform(rotation_matrix)) + face_landmark_68 = transform_points(face_landmark_68, cv2.invertAffineTransform(affine_matrix)) + face_landmark_score_68 = numpy.amax(face_heatmap, axis = (2, 3)) + face_landmark_score_68 = numpy.mean(face_landmark_score_68) + face_landmark_score_68 = numpy.interp(face_landmark_score_68, [ 0, 0.9 ], [ 0, 1 ]) + return face_landmark_68, face_landmark_score_68 + + +def detect_with_peppa_wutz(temp_vision_frame : VisionFrame, bounding_box : BoundingBox, face_angle : Angle) -> Tuple[FaceLandmark68, Score]: + model_size = create_static_model_set('full').get('peppa_wutz').get('size') + scale = 195 / numpy.subtract(bounding_box[2:], bounding_box[:2]).max().clip(1, None) + translation = (model_size[0] - numpy.add(bounding_box[2:], bounding_box[:2]) * scale) * 0.5 + rotation_matrix, rotation_size = create_rotation_matrix_and_size(face_angle, model_size) + crop_vision_frame, affine_matrix = warp_face_by_translation(temp_vision_frame, translation, scale, model_size) + crop_vision_frame = cv2.warpAffine(crop_vision_frame, rotation_matrix, rotation_size) + crop_vision_frame = conditional_optimize_contrast(crop_vision_frame) + crop_vision_frame = crop_vision_frame.transpose(2, 0, 1).astype(numpy.float32) / 255.0 + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0) + prediction = forward_with_peppa_wutz(crop_vision_frame) + face_landmark_68 = prediction.reshape(-1, 3)[:, :2] / 64 * model_size[0] + face_landmark_68 = transform_points(face_landmark_68, cv2.invertAffineTransform(rotation_matrix)) + face_landmark_68 = transform_points(face_landmark_68, cv2.invertAffineTransform(affine_matrix)) + face_landmark_score_68 = prediction.reshape(-1, 3)[:, 2].mean() + face_landmark_score_68 = numpy.interp(face_landmark_score_68, [ 0, 0.95 ], [ 0, 1 ]) + return face_landmark_68, face_landmark_score_68 + + +def conditional_optimize_contrast(crop_vision_frame : VisionFrame) -> VisionFrame: + crop_vision_frame = cv2.cvtColor(crop_vision_frame, cv2.COLOR_RGB2Lab) + if numpy.mean(crop_vision_frame[:, :, 0]) < 30: #type:ignore[arg-type] + crop_vision_frame[:, :, 0] = cv2.createCLAHE(clipLimit = 2).apply(crop_vision_frame[:, :, 0]) + crop_vision_frame = cv2.cvtColor(crop_vision_frame, cv2.COLOR_Lab2RGB) + return crop_vision_frame + + +def estimate_face_landmark_68_5(face_landmark_5 : FaceLandmark5) -> FaceLandmark68: + affine_matrix = estimate_matrix_by_face_landmark_5(face_landmark_5, 'ffhq_512', (1, 1)) + face_landmark_5 = cv2.transform(face_landmark_5.reshape(1, -1, 2), affine_matrix).reshape(-1, 2) + face_landmark_68_5 = forward_fan_68_5(face_landmark_5) + face_landmark_68_5 = cv2.transform(face_landmark_68_5.reshape(1, -1, 2), cv2.invertAffineTransform(affine_matrix)).reshape(-1, 2) + return face_landmark_68_5 + + +def forward_with_2dfan4(crop_vision_frame : VisionFrame) -> Tuple[Prediction, Prediction]: + face_landmarker = get_inference_pool().get('2dfan4') + + with conditional_thread_semaphore(): + prediction = face_landmarker.run(None, + { + 'input': [ crop_vision_frame ] + }) + + return prediction + + +def forward_with_peppa_wutz(crop_vision_frame : VisionFrame) -> Prediction: + face_landmarker = get_inference_pool().get('peppa_wutz') + + with conditional_thread_semaphore(): + prediction = face_landmarker.run(None, + { + 'input': crop_vision_frame + })[0] + + return prediction + + +def forward_fan_68_5(face_landmark_5 : FaceLandmark5) -> FaceLandmark68: + face_landmarker = get_inference_pool().get('fan_68_5') + + with conditional_thread_semaphore(): + face_landmark_68_5 = face_landmarker.run(None, + { + 'input': [ face_landmark_5 ] + })[0][0] + + return face_landmark_68_5 diff --git a/facefusion/face_masker.py b/facefusion/face_masker.py new file mode 100644 index 0000000000000000000000000000000000000000..6603be216c61e91b3960141891b75327a09e732f --- /dev/null +++ b/facefusion/face_masker.py @@ -0,0 +1,279 @@ +from functools import lru_cache +from typing import List, Tuple + +import cv2 +import numpy + +import facefusion.choices +from facefusion import inference_manager, state_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 DownloadScope, DownloadSet, FaceLandmark68, FaceMaskArea, FaceMaskRegion, InferencePool, Mask, ModelSet, Padding, VisionFrame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'xseg_1': + { + '__metadata__': + { + 'vendor': 'DeepFaceLab', + 'license': 'GPL-3.0', + 'year': 2021 + }, + 'hashes': + { + 'face_occluder': + { + 'url': resolve_download_url('models-3.1.0', 'xseg_1.hash'), + 'path': resolve_relative_path('../.assets/models/xseg_1.hash') + } + }, + 'sources': + { + 'face_occluder': + { + 'url': resolve_download_url('models-3.1.0', 'xseg_1.onnx'), + 'path': resolve_relative_path('../.assets/models/xseg_1.onnx') + } + }, + 'size': (256, 256) + }, + 'xseg_2': + { + '__metadata__': + { + 'vendor': 'DeepFaceLab', + 'license': 'GPL-3.0', + 'year': 2021 + }, + 'hashes': + { + 'face_occluder': + { + 'url': resolve_download_url('models-3.1.0', 'xseg_2.hash'), + 'path': resolve_relative_path('../.assets/models/xseg_2.hash') + } + }, + 'sources': + { + 'face_occluder': + { + 'url': resolve_download_url('models-3.1.0', 'xseg_2.onnx'), + 'path': resolve_relative_path('../.assets/models/xseg_2.onnx') + } + }, + 'size': (256, 256) + }, + 'xseg_3': + { + '__metadata__': + { + 'vendor': 'DeepFaceLab', + 'license': 'GPL-3.0', + 'year': 2021 + }, + 'hashes': + { + 'face_occluder': + { + 'url': resolve_download_url('models-3.2.0', 'xseg_3.hash'), + 'path': resolve_relative_path('../.assets/models/xseg_3.hash') + } + }, + 'sources': + { + 'face_occluder': + { + 'url': resolve_download_url('models-3.2.0', 'xseg_3.onnx'), + 'path': resolve_relative_path('../.assets/models/xseg_3.onnx') + } + }, + 'size': (256, 256) + }, + 'bisenet_resnet_18': + { + '__metadata__': + { + 'vendor': 'yakhyo', + 'license': 'MIT', + 'year': 2024 + }, + 'hashes': + { + 'face_parser': + { + 'url': resolve_download_url('models-3.1.0', 'bisenet_resnet_18.hash'), + 'path': resolve_relative_path('../.assets/models/bisenet_resnet_18.hash') + } + }, + 'sources': + { + 'face_parser': + { + 'url': resolve_download_url('models-3.1.0', 'bisenet_resnet_18.onnx'), + 'path': resolve_relative_path('../.assets/models/bisenet_resnet_18.onnx') + } + }, + 'size': (512, 512) + }, + 'bisenet_resnet_34': + { + '__metadata__': + { + 'vendor': 'yakhyo', + 'license': 'MIT', + 'year': 2024 + }, + 'hashes': + { + 'face_parser': + { + 'url': resolve_download_url('models-3.0.0', 'bisenet_resnet_34.hash'), + 'path': resolve_relative_path('../.assets/models/bisenet_resnet_34.hash') + } + }, + 'sources': + { + 'face_parser': + { + 'url': resolve_download_url('models-3.0.0', 'bisenet_resnet_34.onnx'), + 'path': resolve_relative_path('../.assets/models/bisenet_resnet_34.onnx') + } + }, + 'size': (512, 512) + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('face_occluder_model'), state_manager.get_item('face_parser_model') ] + _, model_source_set = collect_model_downloads() + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('face_occluder_model'), state_manager.get_item('face_parser_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_set = create_static_model_set('full') + model_hash_set = {} + model_source_set = {} + + for face_occluder_model in [ 'xseg_1', 'xseg_2', 'xseg_3' ]: + if state_manager.get_item('face_occluder_model') in [ 'many', face_occluder_model ]: + model_hash_set[face_occluder_model] = model_set.get(face_occluder_model).get('hashes').get('face_occluder') + model_source_set[face_occluder_model] = model_set.get(face_occluder_model).get('sources').get('face_occluder') + + for face_parser_model in [ 'bisenet_resnet_18', 'bisenet_resnet_34' ]: + if state_manager.get_item('face_parser_model') == face_parser_model: + model_hash_set[face_parser_model] = model_set.get(face_parser_model).get('hashes').get('face_parser') + model_source_set[face_parser_model] = model_set.get(face_parser_model).get('sources').get('face_parser') + + return model_hash_set, model_source_set + + +def pre_check() -> bool: + model_hash_set, model_source_set = collect_model_downloads() + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def create_box_mask(crop_vision_frame : VisionFrame, face_mask_blur : float, face_mask_padding : Padding) -> Mask: + crop_size = crop_vision_frame.shape[:2][::-1] + blur_amount = int(crop_size[0] * 0.5 * face_mask_blur) + blur_area = max(blur_amount // 2, 1) + box_mask : Mask = numpy.ones(crop_size).astype(numpy.float32) + box_mask[:max(blur_area, int(crop_size[1] * face_mask_padding[0] / 100)), :] = 0 + box_mask[-max(blur_area, int(crop_size[1] * face_mask_padding[2] / 100)):, :] = 0 + box_mask[:, :max(blur_area, int(crop_size[0] * face_mask_padding[3] / 100))] = 0 + box_mask[:, -max(blur_area, int(crop_size[0] * face_mask_padding[1] / 100)):] = 0 + + if blur_amount > 0: + box_mask = cv2.GaussianBlur(box_mask, (0, 0), blur_amount * 0.25) + return box_mask + + +def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask: + temp_masks = [] + + if state_manager.get_item('face_occluder_model') == 'many': + model_names = [ 'xseg_1', 'xseg_2', 'xseg_3' ] + else: + model_names = [ state_manager.get_item('face_occluder_model') ] + + for model_name in model_names: + model_size = create_static_model_set('full').get(model_name).get('size') + prepare_vision_frame = cv2.resize(crop_vision_frame, model_size) + prepare_vision_frame = numpy.expand_dims(prepare_vision_frame, axis = 0).astype(numpy.float32) / 255.0 + prepare_vision_frame = prepare_vision_frame.transpose(0, 1, 2, 3) + temp_mask = forward_occlude_face(prepare_vision_frame, model_name) + temp_mask = temp_mask.transpose(0, 1, 2).clip(0, 1).astype(numpy.float32) + temp_mask = cv2.resize(temp_mask, crop_vision_frame.shape[:2][::-1]) + temp_masks.append(temp_mask) + + occlusion_mask = numpy.minimum.reduce(temp_masks) + occlusion_mask = (cv2.GaussianBlur(occlusion_mask.clip(0, 1), (0, 0), 5).clip(0.5, 1) - 0.5) * 2 + return occlusion_mask + + +def create_area_mask(crop_vision_frame : VisionFrame, face_landmark_68 : FaceLandmark68, face_mask_areas : List[FaceMaskArea]) -> Mask: + crop_size = crop_vision_frame.shape[:2][::-1] + landmark_points = [] + + for face_mask_area in face_mask_areas: + if face_mask_area in facefusion.choices.face_mask_area_set: + landmark_points.extend(facefusion.choices.face_mask_area_set.get(face_mask_area)) + + convex_hull = cv2.convexHull(face_landmark_68[landmark_points].astype(numpy.int32)) + area_mask = numpy.zeros(crop_size).astype(numpy.float32) + cv2.fillConvexPoly(area_mask, convex_hull, 1.0) # type: ignore[call-overload] + area_mask = (cv2.GaussianBlur(area_mask.clip(0, 1), (0, 0), 5).clip(0.5, 1) - 0.5) * 2 + return area_mask + + +def create_region_mask(crop_vision_frame : VisionFrame, face_mask_regions : List[FaceMaskRegion]) -> Mask: + model_name = state_manager.get_item('face_parser_model') + model_size = create_static_model_set('full').get(model_name).get('size') + prepare_vision_frame = cv2.resize(crop_vision_frame, model_size) + prepare_vision_frame = prepare_vision_frame[:, :, ::-1].astype(numpy.float32) / 255.0 + prepare_vision_frame = numpy.subtract(prepare_vision_frame, numpy.array([ 0.485, 0.456, 0.406 ]).astype(numpy.float32)) + prepare_vision_frame = numpy.divide(prepare_vision_frame, numpy.array([ 0.229, 0.224, 0.225 ]).astype(numpy.float32)) + prepare_vision_frame = numpy.expand_dims(prepare_vision_frame, axis = 0) + prepare_vision_frame = prepare_vision_frame.transpose(0, 3, 1, 2) + region_mask = forward_parse_face(prepare_vision_frame) + region_mask = numpy.isin(region_mask.argmax(0), [ facefusion.choices.face_mask_region_set.get(face_mask_region) for face_mask_region in face_mask_regions ]) + region_mask = cv2.resize(region_mask.astype(numpy.float32), crop_vision_frame.shape[:2][::-1]) + region_mask = (cv2.GaussianBlur(region_mask.clip(0, 1), (0, 0), 5).clip(0.5, 1) - 0.5) * 2 + return region_mask + + +def forward_occlude_face(prepare_vision_frame : VisionFrame, model_name : str) -> Mask: + face_occluder = get_inference_pool().get(model_name) + + with conditional_thread_semaphore(): + occlusion_mask : Mask = face_occluder.run(None, + { + 'input': prepare_vision_frame + })[0][0] + + return occlusion_mask + + +def forward_parse_face(prepare_vision_frame : VisionFrame) -> Mask: + model_name = state_manager.get_item('face_parser_model') + face_parser = get_inference_pool().get(model_name) + + with conditional_thread_semaphore(): + region_mask : Mask = face_parser.run(None, + { + 'input': prepare_vision_frame + })[0][0] + + return region_mask diff --git a/facefusion/face_recognizer.py b/facefusion/face_recognizer.py new file mode 100644 index 0000000000000000000000000000000000000000..99012219d1e3c32aec151919429141aeaa8e290e --- /dev/null +++ b/facefusion/face_recognizer.py @@ -0,0 +1,93 @@ +from functools import lru_cache +from typing import Tuple + +import numpy + +from facefusion import inference_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.face_helper import warp_face_by_face_landmark_5 +from facefusion.filesystem import resolve_relative_path +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.types import DownloadScope, Embedding, FaceLandmark5, InferencePool, ModelOptions, ModelSet, VisionFrame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'arcface': + { + '__metadata__': + { + 'vendor': 'InsightFace', + 'license': 'Non-Commercial', + 'year': 2018 + }, + 'hashes': + { + 'face_recognizer': + { + 'url': resolve_download_url('models-3.0.0', 'arcface_w600k_r50.hash'), + 'path': resolve_relative_path('../.assets/models/arcface_w600k_r50.hash') + } + }, + 'sources': + { + 'face_recognizer': + { + 'url': resolve_download_url('models-3.0.0', 'arcface_w600k_r50.onnx'), + 'path': resolve_relative_path('../.assets/models/arcface_w600k_r50.onnx') + } + }, + 'template': 'arcface_112_v2', + 'size': (112, 112) + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ 'arcface' ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ 'arcface' ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + return create_static_model_set('full').get('arcface') + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def calculate_face_embedding(temp_vision_frame : VisionFrame, face_landmark_5 : FaceLandmark5) -> Tuple[Embedding, Embedding]: + model_template = get_model_options().get('template') + model_size = get_model_options().get('size') + crop_vision_frame, matrix = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5, model_template, model_size) + crop_vision_frame = crop_vision_frame / 127.5 - 1 + crop_vision_frame = crop_vision_frame[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0) + face_embedding = forward(crop_vision_frame) + face_embedding = face_embedding.ravel() + face_embedding_norm = face_embedding / numpy.linalg.norm(face_embedding) + return face_embedding, face_embedding_norm + + +def forward(crop_vision_frame : VisionFrame) -> Embedding: + face_recognizer = get_inference_pool().get('face_recognizer') + + with conditional_thread_semaphore(): + face_embedding = face_recognizer.run(None, + { + 'input': crop_vision_frame + })[0] + + return face_embedding diff --git a/facefusion/face_selector.py b/facefusion/face_selector.py new file mode 100644 index 0000000000000000000000000000000000000000..c4246faf66285eb05f146b0b060d8f28f5d9f38f --- /dev/null +++ b/facefusion/face_selector.py @@ -0,0 +1,130 @@ +from typing import List + +import numpy + +from facefusion import state_manager +from facefusion.face_analyser import get_many_faces, get_one_face +from facefusion.types import Face, FaceSelectorOrder, Gender, Race, Score, VisionFrame + + +def select_faces(reference_vision_frame : VisionFrame, target_vision_frame : VisionFrame) -> List[Face]: + target_faces = get_many_faces([ target_vision_frame ]) + + if state_manager.get_item('face_selector_mode') == 'many': + return sort_and_filter_faces(target_faces) + + if state_manager.get_item('face_selector_mode') == 'one': + target_face = get_one_face(sort_and_filter_faces(target_faces)) + if target_face: + return [ target_face ] + + if state_manager.get_item('face_selector_mode') == 'reference': + reference_faces = get_many_faces([ reference_vision_frame ]) + reference_faces = sort_and_filter_faces(reference_faces) + reference_face = get_one_face(reference_faces, state_manager.get_item('reference_face_position')) + if reference_face: + match_faces = find_match_faces([ reference_face ], target_faces, state_manager.get_item('reference_face_distance')) + return match_faces + + return [] + + +def find_match_faces(reference_faces : List[Face], target_faces : List[Face], face_distance : float) -> List[Face]: + match_faces : List[Face] = [] + + for reference_face in reference_faces: + if reference_face: + for index, target_face in enumerate(target_faces): + if compare_faces(target_face, reference_face, face_distance): + match_faces.append(target_faces[index]) + + return match_faces + + +def compare_faces(face : Face, reference_face : Face, face_distance : float) -> bool: + current_face_distance = calculate_face_distance(face, reference_face) + current_face_distance = float(numpy.interp(current_face_distance, [ 0, 2 ], [ 0, 1 ])) + return current_face_distance < face_distance + + +def calculate_face_distance(face : Face, reference_face : Face) -> float: + if hasattr(face, 'embedding_norm') and hasattr(reference_face, 'embedding_norm'): + return 1 - numpy.dot(face.embedding_norm, reference_face.embedding_norm) + return 0 + + +def sort_and_filter_faces(faces : List[Face]) -> List[Face]: + if faces: + if state_manager.get_item('face_selector_order'): + faces = sort_faces_by_order(faces, state_manager.get_item('face_selector_order')) + if state_manager.get_item('face_selector_gender'): + faces = filter_faces_by_gender(faces, state_manager.get_item('face_selector_gender')) + if state_manager.get_item('face_selector_race'): + faces = filter_faces_by_race(faces, state_manager.get_item('face_selector_race')) + if state_manager.get_item('face_selector_age_start') or state_manager.get_item('face_selector_age_end'): + faces = filter_faces_by_age(faces, state_manager.get_item('face_selector_age_start'), state_manager.get_item('face_selector_age_end')) + return faces + + +def sort_faces_by_order(faces : List[Face], order : FaceSelectorOrder) -> List[Face]: + if order == 'left-right': + return sorted(faces, key = get_bounding_box_left) + if order == 'right-left': + return sorted(faces, key = get_bounding_box_left, reverse = True) + if order == 'top-bottom': + return sorted(faces, key = get_bounding_box_top) + if order == 'bottom-top': + return sorted(faces, key = get_bounding_box_top, reverse = True) + if order == 'small-large': + return sorted(faces, key = get_bounding_box_area) + if order == 'large-small': + return sorted(faces, key = get_bounding_box_area, reverse = True) + if order == 'best-worst': + return sorted(faces, key = get_face_detector_score, reverse = True) + if order == 'worst-best': + return sorted(faces, key = get_face_detector_score) + return faces + + +def get_bounding_box_left(face : Face) -> float: + return face.bounding_box[0] + + +def get_bounding_box_top(face : Face) -> float: + return face.bounding_box[1] + + +def get_bounding_box_area(face : Face) -> float: + return (face.bounding_box[2] - face.bounding_box[0]) * (face.bounding_box[3] - face.bounding_box[1]) + + +def get_face_detector_score(face : Face) -> Score: + return face.score_set.get('detector') + + +def filter_faces_by_gender(faces : List[Face], gender : Gender) -> List[Face]: + filter_faces = [] + + for face in faces: + if face.gender == gender: + filter_faces.append(face) + return filter_faces + + +def filter_faces_by_age(faces : List[Face], face_selector_age_start : int, face_selector_age_end : int) -> List[Face]: + filter_faces = [] + age = range(face_selector_age_start, face_selector_age_end) + + for face in faces: + if set(face.age) & set(age): + filter_faces.append(face) + return filter_faces + + +def filter_faces_by_race(faces : List[Face], race : Race) -> List[Face]: + filter_faces = [] + + for face in faces: + if face.race == race: + filter_faces.append(face) + return filter_faces diff --git a/facefusion/face_store.py b/facefusion/face_store.py new file mode 100644 index 0000000000000000000000000000000000000000..2e7c96ac0b302a7b37731bf5856f47ba8561a31f --- /dev/null +++ b/facefusion/face_store.py @@ -0,0 +1,28 @@ +from typing import List, Optional + +from facefusion.hash_helper import create_hash +from facefusion.types import Face, FaceStore, VisionFrame + +FACE_STORE : FaceStore =\ +{ + 'static_faces': {} +} + + +def get_face_store() -> FaceStore: + return FACE_STORE + + +def get_static_faces(vision_frame : VisionFrame) -> Optional[List[Face]]: + vision_hash = create_hash(vision_frame.tobytes()) + return FACE_STORE.get('static_faces').get(vision_hash) + + +def set_static_faces(vision_frame : VisionFrame, faces : List[Face]) -> None: + vision_hash = create_hash(vision_frame.tobytes()) + if vision_hash: + FACE_STORE['static_faces'][vision_hash] = faces + + +def clear_static_faces() -> None: + FACE_STORE['static_faces'].clear() diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py new file mode 100644 index 0000000000000000000000000000000000000000..a43b887c718730d0835224e8e995bfb1b3dba5ab --- /dev/null +++ b/facefusion/ffmpeg.py @@ -0,0 +1,291 @@ +import os +import subprocess +import tempfile +from functools import partial +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.filesystem import get_file_format, remove_file +from facefusion.temp_helper import get_temp_file_path, get_temp_frames_pattern +from facefusion.types import AudioBuffer, AudioEncoder, Command, EncoderSet, Fps, Resolution, UpdateProgress, VideoEncoder, VideoFormat +from facefusion.vision import detect_video_duration, detect_video_fps, pack_resolution, predict_video_frame_total + + +def run_ffmpeg_with_progress(commands : List[Command], update_progress : UpdateProgress) -> subprocess.Popen[bytes]: + log_level = state_manager.get_item('log_level') + commands.extend(ffmpeg_builder.set_progress()) + commands.extend(ffmpeg_builder.cast_stream()) + commands = ffmpeg_builder.run(commands) + process = subprocess.Popen(commands, stderr = subprocess.PIPE, stdout = subprocess.PIPE) + + while process_manager.is_processing(): + try: + while __line__ := process.stdout.readline().decode().lower(): + if process_manager.is_stopping(): + process.terminate() + + if 'frame=' in __line__: + _, frame_number = __line__.split('frame=') + update_progress(int(frame_number)) + + if log_level == 'debug': + log_debug(process) + process.wait(timeout = 0.5) + except subprocess.TimeoutExpired: + continue + return process + + return process + + +def update_progress(progress : tqdm, frame_number : int) -> None: + progress.update(frame_number - progress.n) + + +def run_ffmpeg(commands : List[Command]) -> subprocess.Popen[bytes]: + log_level = state_manager.get_item('log_level') + commands = ffmpeg_builder.run(commands) + process = subprocess.Popen(commands, stderr = subprocess.PIPE, stdout = subprocess.PIPE) + + while process_manager.is_processing(): + try: + if log_level == 'debug': + log_debug(process) + process.wait(timeout = 0.5) + except subprocess.TimeoutExpired: + continue + return process + + if process_manager.is_stopping(): + process.terminate() + + return process + + +def open_ffmpeg(commands : List[Command]) -> subprocess.Popen[bytes]: + commands = ffmpeg_builder.run(commands) + return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.PIPE) + + +def log_debug(process : subprocess.Popen[bytes]) -> None: + _, stderr = process.communicate() + errors = stderr.decode().split(os.linesep) + + for error in errors: + if error.strip(): + logger.debug(error.strip(), __name__) + + +def get_available_encoder_set() -> EncoderSet: + available_encoder_set : EncoderSet =\ + { + 'audio': [], + 'video': [] + } + commands = ffmpeg_builder.chain( + ffmpeg_builder.get_encoders() + ) + process = run_ffmpeg(commands) + + while line := process.stdout.readline().decode().lower(): + if line.startswith(' a'): + audio_encoder = line.split()[1] + + if audio_encoder in facefusion.choices.output_audio_encoders: + index = facefusion.choices.output_audio_encoders.index(audio_encoder) #type:ignore[arg-type] + available_encoder_set['audio'].insert(index, audio_encoder) #type:ignore[arg-type] + if line.startswith(' v'): + video_encoder = line.split()[1] + + if video_encoder in facefusion.choices.output_video_encoders: + index = facefusion.choices.output_video_encoders.index(video_encoder) #type:ignore[arg-type] + available_encoder_set['video'].insert(index, video_encoder) #type:ignore[arg-type] + + return available_encoder_set + + +def extract_frames(target_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) + temp_frames_pattern = get_temp_frames_pattern(target_path, '%08d') + commands = ffmpeg_builder.chain( + ffmpeg_builder.set_input(target_path), + ffmpeg_builder.set_media_resolution(pack_resolution(temp_video_resolution)), + ffmpeg_builder.set_frame_quality(0), + ffmpeg_builder.select_frame_range(trim_frame_start, trim_frame_end, temp_video_fps), + ffmpeg_builder.prevent_frame_drop(), + ffmpeg_builder.set_output(temp_frames_pattern) + ) + + with tqdm(total = extract_frame_total, desc = translator.get('extracting'), unit = 'frame', ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: + process = run_ffmpeg_with_progress(commands, partial(update_progress, progress)) + return process.returncode == 0 + + +def copy_image(target_path : str, temp_image_resolution : Resolution) -> bool: + temp_image_path = get_temp_file_path(target_path) + commands = ffmpeg_builder.chain( + ffmpeg_builder.set_input(target_path), + ffmpeg_builder.set_media_resolution(pack_resolution(temp_image_resolution)), + ffmpeg_builder.set_image_quality(target_path, 100), + ffmpeg_builder.force_output(temp_image_path) + ) + return run_ffmpeg(commands).returncode == 0 + + +def finalize_image(target_path : str, output_path : str, output_image_resolution : Resolution) -> bool: + output_image_quality = state_manager.get_item('output_image_quality') + temp_image_path = get_temp_file_path(target_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_image_quality(target_path, output_image_quality), + ffmpeg_builder.force_output(output_path) + ) + return run_ffmpeg(commands).returncode == 0 + + +def read_audio_buffer(target_path : str, audio_sample_rate : int, audio_sample_size : int, audio_channel_total : int) -> Optional[AudioBuffer]: + commands = ffmpeg_builder.chain( + ffmpeg_builder.set_input(target_path), + ffmpeg_builder.ignore_video_stream(), + ffmpeg_builder.set_audio_sample_rate(audio_sample_rate), + ffmpeg_builder.set_audio_sample_size(audio_sample_size), + ffmpeg_builder.set_audio_channel_total(audio_channel_total), + ffmpeg_builder.cast_stream() + ) + + process = open_ffmpeg(commands) + audio_buffer, _ = process.communicate() + if process.returncode == 0: + return audio_buffer + return None + + +def restore_audio(target_path : str, output_path : str, trim_frame_start : int, trim_frame_end : int) -> bool: + 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) + temp_video_path = get_temp_file_path(target_path) + temp_video_format = cast(VideoFormat, get_file_format(temp_video_path)) + temp_video_duration = detect_video_duration(temp_video_path) + + output_audio_encoder = fix_audio_encoder(temp_video_format, output_audio_encoder) + commands = ffmpeg_builder.chain( + ffmpeg_builder.set_input(temp_video_path), + ffmpeg_builder.select_media_range(trim_frame_start, trim_frame_end, target_video_fps), + ffmpeg_builder.set_input(target_path), + ffmpeg_builder.copy_video_encoder(), + ffmpeg_builder.set_audio_encoder(output_audio_encoder), + ffmpeg_builder.set_audio_quality(output_audio_encoder, output_audio_quality), + ffmpeg_builder.set_audio_volume(output_audio_volume), + ffmpeg_builder.select_media_stream('0:v:0'), + ffmpeg_builder.select_media_stream('1:a:0'), + ffmpeg_builder.set_video_duration(temp_video_duration), + ffmpeg_builder.force_output(output_path) + ) + return run_ffmpeg(commands).returncode == 0 + + +def replace_audio(target_path : str, audio_path : str, output_path : str) -> bool: + 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') + temp_video_path = get_temp_file_path(target_path) + temp_video_format = cast(VideoFormat, get_file_format(temp_video_path)) + temp_video_duration = detect_video_duration(temp_video_path) + + output_audio_encoder = fix_audio_encoder(temp_video_format, output_audio_encoder) + commands = ffmpeg_builder.chain( + ffmpeg_builder.set_input(temp_video_path), + ffmpeg_builder.set_input(audio_path), + ffmpeg_builder.copy_video_encoder(), + ffmpeg_builder.set_audio_encoder(output_audio_encoder), + ffmpeg_builder.set_audio_quality(output_audio_encoder, output_audio_quality), + ffmpeg_builder.set_audio_volume(output_audio_volume), + ffmpeg_builder.set_video_duration(temp_video_duration), + ffmpeg_builder.force_output(output_path) + ) + return run_ffmpeg(commands).returncode == 0 + + +def merge_video(target_path : str, temp_video_fps : Fps, output_video_resolution : Resolution, output_video_fps : Fps, trim_frame_start : int, trim_frame_end : int) -> bool: + 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) + temp_video_path = get_temp_file_path(target_path) + temp_video_format = cast(VideoFormat, get_file_format(temp_video_path)) + temp_frames_pattern = get_temp_frames_pattern(target_path, '%08d') + + output_video_encoder = fix_video_encoder(temp_video_format, output_video_encoder) + commands = ffmpeg_builder.chain( + ffmpeg_builder.set_input_fps(temp_video_fps), + ffmpeg_builder.set_input(temp_frames_pattern), + ffmpeg_builder.set_media_resolution(pack_resolution(output_video_resolution)), + ffmpeg_builder.set_video_encoder(output_video_encoder), + 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.set_pixel_format(output_video_encoder), + ffmpeg_builder.force_output(temp_video_path) + ) + + with tqdm(total = merge_frame_total, desc = translator.get('merging'), unit = 'frame', ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: + process = run_ffmpeg_with_progress(commands, partial(update_progress, progress)) + return process.returncode == 0 + + +def concat_video(output_path : str, temp_output_paths : List[str]) -> bool: + concat_video_path = tempfile.mktemp() + + with open(concat_video_path, 'w') as concat_video_file: + for temp_output_path in temp_output_paths: + concat_video_file.write('file \'' + os.path.abspath(temp_output_path) + '\'' + os.linesep) + concat_video_file.flush() + concat_video_file.close() + + output_path = os.path.abspath(output_path) + commands = ffmpeg_builder.chain( + ffmpeg_builder.unsafe_concat(), + ffmpeg_builder.set_input(concat_video_file.name), + ffmpeg_builder.copy_video_encoder(), + ffmpeg_builder.copy_audio_encoder(), + ffmpeg_builder.force_output(output_path) + ) + process = run_ffmpeg(commands) + process.communicate() + remove_file(concat_video_path) + return process.returncode == 0 + + +def fix_audio_encoder(video_format : VideoFormat, audio_encoder : AudioEncoder) -> AudioEncoder: + if video_format == 'avi' and audio_encoder == 'libopus': + return 'aac' + if video_format in [ 'm4v', 'mpeg', 'wmv' ]: + return 'aac' + if video_format == 'mov' and audio_encoder in [ 'flac', 'libopus' ]: + return 'aac' + if video_format == 'mxf': + return 'pcm_s16le' + if video_format == 'webm': + return 'libopus' + return audio_encoder + + +def fix_video_encoder(video_format : VideoFormat, video_encoder : VideoEncoder) -> VideoEncoder: + if video_format in [ 'm4v', 'mpeg', 'mxf', 'wmv' ]: + return 'libx264' + if video_format in [ 'mkv', 'mp4' ] and video_encoder == 'rawvideo': + return 'libx264' + if video_format == 'mov' and video_encoder == 'libvpx-vp9': + return 'libx264' + if video_format == 'webm': + return 'libvpx-vp9' + return video_encoder diff --git a/facefusion/ffmpeg_builder.py b/facefusion/ffmpeg_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..0ac1840112d2e2aca20396d2cb61776916f79817 --- /dev/null +++ b/facefusion/ffmpeg_builder.py @@ -0,0 +1,267 @@ +import itertools +import shutil +from typing import List, Optional + +import numpy + +from facefusion.filesystem import get_file_format +from facefusion.types import AudioEncoder, Command, CommandSet, Duration, Fps, StreamMode, VideoEncoder, VideoPreset + + +def run(commands : List[Command]) -> List[Command]: + return [ shutil.which('ffmpeg'), '-loglevel', 'error' ] + commands + + +def chain(*commands : List[Command]) -> List[Command]: + return list(itertools.chain(*commands)) + + +def concat(*__commands__ : List[Command]) -> List[Command]: + commands = [] + command_set : CommandSet = {} + + for command in __commands__: + for argument, value in zip(command[::2], command[1::2]): + command_set.setdefault(argument, []).append(value) + + for argument, values in command_set.items(): + commands.append(argument) + commands.append(','.join(values)) + + return commands + + +def get_encoders() -> List[Command]: + return [ '-encoders' ] + + +def set_hardware_accelerator(value : str) -> List[Command]: + return [ '-hwaccel', value ] + + +def set_progress() -> List[Command]: + return [ '-progress' ] + + +def set_input(input_path : str) -> List[Command]: + return [ '-i', input_path ] + + +def set_input_fps(input_fps : Fps) -> List[Command]: + return [ '-r', str(input_fps)] + + +def set_output(output_path : str) -> List[Command]: + return [ output_path ] + + +def force_output(output_path : str) -> List[Command]: + return [ '-y', output_path ] + + +def cast_stream() -> List[Command]: + return [ '-' ] + + +def set_stream_mode(stream_mode : StreamMode) -> List[Command]: + if stream_mode == 'udp': + return [ '-f', 'mpegts' ] + if stream_mode == 'v4l2': + return [ '-f', 'v4l2' ] + return [] + + +def set_stream_quality(stream_quality : int) -> List[Command]: + return [ '-b:v', str(stream_quality) + 'k' ] + + +def unsafe_concat() -> List[Command]: + return [ '-f', 'concat', '-safe', '0' ] + + +def set_pixel_format(video_encoder : VideoEncoder) -> List[Command]: + if video_encoder == 'rawvideo': + return [ '-pix_fmt', 'rgb24' ] + if video_encoder == 'libvpx-vp9': + return [ '-pix_fmt', 'yuva420p' ] + return [ '-pix_fmt', 'yuv420p' ] + + +def set_frame_quality(frame_quality : int) -> List[Command]: + return [ '-q:v', str(frame_quality) ] + + +def select_frame_range(frame_start : int, frame_end : int, video_fps : Fps) -> List[Command]: + if isinstance(frame_start, int) and isinstance(frame_end, int): + return [ '-vf', 'trim=start_frame=' + str(frame_start) + ':end_frame=' + str(frame_end) + ',fps=' + str(video_fps) ] + if isinstance(frame_start, int): + return [ '-vf', 'trim=start_frame=' + str(frame_start) + ',fps=' + str(video_fps) ] + if isinstance(frame_end, int): + return [ '-vf', 'trim=end_frame=' + str(frame_end) + ',fps=' + str(video_fps) ] + return [ '-vf', 'fps=' + str(video_fps) ] + + +def prevent_frame_drop() -> List[Command]: + return [ '-vsync', '0' ] + + +def select_media_range(frame_start : int, frame_end : int, media_fps : Fps) -> List[Command]: + commands = [] + + if isinstance(frame_start, int): + commands.extend([ '-ss', str(frame_start / media_fps) ]) + if isinstance(frame_end, int): + commands.extend([ '-to', str(frame_end / media_fps) ]) + return commands + + +def select_media_stream(media_stream : str) -> List[Command]: + return [ '-map', media_stream ] + + +def set_media_resolution(video_resolution : str) -> List[Command]: + return [ '-s', video_resolution ] + + +def set_image_quality(image_path : str, image_quality : int) -> List[Command]: + if get_file_format(image_path) == 'webp': + return [ '-q:v', str(image_quality) ] + + image_compression = round(31 - (image_quality * 0.31)) + return [ '-q:v', str(image_compression) ] + + +def set_audio_encoder(audio_codec : str) -> List[Command]: + return [ '-c:a', audio_codec ] + + +def copy_audio_encoder() -> List[Command]: + return set_audio_encoder('copy') + + +def set_audio_sample_rate(audio_sample_rate : int) -> List[Command]: + return [ '-ar', str(audio_sample_rate) ] + + +def set_audio_sample_size(audio_sample_size : int) -> List[Command]: + if audio_sample_size == 16: + return [ '-f', 's16le' ] + if audio_sample_size == 32: + return [ '-f', 's32le' ] + return [] + + +def set_audio_channel_total(audio_channel_total : int) -> List[Command]: + return [ '-ac', str(audio_channel_total) ] + + +def set_audio_quality(audio_encoder : AudioEncoder, audio_quality : int) -> List[Command]: + if audio_encoder == 'aac': + audio_compression = numpy.round(numpy.interp(audio_quality, [ 0, 100 ], [ 0.1, 2.0 ]), 1).astype(float).item() + return [ '-q:a', str(audio_compression) ] + if audio_encoder == 'libmp3lame': + audio_compression = numpy.round(numpy.interp(audio_quality, [ 0, 100 ], [ 9, 0 ])).astype(int).item() + return [ '-q:a', str(audio_compression) ] + if audio_encoder == 'libopus': + audio_bit_rate = numpy.round(numpy.interp(audio_quality, [ 0, 100 ], [ 64, 256 ])).astype(int).item() + return [ '-b:a', str(audio_bit_rate) + 'k' ] + if audio_encoder == 'libvorbis': + audio_compression = numpy.round(numpy.interp(audio_quality, [ 0, 100 ], [ -1, 10 ]), 1).astype(float).item() + return [ '-q:a', str(audio_compression) ] + return [] + + +def set_audio_volume(audio_volume : int) -> List[Command]: + return [ '-filter:a', 'volume=' + str(audio_volume / 100) ] + + +def set_video_encoder(video_encoder : str) -> List[Command]: + return [ '-c:v', video_encoder ] + + +def copy_video_encoder() -> List[Command]: + return set_video_encoder('copy') + + +def set_video_quality(video_encoder : VideoEncoder, video_quality : int) -> List[Command]: + if video_encoder in [ 'libx264', 'libx264rgb', 'libx265' ]: + video_compression = numpy.round(numpy.interp(video_quality, [ 0, 100 ], [ 51, 0 ])).astype(int).item() + return [ '-crf', str(video_compression) ] + if video_encoder == 'libvpx-vp9': + video_compression = numpy.round(numpy.interp(video_quality, [ 0, 100 ], [ 63, 0 ])).astype(int).item() + return [ '-crf', str(video_compression) ] + if video_encoder in [ 'h264_nvenc', 'hevc_nvenc' ]: + video_compression = numpy.round(numpy.interp(video_quality, [ 0, 100 ], [ 51, 0 ])).astype(int).item() + return [ '-cq', str(video_compression) ] + if video_encoder in [ 'h264_amf', 'hevc_amf' ]: + video_compression = numpy.round(numpy.interp(video_quality, [ 0, 100 ], [ 51, 0 ])).astype(int).item() + return [ '-qp_i', str(video_compression), '-qp_p', str(video_compression), '-qp_b', str(video_compression) ] + if video_encoder in [ 'h264_qsv', 'hevc_qsv' ]: + video_compression = numpy.round(numpy.interp(video_quality, [ 0, 100 ], [ 51, 0 ])).astype(int).item() + return [ '-qp', str(video_compression) ] + if video_encoder in [ 'h264_videotoolbox', 'hevc_videotoolbox' ]: + video_bit_rate = numpy.round(numpy.interp(video_quality, [ 0, 100 ], [ 1024, 50512 ])).astype(int).item() + return [ '-b:v', str(video_bit_rate) + 'k' ] + return [] + + +def set_video_preset(video_encoder : VideoEncoder, video_preset : VideoPreset) -> List[Command]: + if video_encoder in [ 'libx264', 'libx264rgb', 'libx265' ]: + return [ '-preset', video_preset ] + if video_encoder in [ 'h264_nvenc', 'hevc_nvenc' ]: + return [ '-preset', map_nvenc_preset(video_preset) ] + if video_encoder in [ 'h264_amf', 'hevc_amf' ]: + return [ '-quality', map_amf_preset(video_preset) ] + if video_encoder in [ 'h264_qsv', 'hevc_qsv' ]: + return [ '-preset', map_qsv_preset(video_preset) ] + return [] + + +def set_video_fps(video_fps : Fps) -> List[Command]: + return [ '-vf', 'fps=' + str(video_fps) ] + + +def set_video_duration(video_duration : Duration) -> List[Command]: + return [ '-t', str(video_duration) ] + + +def keep_video_alpha(video_encoder : VideoEncoder) -> List[Command]: + if video_encoder == 'libvpx-vp9': + return [ '-vf', 'format=yuva420p' ] + return [] + + +def capture_video() -> List[Command]: + return [ '-f', 'rawvideo', '-pix_fmt', 'rgb24' ] + + +def ignore_video_stream() -> List[Command]: + return [ '-vn' ] + + +def map_nvenc_preset(video_preset : VideoPreset) -> Optional[str]: + if video_preset in [ 'ultrafast', 'superfast', 'veryfast', 'faster', 'fast' ]: + return 'fast' + if video_preset == 'medium': + return 'medium' + if video_preset in [ 'slow', 'slower', 'veryslow' ]: + return 'slow' + return None + + +def map_amf_preset(video_preset : VideoPreset) -> Optional[str]: + if video_preset in [ 'ultrafast', 'superfast', 'veryfast' ]: + return 'speed' + if video_preset in [ 'faster', 'fast', 'medium' ]: + return 'balanced' + if video_preset in [ 'slow', 'slower', 'veryslow' ]: + return 'quality' + return None + + +def map_qsv_preset(video_preset : VideoPreset) -> Optional[str]: + if video_preset in [ 'ultrafast', 'superfast', 'veryfast' ]: + return 'veryfast' + if video_preset in [ 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow' ]: + return video_preset + return None diff --git a/facefusion/filesystem.py b/facefusion/filesystem.py new file mode 100644 index 0000000000000000000000000000000000000000..2dbd0934ebc39f80a4044f1d3e87065879b7c83c --- /dev/null +++ b/facefusion/filesystem.py @@ -0,0 +1,190 @@ +import glob +import os +import shutil +from typing import List, Optional + +import facefusion.choices + + +def get_file_size(file_path : str) -> int: + if is_file(file_path): + return os.path.getsize(file_path) + return 0 + + +def get_file_name(file_path : str) -> Optional[str]: + file_name, _ = os.path.splitext(os.path.basename(file_path)) + + if file_name: + return file_name + return None + + +def get_file_extension(file_path : str) -> Optional[str]: + _, file_extension = os.path.splitext(file_path) + + if file_extension: + return file_extension.lower() + return None + + +def get_file_format(file_path : str) -> Optional[str]: + file_extension = get_file_extension(file_path) + + if file_extension: + if file_extension == '.jpg': + return 'jpeg' + if file_extension == '.tif': + return 'tiff' + if file_extension == '.mpg': + return 'mpeg' + return file_extension.lstrip('.') + return None + + +def same_file_extension(first_file_path : str, second_file_path : str) -> bool: + first_file_extension = get_file_extension(first_file_path) + second_file_extension = get_file_extension(second_file_path) + + if first_file_extension and second_file_extension: + return get_file_extension(first_file_path) == get_file_extension(second_file_path) + return False + + +def is_file(file_path : str) -> bool: + if file_path: + return os.path.isfile(file_path) + return False + + +def is_audio(audio_path : str) -> bool: + return is_file(audio_path) and get_file_format(audio_path) in facefusion.choices.audio_formats + + +def has_audio(audio_paths : List[str]) -> bool: + if audio_paths: + return any(map(is_audio, audio_paths)) + return False + + +def are_audios(audio_paths : List[str]) -> bool: + if audio_paths: + return all(map(is_audio, audio_paths)) + return False + + +def is_image(image_path : str) -> bool: + return is_file(image_path) and get_file_format(image_path) in facefusion.choices.image_formats + + +def has_image(image_paths : List[str]) -> bool: + if image_paths: + return any(is_image(image_path) for image_path in image_paths) + return False + + +def are_images(image_paths : List[str]) -> bool: + if image_paths: + return all(map(is_image, image_paths)) + return False + + +def is_video(video_path : str) -> bool: + return is_file(video_path) and get_file_format(video_path) in facefusion.choices.video_formats + + +def has_video(video_paths : List[str]) -> bool: + if video_paths: + return any(map(is_video, video_paths)) + return False + + +def are_videos(video_paths : List[str]) -> bool: + if video_paths: + return all(map(is_video, video_paths)) + return False + + +def filter_audio_paths(paths : List[str]) -> List[str]: + if paths: + return [ path for path in paths if is_audio(path) ] + return [] + + +def filter_image_paths(paths : List[str]) -> List[str]: + if paths: + return [ path for path in paths if is_image(path) ] + return [] + + +def copy_file(file_path : str, move_path : str) -> bool: + if is_file(file_path): + shutil.copy(file_path, move_path) + return is_file(move_path) + return False + + +def move_file(file_path : str, move_path : str) -> bool: + if is_file(file_path): + shutil.move(file_path, move_path) + return not is_file(file_path) and is_file(move_path) + return False + + +def remove_file(file_path : str) -> bool: + if is_file(file_path): + os.remove(file_path) + return not is_file(file_path) + return False + + +def resolve_file_paths(directory_path : str) -> List[str]: + file_paths : List[str] = [] + + if is_directory(directory_path): + file_names_and_extensions = sorted(os.listdir(directory_path)) + + for file_name_and_extension in file_names_and_extensions: + if not file_name_and_extension.startswith(('.', '__')): + file_path = os.path.join(directory_path, file_name_and_extension) + file_paths.append(file_path) + + return file_paths + + +def resolve_file_pattern(file_pattern : str) -> List[str]: + if in_directory(file_pattern): + return sorted(glob.glob(file_pattern)) + return [] + + +def is_directory(directory_path : str) -> bool: + if directory_path: + return os.path.isdir(directory_path) + return False + + +def in_directory(file_path : str) -> bool: + if file_path: + directory_path = os.path.dirname(file_path) + if directory_path: + return not is_directory(file_path) and is_directory(directory_path) + return False + + +def create_directory(directory_path : str) -> bool: + if directory_path and not is_file(directory_path): + os.makedirs(directory_path, exist_ok = True) + return is_directory(directory_path) + return False + + +def remove_directory(directory_path : str) -> bool: + if is_directory(directory_path): + shutil.rmtree(directory_path, ignore_errors = True) + return not is_directory(directory_path) + return False + + +def resolve_relative_path(path : str) -> str: + return os.path.abspath(os.path.join(os.path.dirname(__file__), path)) diff --git a/facefusion/hash_helper.py b/facefusion/hash_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..f190fcc3eec9412807dd0b276b22b629d0352190 --- /dev/null +++ b/facefusion/hash_helper.py @@ -0,0 +1,32 @@ +import os +import zlib +from typing import Optional + +from facefusion.filesystem import get_file_name, is_file + + +def create_hash(content : bytes) -> str: + return format(zlib.crc32(content), '08x') + + +def validate_hash(validate_path : str) -> bool: + hash_path = get_hash_path(validate_path) + + if is_file(hash_path): + with open(hash_path) as hash_file: + hash_content = hash_file.read() + + with open(validate_path, 'rb') as validate_file: + validate_content = validate_file.read() + + return create_hash(validate_content) == hash_content + return False + + +def get_hash_path(validate_path : str) -> Optional[str]: + if is_file(validate_path): + validate_directory_path, file_name_and_extension = os.path.split(validate_path) + validate_file_name = get_file_name(file_name_and_extension) + + return os.path.join(validate_directory_path, validate_file_name + '.hash') + return None diff --git a/facefusion/inference_manager.py b/facefusion/inference_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..4213bb0d3fbc759f71d746eee87822756813000c --- /dev/null +++ b/facefusion/inference_manager.py @@ -0,0 +1,95 @@ +import importlib +import random +from time import sleep, time +from typing import List + +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_session_providers, 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 +from facefusion.types import DownloadSet, ExecutionProvider, InferencePool, InferencePoolSet + +INFERENCE_POOL_SET : InferencePoolSet =\ +{ + 'cli': {}, + 'ui': {} +} + + +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 = resolve_execution_providers(module_name) + 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('ui').get(inference_context): + INFERENCE_POOL_SET['cli'][inference_context] = INFERENCE_POOL_SET.get('ui').get(inference_context) + if app_context == 'ui' and INFERENCE_POOL_SET.get('cli').get(inference_context): + INFERENCE_POOL_SET['ui'][inference_context] = INFERENCE_POOL_SET.get('cli').get(inference_context) + if not INFERENCE_POOL_SET.get(app_context).get(inference_context): + INFERENCE_POOL_SET[app_context][inference_context] = create_inference_pool(model_source_set, execution_device_id, execution_providers) + + current_inference_context = get_inference_context(module_name, model_names, random.choice(execution_device_ids), execution_providers) + return INFERENCE_POOL_SET.get(app_context).get(current_inference_context) + + +def create_inference_pool(model_source_set : DownloadSet, execution_device_id : int, execution_providers : List[ExecutionProvider]) -> InferencePool: + inference_pool : InferencePool = {} + + 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, execution_device_id, execution_providers) + + return inference_pool + + +def clear_inference_pool(module_name : str, model_names : List[str]) -> None: + execution_device_ids = state_manager.get_item('execution_device_ids') + execution_providers = resolve_execution_providers(module_name) + app_context = detect_app_context() + + if is_windows() and has_execution_provider('directml'): + INFERENCE_POOL_SET[app_context].clear() + + 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] + + +def create_inference_session(model_path : str, execution_device_id : int, execution_providers : List[ExecutionProvider]) -> InferenceSession: + model_file_name = get_file_name(model_path) + start_time = time() + + try: + inference_session_providers = create_inference_session_providers(execution_device_id, execution_providers) + inference_session = InferenceSession(model_path, providers = inference_session_providers) + logger.debug(translator.get('loading_model_succeeded').format(model_name = model_file_name, seconds = calculate_end_time(start_time)), __name__) + return inference_session + + except Exception: + logger.error(translator.get('loading_model_failed').format(model_name = model_file_name), __name__) + fatal_exit(1) + + +def get_inference_context(module_name : str, model_names : List[str], execution_device_id : int, execution_providers : List[ExecutionProvider]) -> str: + inference_context = '.'.join([ module_name ] + model_names + [ str(execution_device_id) ] + list(execution_providers)) + return inference_context + + +def resolve_execution_providers(module_name : str) -> List[ExecutionProvider]: + module = importlib.import_module(module_name) + + if hasattr(module, 'resolve_execution_providers'): + return getattr(module, 'resolve_execution_providers')() + return state_manager.get_item('execution_providers') diff --git a/facefusion/installer.py b/facefusion/installer.py new file mode 100644 index 0000000000000000000000000000000000000000..09d84e99a1363a867ed4f0b5e982094a6858ec8b --- /dev/null +++ b/facefusion/installer.py @@ -0,0 +1,109 @@ +import os +import shutil +import signal +import subprocess +import sys +from argparse import ArgumentParser, HelpFormatter +from functools import partial +from types import FrameType + +from facefusion import metadata +from facefusion.common_helper import is_linux, is_windows + +LOCALES =\ +{ + 'install_dependency': 'install the {dependency} package', + 'force_reinstall': 'force reinstall of packages', + 'skip_conda': 'skip the conda environment check', + 'conda_not_activated': 'conda is not activated' +} +ONNXRUNTIME_SET =\ +{ + 'default': ('onnxruntime', '1.23.2') +} +if is_windows() or is_linux(): + ONNXRUNTIME_SET['cuda'] = ('onnxruntime-gpu', '1.23.2') + ONNXRUNTIME_SET['openvino'] = ('onnxruntime-openvino', '1.23.0') +if is_windows(): + ONNXRUNTIME_SET['directml'] = ('onnxruntime-directml', '1.23.0') +if is_linux(): + ONNXRUNTIME_SET['migraphx'] = ('onnxruntime-migraphx', '1.23.0') + ONNXRUNTIME_SET['rocm'] = ('onnxruntime_rocm', '1.22.1', '7.0.2') #type:ignore[assignment] + + +def cli() -> None: + signal.signal(signal.SIGINT, signal_exit) + program = ArgumentParser(formatter_class = partial(HelpFormatter, max_help_position = 50)) + program.add_argument('--onnxruntime', help = LOCALES.get('install_dependency').format(dependency = 'onnxruntime'), choices = ONNXRUNTIME_SET.keys(), required = True) + program.add_argument('--force-reinstall', help = LOCALES.get('force_reinstall'), action = 'store_true') + program.add_argument('--skip-conda', help = LOCALES.get('skip_conda'), action = 'store_true') + program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version') + run(program) + + +def signal_exit(signum : int, frame : FrameType) -> None: + sys.exit(0) + + +def run(program : ArgumentParser) -> None: + args = program.parse_args() + has_conda = 'CONDA_PREFIX' in os.environ + commands = [ shutil.which('pip'), 'install' ] + + if args.force_reinstall: + commands.append('--force-reinstall') + + if not args.skip_conda and not has_conda: + sys.stdout.write(LOCALES.get('conda_not_activated') + os.linesep) + sys.exit(1) + + with open('requirements.txt') as file: + + for line in file.readlines(): + __line__ = line.strip() + if not __line__.startswith('onnxruntime'): + commands.append(__line__) + + if args.onnxruntime == 'rocm': + onnxruntime_name, onnxruntime_version, rocm_version = ONNXRUNTIME_SET.get(args.onnxruntime) #type:ignore[misc] + python_id = 'cp' + str(sys.version_info.major) + str(sys.version_info.minor) + + if python_id in [ 'cp310', 'cp312' ]: + wheel_name = onnxruntime_name + '-' + onnxruntime_version + '-' + python_id + '-' + python_id + '-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl' + wheel_url = 'https://repo.radeon.com/rocm/manylinux/rocm-rel-' + rocm_version + '/' + wheel_name + commands.append(wheel_url) + else: + onnxruntime_name, onnxruntime_version = ONNXRUNTIME_SET.get(args.onnxruntime) + commands.append(onnxruntime_name + '==' + onnxruntime_version) + + subprocess.call(commands) + + if args.onnxruntime == 'cuda' and has_conda: + library_paths = [] + + if is_linux(): + if os.getenv('LD_LIBRARY_PATH'): + library_paths = os.getenv('LD_LIBRARY_PATH').split(os.pathsep) + + python_id = 'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor) + library_paths.extend( + [ + os.path.join(os.getenv('CONDA_PREFIX'), 'lib'), + os.path.join(os.getenv('CONDA_PREFIX'), 'lib', python_id, 'site-packages', 'tensorrt_libs') + ]) + library_paths = list(dict.fromkeys([ library_path for library_path in library_paths if os.path.exists(library_path) ])) + + subprocess.call([ shutil.which('conda'), 'env', 'config', 'vars', 'set', 'LD_LIBRARY_PATH=' + os.pathsep.join(library_paths) ]) + + if is_windows(): + if os.getenv('PATH'): + library_paths = os.getenv('PATH').split(os.pathsep) + + library_paths.extend( + [ + os.path.join(os.getenv('CONDA_PREFIX'), 'Lib'), + os.path.join(os.getenv('CONDA_PREFIX'), 'Lib', 'site-packages', 'tensorrt_libs') + ]) + library_paths = list(dict.fromkeys([ library_path for library_path in library_paths if os.path.exists(library_path) ])) + + subprocess.call([ shutil.which('conda'), 'env', 'config', 'vars', 'set', 'PATH=' + os.pathsep.join(library_paths) ]) diff --git a/facefusion/jobs/__init__.py b/facefusion/jobs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/jobs/job_helper.py b/facefusion/jobs/job_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..5337baca87f4433d9057531e6530079221891967 --- /dev/null +++ b/facefusion/jobs/job_helper.py @@ -0,0 +1,20 @@ +import os +from datetime import datetime +from typing import Optional + +from facefusion.filesystem import get_file_extension, get_file_name + + +def get_step_output_path(job_id : str, step_index : int, output_path : str) -> Optional[str]: + if output_path: + output_directory_path, output_file_path = os.path.split(output_path) + output_file_name = get_file_name(output_file_path) + output_file_extension = get_file_extension(output_file_path) + + if output_file_name and output_file_extension: + return os.path.join(output_directory_path, output_file_name + '-' + job_id + '-' + str(step_index) + output_file_extension) + return None + + +def suggest_job_id(job_prefix : str = 'job') -> str: + return job_prefix + '-' + datetime.now().strftime('%Y-%m-%d-%H-%M-%S') diff --git a/facefusion/jobs/job_list.py b/facefusion/jobs/job_list.py new file mode 100644 index 0000000000000000000000000000000000000000..95d21be952ce08aeeba2c8cb9baea81c6c5fcb38 --- /dev/null +++ b/facefusion/jobs/job_list.py @@ -0,0 +1,34 @@ +from datetime import datetime +from typing import List, Optional, Tuple + +from facefusion.jobs import job_manager +from facefusion.time_helper import describe_time_ago +from facefusion.types import JobStatus, TableContent, TableHeader + + +def compose_job_list(job_status : JobStatus) -> Tuple[List[TableHeader], List[List[TableContent]]]: + jobs = job_manager.find_jobs(job_status) + job_headers : List[TableHeader] = [ 'job id', 'steps', 'date created', 'date updated', 'job status' ] + job_contents : List[List[TableContent]] = [] + + for index, job_id in enumerate(jobs): + if job_manager.validate_job(job_id): + job = jobs[job_id] + step_total = job_manager.count_step_total(job_id) + date_created = prepare_describe_datetime(job.get('date_created')) + date_updated = prepare_describe_datetime(job.get('date_updated')) + job_contents.append( + [ + job_id, + step_total, + date_created, + date_updated, + job_status + ]) + return job_headers, job_contents + + +def prepare_describe_datetime(date_time : Optional[str]) -> Optional[str]: + if date_time: + return describe_time_ago(datetime.fromisoformat(date_time)) + return None diff --git a/facefusion/jobs/job_manager.py b/facefusion/jobs/job_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..a1aec04e48b844ac145ca9733b4b1e4c19aee678 --- /dev/null +++ b/facefusion/jobs/job_manager.py @@ -0,0 +1,265 @@ +import os +from copy import copy +from typing import List, Optional + +import facefusion.choices +from facefusion.filesystem import create_directory, get_file_name, is_directory, is_file, move_file, remove_directory, remove_file, resolve_file_pattern +from facefusion.jobs.job_helper import get_step_output_path +from facefusion.json import read_json, write_json +from facefusion.time_helper import get_current_date_time +from facefusion.types import Args, Job, JobSet, JobStatus, JobStep, JobStepStatus + +JOBS_PATH : Optional[str] = None + + +def init_jobs(jobs_path : str) -> bool: + global JOBS_PATH + + JOBS_PATH = jobs_path + job_status_paths = [ os.path.join(JOBS_PATH, job_status) for job_status in facefusion.choices.job_statuses ] + + for job_status_path in job_status_paths: + create_directory(job_status_path) + return all(is_directory(status_path) for status_path in job_status_paths) + + +def clear_jobs(jobs_path : str) -> bool: + return remove_directory(jobs_path) + + +def create_job(job_id : str) -> bool: + job : Job =\ + { + 'version': '1', + 'date_created': get_current_date_time().isoformat(), + 'date_updated': None, + 'steps': [] + } + + return create_job_file(job_id, job) + + +def submit_job(job_id : str) -> bool: + drafted_job_ids = find_job_ids('drafted') + steps = get_steps(job_id) + + if job_id in drafted_job_ids and steps: + return set_steps_status(job_id, 'queued') and move_job_file(job_id, 'queued') + return False + + +def submit_jobs(halt_on_error : bool) -> bool: + drafted_job_ids = find_job_ids('drafted') + has_error = False + + if drafted_job_ids: + for job_id in drafted_job_ids: + if not submit_job(job_id): + has_error = True + if halt_on_error: + return False + return not has_error + return False + + +def delete_job(job_id : str) -> bool: + return delete_job_file(job_id) + + +def delete_jobs(halt_on_error : bool) -> bool: + job_ids = find_job_ids('drafted') + find_job_ids('queued') + find_job_ids('failed') + find_job_ids('completed') + has_error = False + + if job_ids: + for job_id in job_ids: + if not delete_job(job_id): + has_error = True + if halt_on_error: + return False + return not has_error + return False + + +def find_jobs(job_status : JobStatus) -> JobSet: + job_ids = find_job_ids(job_status) + job_set : JobSet = {} + + for job_id in job_ids: + job_set[job_id] = read_job_file(job_id) + return job_set + + +def find_job_ids(job_status : JobStatus) -> List[str]: + job_pattern = os.path.join(JOBS_PATH, job_status, '*.json') + job_paths = resolve_file_pattern(job_pattern) + job_paths.sort(key = os.path.getmtime) + job_ids = [] + + for job_path in job_paths: + job_id = get_file_name(job_path) + job_ids.append(job_id) + return job_ids + + +def validate_job(job_id : str) -> bool: + job = read_job_file(job_id) + return bool(job and 'version' in job and 'date_created' in job and 'date_updated' in job and 'steps' in job) + + +def has_step(job_id : str, step_index : int) -> bool: + step_total = count_step_total(job_id) + return step_index in range(step_total) + + +def add_step(job_id : str, step_args : Args) -> bool: + job = read_job_file(job_id) + + if job: + job.get('steps').append( + { + 'args': step_args, + 'status': 'drafted' + }) + return update_job_file(job_id, job) + return False + + +def remix_step(job_id : str, step_index : int, step_args : Args) -> bool: + steps = get_steps(job_id) + step_args = copy(step_args) + + if step_index and step_index < 0: + step_index = count_step_total(job_id) - 1 + + if has_step(job_id, step_index): + output_path = steps[step_index].get('args').get('output_path') + step_args['target_path'] = get_step_output_path(job_id, step_index, output_path) + return add_step(job_id, step_args) + return False + + +def insert_step(job_id : str, step_index : int, step_args : Args) -> bool: + job = read_job_file(job_id) + step_args = copy(step_args) + + if step_index and step_index < 0: + step_index = count_step_total(job_id) - 1 + + if job and has_step(job_id, step_index): + job.get('steps').insert(step_index, + { + 'args': step_args, + 'status': 'drafted' + }) + return update_job_file(job_id, job) + return False + + +def remove_step(job_id : str, step_index : int) -> bool: + job = read_job_file(job_id) + + if step_index and step_index < 0: + step_index = count_step_total(job_id) - 1 + + if job and has_step(job_id, step_index): + job.get('steps').pop(step_index) + return update_job_file(job_id, job) + return False + + +def get_steps(job_id : str) -> List[JobStep]: + job = read_job_file(job_id) + + if job: + return job.get('steps') + return [] + + +def count_step_total(job_id : str) -> int: + steps = get_steps(job_id) + + if steps: + return len(steps) + return 0 + + +def set_step_status(job_id : str, step_index : int, step_status : JobStepStatus) -> bool: + job = read_job_file(job_id) + + if job: + steps = job.get('steps') + if has_step(job_id, step_index): + steps[step_index]['status'] = step_status + return update_job_file(job_id, job) + return False + + +def set_steps_status(job_id : str, step_status : JobStepStatus) -> bool: + job = read_job_file(job_id) + + if job: + for step in job.get('steps'): + step['status'] = step_status + return update_job_file(job_id, job) + return False + + +def read_job_file(job_id : str) -> Optional[Job]: + job_path = find_job_path(job_id) + return read_json(job_path) #type:ignore[return-value] + + +def create_job_file(job_id : str, job : Job) -> bool: + job_path = find_job_path(job_id) + + if not is_file(job_path): + job_create_path = suggest_job_path(job_id, 'drafted') + return write_json(job_create_path, job) #type:ignore[arg-type] + return False + + +def update_job_file(job_id : str, job : Job) -> bool: + job_path = find_job_path(job_id) + + if is_file(job_path): + job['date_updated'] = get_current_date_time().isoformat() + return write_json(job_path, job) #type:ignore[arg-type] + return False + + +def move_job_file(job_id : str, job_status : JobStatus) -> bool: + job_path = find_job_path(job_id) + job_move_path = suggest_job_path(job_id, job_status) + return move_file(job_path, job_move_path) + + +def delete_job_file(job_id : str) -> bool: + job_path = find_job_path(job_id) + return remove_file(job_path) + + +def suggest_job_path(job_id : str, job_status : JobStatus) -> Optional[str]: + job_file_name = get_job_file_name(job_id) + + if job_file_name: + return os.path.join(JOBS_PATH, job_status, job_file_name) + return None + + +def find_job_path(job_id : str) -> Optional[str]: + job_file_name = get_job_file_name(job_id) + + if job_file_name: + for job_status in facefusion.choices.job_statuses: + job_pattern = os.path.join(JOBS_PATH, job_status, job_file_name) + job_paths = resolve_file_pattern(job_pattern) + + for job_path in job_paths: + return job_path + return None + + +def get_job_file_name(job_id : str) -> Optional[str]: + if job_id: + return job_id + '.json' + return None diff --git a/facefusion/jobs/job_runner.py b/facefusion/jobs/job_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..ee6be8c78ba90b3747e39acb42ca8d3175f723b1 --- /dev/null +++ b/facefusion/jobs/job_runner.py @@ -0,0 +1,112 @@ +from facefusion.ffmpeg import concat_video +from facefusion.filesystem import are_images, are_videos, move_file, remove_file +from facefusion.jobs import job_helper, job_manager +from facefusion.types import JobOutputSet, JobStep, ProcessStep + + +def run_job(job_id : str, process_step : ProcessStep) -> bool: + queued_job_ids = job_manager.find_job_ids('queued') + + if job_id in queued_job_ids: + if run_steps(job_id, process_step) and finalize_steps(job_id): + clean_steps(job_id) + return job_manager.move_job_file(job_id, 'completed') + clean_steps(job_id) + job_manager.move_job_file(job_id, 'failed') + return False + + +def run_jobs(process_step : ProcessStep, halt_on_error : bool) -> bool: + queued_job_ids = job_manager.find_job_ids('queued') + has_error = False + + if queued_job_ids: + for job_id in queued_job_ids: + if not run_job(job_id, process_step): + has_error = True + if halt_on_error: + return False + return not has_error + return False + + +def retry_job(job_id : str, process_step : ProcessStep) -> bool: + failed_job_ids = job_manager.find_job_ids('failed') + + if job_id in failed_job_ids: + return job_manager.set_steps_status(job_id, 'queued') and job_manager.move_job_file(job_id, 'queued') and run_job(job_id, process_step) + return False + + +def retry_jobs(process_step : ProcessStep, halt_on_error : bool) -> bool: + failed_job_ids = job_manager.find_job_ids('failed') + has_error = False + + if failed_job_ids: + for job_id in failed_job_ids: + if not retry_job(job_id, process_step): + has_error = True + if halt_on_error: + return False + return not has_error + return False + + +def run_step(job_id : str, step_index : int, step : JobStep, process_step : ProcessStep) -> bool: + step_args = step.get('args') + + if job_manager.set_step_status(job_id, step_index, 'started') and process_step(job_id, step_index, step_args): + output_path = step_args.get('output_path') + step_output_path = job_helper.get_step_output_path(job_id, step_index, output_path) + + return move_file(output_path, step_output_path) and job_manager.set_step_status(job_id, step_index, 'completed') + job_manager.set_step_status(job_id, step_index, 'failed') + return False + + +def run_steps(job_id : str, process_step : ProcessStep) -> bool: + steps = job_manager.get_steps(job_id) + + if steps: + for index, step in enumerate(steps): + if not run_step(job_id, index, step, process_step): + return False + return True + return False + + +def finalize_steps(job_id : str) -> bool: + output_set = collect_output_set(job_id) + + for output_path, temp_output_paths in output_set.items(): + if are_videos(temp_output_paths): + if not concat_video(output_path, temp_output_paths): + return False + if are_images(temp_output_paths): + for temp_output_path in temp_output_paths: + if not move_file(temp_output_path, output_path): + return False + return True + + +def clean_steps(job_id: str) -> bool: + output_set = collect_output_set(job_id) + + for temp_output_paths in output_set.values(): + for temp_output_path in temp_output_paths: + if not remove_file(temp_output_path): + return False + return True + + +def collect_output_set(job_id : str) -> JobOutputSet: + steps = job_manager.get_steps(job_id) + job_output_set : JobOutputSet = {} + + for index, step in enumerate(steps): + output_path = step.get('args').get('output_path') + + if output_path: + step_output_path = job_manager.get_step_output_path(job_id, index, output_path) + job_output_set.setdefault(output_path, []).append(step_output_path) + return job_output_set diff --git a/facefusion/jobs/job_store.py b/facefusion/jobs/job_store.py new file mode 100644 index 0000000000000000000000000000000000000000..e44c091e111b52d9d34d68805f4def86b07a52b2 --- /dev/null +++ b/facefusion/jobs/job_store.py @@ -0,0 +1,27 @@ +from typing import List + +from facefusion.types import JobStore + +JOB_STORE : JobStore =\ +{ + 'job_keys': [], + 'step_keys': [] +} + + +def get_job_keys() -> List[str]: + return JOB_STORE.get('job_keys') + + +def get_step_keys() -> List[str]: + return JOB_STORE.get('step_keys') + + +def register_job_keys(job_keys : List[str]) -> None: + for job_key in job_keys: + JOB_STORE['job_keys'].append(job_key) + + +def register_step_keys(step_keys : List[str]) -> None: + for step_key in step_keys: + JOB_STORE['step_keys'].append(step_key) diff --git a/facefusion/json.py b/facefusion/json.py new file mode 100644 index 0000000000000000000000000000000000000000..944df65c3370110b070ceada4c519b712fd84ccf --- /dev/null +++ b/facefusion/json.py @@ -0,0 +1,22 @@ +import json +from json import JSONDecodeError +from typing import Optional + +from facefusion.filesystem import is_file +from facefusion.types import Content + + +def read_json(json_path : str) -> Optional[Content]: + if is_file(json_path): + try: + with open(json_path) as json_file: + return json.load(json_file) + except JSONDecodeError: + pass + return None + + +def write_json(json_path : str, content : Content) -> bool: + with open(json_path, 'w') as json_file: + json.dump(content, json_file, indent = 4) + return is_file(json_path) diff --git a/facefusion/locales.py b/facefusion/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..0fdce05dcf19a0baad3c95e28676ad04ed1bd62c --- /dev/null +++ b/facefusion/locales.py @@ -0,0 +1,274 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + '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', + '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', + 'extracting_frames_failed': 'extracting frames failed', + 'analysing': 'analysing', + 'extracting': 'extracting', + 'streaming': 'streaming', + 'processing': 'processing', + 'merging': 'merging', + 'downloading': 'downloading', + 'temp_frames_not_found': 'temporary frames not found', + 'copying_image': 'copying image with a resolution of {resolution}', + 'copying_image_succeeded': 'copying image succeeded', + 'copying_image_failed': 'copying image failed', + 'finalizing_image': 'finalizing image with a resolution of {resolution}', + 'finalizing_image_succeeded': 'finalizing image succeeded', + 'finalizing_image_skipped': 'finalizing image skipped', + 'merging_video': 'merging video with a resolution of {resolution} and {fps} frames per second', + 'merging_video_succeeded': 'merging video succeeded', + 'merging_video_failed': 'merging video failed', + 'skipping_audio': 'skipping audio', + 'replacing_audio_succeeded': 'replacing audio succeeded', + 'replacing_audio_skipped': 'replacing audio skipped', + 'restoring_audio_succeeded': 'restoring audio succeeded', + 'restoring_audio_skipped': 'restoring audio skipped', + 'clearing_temp': 'clearing temporary resources', + 'processing_stopped': 'processing stopped', + 'processing_image_succeeded': 'processing to image succeeded in {seconds} seconds', + 'processing_image_failed': 'processing to image failed', + 'processing_video_succeeded': 'processing to video succeeded in {seconds} seconds', + 'processing_video_failed': 'processing to video failed', + 'choose_image_source': 'choose an image for the source', + 'choose_audio_source': 'choose an audio for the source', + 'choose_video_target': 'choose a video for the target', + 'choose_image_or_video_target': 'choose an image or video for the target', + 'specify_image_or_video_output': 'specify the output image or video within a directory', + 'match_target_and_output_extension': 'match the target and output extension', + 'no_source_face_detected': 'no source face detected', + 'processor_not_loaded': 'processor {processor} could not be loaded', + 'processor_not_implemented': 'processor {processor} not implemented correctly', + 'ui_layout_not_loaded': 'ui layout {ui_layout} could not be loaded', + 'ui_layout_not_implemented': 'ui layout {ui_layout} not implemented correctly', + 'stream_not_loaded': 'stream {stream_mode} could not be loaded', + 'stream_not_supported': 'stream not supported', + 'job_created': 'job {job_id} created', + 'job_not_created': 'job {job_id} not created', + 'job_submitted': 'job {job_id} submitted', + 'job_not_submitted': 'job {job_id} not submitted', + 'job_all_submitted': 'jobs submitted', + 'job_all_not_submitted': 'jobs not submitted', + 'job_deleted': 'job {job_id} deleted', + 'job_not_deleted': 'job {job_id} not deleted', + 'job_all_deleted': 'jobs deleted', + 'job_all_not_deleted': 'jobs not deleted', + 'job_step_added': 'step added to job {job_id}', + 'job_step_not_added': 'step not added to job {job_id}', + 'job_remix_step_added': 'step {step_index} remixed from job {job_id}', + 'job_remix_step_not_added': 'step {step_index} not remixed from job {job_id}', + 'job_step_inserted': 'step {step_index} inserted to job {job_id}', + 'job_step_not_inserted': 'step {step_index} not inserted to job {job_id}', + 'job_step_removed': 'step {step_index} removed from job {job_id}', + 'job_step_not_removed': 'step {step_index} not removed from job {job_id}', + 'running_job': 'running queued job {job_id}', + 'running_jobs': 'running all queued jobs', + 'retrying_job': 'retrying failed job {job_id}', + 'retrying_jobs': 'retrying all failed jobs', + 'processing_job_succeeded': 'processing of job {job_id} succeeded', + 'processing_jobs_succeeded': 'processing of all jobs succeeded', + 'processing_job_failed': 'processing of job {job_id} failed', + 'processing_jobs_failed': 'processing of all jobs failed', + 'processing_step': 'processing step {step_current} of {step_total}', + 'validating_hash_succeeded': 'validating hash for {hash_file_name} succeeded', + 'validating_hash_failed': 'validating hash for {hash_file_name} failed', + 'validating_source_succeeded': 'validating source for {source_file_name} succeeded', + 'validating_source_failed': 'validating source for {source_file_name} failed', + 'deleting_corrupt_source': 'deleting corrupt source for {source_file_name}', + 'loading_model_succeeded': 'loading model {model_name} succeeded in {seconds} seconds', + 'loading_model_failed': 'loading model {model_name} failed', + 'time_ago_now': 'just now', + 'time_ago_minutes': '{minutes} minutes ago', + 'time_ago_hours': '{hours} hours and {minutes} minutes ago', + 'time_ago_days': '{days} days, {hours} hours and {minutes} minutes ago', + 'point': '.', + 'comma': ',', + 'colon': ':', + 'question_mark': '?', + 'exclamation_mark': '!', + 'help': + { + 'install_dependency': 'choose the variant of {dependency} to install', + 'skip_conda': 'skip the conda environment check', + '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', + 'source_paths': 'choose the image or audio paths', + 'target_path': 'choose the image or video path', + 'output_path': 'specify the image or video within a directory', + 'source_pattern': 'choose the image or audio pattern', + 'target_pattern': 'choose the image or video pattern', + 'output_pattern': 'specify the image or video pattern', + 'face_detector_model': 'choose the model responsible for detecting the faces', + 'face_detector_size': 'specify the frame size provided to the face detector', + 'face_detector_margin': 'apply top, right, bottom and left margin to the frame', + 'face_detector_angles': 'specify the angles to rotate the frame before detecting faces', + 'face_detector_score': 'filter the detected faces based on the confidence score', + 'face_landmarker_model': 'choose the model responsible for detecting the face landmarks', + 'face_landmarker_score': 'filter the detected face landmarks based on the confidence score', + 'face_selector_mode': 'use reference based tracking or simple matching', + 'face_selector_order': 'specify the order of the detected faces', + 'face_selector_age_start': 'filter the detected faces based on the starting age', + 'face_selector_age_end': 'filter the detected faces based on the ending age', + 'face_selector_gender': 'filter the detected faces based on their gender', + 'face_selector_race': 'filter the detected faces based on their race', + 'reference_face_position': 'specify the position used to create the reference face', + 'reference_face_distance': 'specify the similarity between the reference face and target face', + 'reference_frame_number': 'specify the frame used to create the reference face', + 'face_occluder_model': 'choose the model responsible for the occlusion mask', + 'face_parser_model': 'choose the model responsible for the region mask', + 'face_mask_types': 'mix and match different face mask types (choices: {choices})', + 'face_mask_areas': 'choose the items used for the area mask (choices: {choices})', + 'face_mask_regions': 'choose the items used for the region mask (choices: {choices})', + 'face_mask_blur': 'specify the degree of blur applied to the box mask', + 'face_mask_padding': 'apply top, right, bottom and left padding to the box mask', + 'voice_extractor_model': 'choose the model responsible for extracting the voices', + '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', + 'keep_temp': 'keep the temporary resources after processing', + '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', + 'output_audio_encoder': 'specify the encoder used for the audio', + 'output_audio_quality': 'specify the audio quality which translates to the audio compression', + 'output_audio_volume': 'specify the audio volume based on the target video', + 'output_video_encoder': 'specify the encoder used for the video', + 'output_video_preset': 'balance fast video processing and video file size', + 'output_video_quality': 'specify the video quality which translates to the video compression', + 'output_video_scale': 'specify the video scale based on the target video', + 'output_video_fps': 'specify the video fps based on the target video', + 'processors': 'load a single or multiple processors (choices: {choices}, ...)', + 'background-remover-model': 'choose the model responsible for removing the background', + 'background-remover-color': 'apply red, green blue and alpha values of the background', + 'open_browser': 'open the browser once the program is ready', + 'ui_layouts': 'launch a single or multiple UI layouts (choices: {choices}, ...)', + 'ui_workflow': 'choose the ui workflow', + 'download_providers': 'download using different providers (choices: {choices}, ...)', + 'download_scope': 'specify the download scope', + 'benchmark_mode': 'choose the benchmark mode', + 'benchmark_resolutions': 'choose the resolutions for the benchmarks (choices: {choices}, ...)', + 'benchmark_cycle_count': 'specify the amount of cycles per benchmark', + 'execution_device_ids': 'specify the devices used for processing', + 'execution_providers': 'inference using different providers (choices: {choices}, ...)', + 'execution_thread_count': 'specify the amount of parallel threads while processing', + 'video_memory_strategy': 'balance fast processing and low VRAM usage', + 'system_memory_limit': 'limit the available RAM that can be used while processing', + 'log_level': 'adjust the message severity displayed in the terminal', + 'halt_on_error': 'halt the program once an error occurred', + 'run': 'run the program', + 'headless_run': 'run the program in headless mode', + 'batch_run': 'run the program in batch mode', + 'force_download': 'force automate downloads and exit', + 'benchmark': 'benchmark the program', + 'job_id': 'specify the job id', + 'job_status': 'specify the job status', + 'step_index': 'specify the step index', + 'job_list': 'list jobs by status', + 'job_create': 'create a drafted job', + 'job_submit': 'submit a drafted job to become a queued job', + 'job_submit_all': 'submit all drafted jobs to become a queued jobs', + 'job_delete': 'delete a drafted, queued, failed or completed job', + 'job_delete_all': 'delete all drafted, queued, failed and completed jobs', + 'job_add_step': 'add a step to a drafted job', + 'job_remix_step': 'remix a previous step from a drafted job', + 'job_insert_step': 'insert a step to a drafted job', + 'job_remove_step': 'remove a step from a drafted job', + 'job_run': 'run a queued job', + 'job_run_all': 'run all queued jobs', + 'job_retry': 'retry a failed job', + 'job_retry_all': 'retry all failed jobs' + }, + 'about': + { + 'fund': 'fund training server', + 'subscribe': 'become a member', + 'join': 'join our community' + }, + 'uis': + { + 'apply_button': 'APPLY', + 'benchmark_mode_dropdown': 'BENCHMARK MODE', + 'benchmark_cycle_count_slider': 'BENCHMARK CYCLE COUNT', + 'benchmark_resolutions_checkbox_group': 'BENCHMARK RESOLUTIONS', + 'clear_button': 'CLEAR', + 'common_options_checkbox_group': 'OPTIONS', + 'download_providers_checkbox_group': 'DOWNLOAD PROVIDERS', + 'execution_providers_checkbox_group': 'EXECUTION PROVIDERS', + 'execution_thread_count_slider': 'EXECUTION THREAD COUNT', + 'face_detector_angles_checkbox_group': 'FACE DETECTOR ANGLES', + 'face_detector_model_dropdown': 'FACE DETECTOR MODEL', + 'face_detector_margin_slider': 'FACE DETECTOR MARGIN', + 'face_detector_score_slider': 'FACE DETECTOR SCORE', + 'face_detector_size_dropdown': 'FACE DETECTOR SIZE', + 'face_landmarker_model_dropdown': 'FACE LANDMARKER MODEL', + 'face_landmarker_score_slider': 'FACE LANDMARKER SCORE', + 'face_mask_blur_slider': 'FACE MASK BLUR', + 'face_mask_padding_bottom_slider': 'FACE MASK PADDING BOTTOM', + 'face_mask_padding_left_slider': 'FACE MASK PADDING LEFT', + 'face_mask_padding_right_slider': 'FACE MASK PADDING RIGHT', + 'face_mask_padding_top_slider': 'FACE MASK PADDING TOP', + 'face_mask_areas_checkbox_group': 'FACE MASK AREAS', + 'face_mask_regions_checkbox_group': 'FACE MASK REGIONS', + 'face_mask_types_checkbox_group': 'FACE MASK TYPES', + 'face_selector_age_range_slider': 'FACE SELECTOR AGE', + 'face_selector_gender_dropdown': 'FACE SELECTOR GENDER', + 'face_selector_mode_dropdown': 'FACE SELECTOR MODE', + 'face_selector_order_dropdown': 'FACE SELECTOR ORDER', + 'face_selector_race_dropdown': 'FACE SELECTOR RACE', + 'face_occluder_model_dropdown': 'FACE OCCLUDER MODEL', + 'face_parser_model_dropdown': 'FACE PARSER MODEL', + 'voice_extractor_model_dropdown': 'VOICE EXTRACTOR MODEL', + 'job_list_status_checkbox_group': 'JOB STATUS', + 'job_manager_job_action_dropdown': 'JOB_ACTION', + 'job_manager_job_id_dropdown': 'JOB ID', + 'job_manager_step_index_dropdown': 'STEP INDEX', + 'job_runner_job_action_dropdown': 'JOB ACTION', + 'job_runner_job_id_dropdown': 'JOB ID', + 'log_level_dropdown': 'LOG LEVEL', + 'output_audio_encoder_dropdown': 'OUTPUT AUDIO ENCODER', + 'output_audio_quality_slider': 'OUTPUT AUDIO QUALITY', + 'output_audio_volume_slider': 'OUTPUT AUDIO VOLUME', + 'output_image_or_video': 'OUTPUT', + 'output_image_quality_slider': 'OUTPUT IMAGE QUALITY', + 'output_image_scale_slider': 'OUTPUT IMAGE SCALE', + 'output_path_textbox': 'OUTPUT PATH', + 'output_video_encoder_dropdown': 'OUTPUT VIDEO ENCODER', + 'output_video_fps_slider': 'OUTPUT VIDEO FPS', + 'output_video_preset_dropdown': 'OUTPUT VIDEO PRESET', + 'output_video_quality_slider': 'OUTPUT VIDEO QUALITY', + 'output_video_scale_slider': 'OUTPUT VIDEO SCALE', + 'preview_frame_slider': 'PREVIEW FRAME', + 'preview_image': 'PREVIEW', + 'preview_mode_dropdown': 'PREVIEW MODE', + 'preview_resolution_dropdown': 'PREVIEW RESOLUTION', + 'processors_checkbox_group': 'PROCESSORS', + 'reference_face_distance_slider': 'REFERENCE FACE DISTANCE', + 'reference_face_gallery': 'REFERENCE FACE', + 'refresh_button': 'REFRESH', + 'source_file': 'SOURCE', + 'start_button': 'START', + 'stop_button': 'STOP', + 'system_memory_limit_slider': 'SYSTEM MEMORY LIMIT', + 'target_file': 'TARGET', + 'temp_frame_format_dropdown': 'TEMP FRAME FORMAT', + 'terminal_textbox': 'TERMINAL', + 'trim_frame_slider': 'TRIM FRAME', + 'ui_workflow': 'UI WORKFLOW', + 'video_memory_strategy_dropdown': 'VIDEO MEMORY STRATEGY', + 'webcam_fps_slider': 'WEBCAM FPS', + 'webcam_image': 'WEBCAM', + 'webcam_device_id_dropdown': 'WEBCAM DEVICE ID', + 'webcam_mode_radio': 'WEBCAM MODE', + 'webcam_resolution_dropdown': 'WEBCAM RESOLUTION' + } + } +} diff --git a/facefusion/logger.py b/facefusion/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..92e6663c1b2ced0ee5289e2315b3391479348a0d --- /dev/null +++ b/facefusion/logger.py @@ -0,0 +1,48 @@ +from logging import Logger, basicConfig, getLogger + +import facefusion.choices +from facefusion.common_helper import get_first, get_last +from facefusion.types import LogLevel + + +def init(log_level : LogLevel) -> None: + basicConfig(format = '%(message)s') + get_package_logger().setLevel(facefusion.choices.log_level_set.get(log_level)) + + +def get_package_logger() -> Logger: + return getLogger('facefusion') + + +def debug(message : str, module_name : str) -> None: + get_package_logger().debug(create_message(message, module_name)) + + +def info(message : str, module_name : str) -> None: + get_package_logger().info(create_message(message, module_name)) + + +def warn(message : str, module_name : str) -> None: + get_package_logger().warning(create_message(message, module_name)) + + +def error(message : str, module_name : str) -> None: + get_package_logger().error(create_message(message, module_name)) + + +def create_message(message : str, module_name : str) -> str: + module_names = module_name.split('.') + first_module_name = get_first(module_names) + last_module_name = get_last(module_names) + + if first_module_name and last_module_name: + return '[' + first_module_name.upper() + '.' + last_module_name.upper() + '] ' + message + return message + + +def enable() -> None: + get_package_logger().disabled = False + + +def disable() -> None: + get_package_logger().disabled = True diff --git a/facefusion/memory.py b/facefusion/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..dacfebc7cf19f3751caa4c56d7365bba7d52cec2 --- /dev/null +++ b/facefusion/memory.py @@ -0,0 +1,21 @@ +from facefusion.common_helper import is_macos, is_windows + +if is_windows(): + import ctypes +else: + import resource + + +def limit_system_memory(system_memory_limit : int = 1) -> bool: + if is_macos(): + system_memory_limit = system_memory_limit * (1024 ** 6) + else: + system_memory_limit = system_memory_limit * (1024 ** 3) + try: + if is_windows(): + ctypes.windll.kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(system_memory_limit), ctypes.c_size_t(system_memory_limit)) #type:ignore[attr-defined] + else: + resource.setrlimit(resource.RLIMIT_DATA, (system_memory_limit, system_memory_limit)) + return True + except Exception: + return False diff --git a/facefusion/metadata.py b/facefusion/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..644c750d26b529152d5419a6a192cd814b77ecf7 --- /dev/null +++ b/facefusion/metadata.py @@ -0,0 +1,15 @@ +from typing import Optional + +METADATA =\ +{ + 'name': 'FaceFusion', + 'description': 'Industry leading face manipulation platform', + 'version': '3.5.2', + 'license': 'OpenRAIL-AS', + 'author': 'Henry Ruhs', + 'url': 'https://facefusion.io' +} + + +def get(key : str) -> Optional[str]: + return METADATA.get(key) diff --git a/facefusion/model_helper.py b/facefusion/model_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..22b3d0d776615bd73f87e7ba03a1698552ccd382 --- /dev/null +++ b/facefusion/model_helper.py @@ -0,0 +1,11 @@ +from functools import lru_cache + +import onnx + +from facefusion.types import ModelInitializer + + +@lru_cache() +def get_static_model_initializer(model_path : str) -> ModelInitializer: + model = onnx.load(model_path) + return onnx.numpy_helper.to_array(model.graph.initializer[-1]) diff --git a/facefusion/normalizer.py b/facefusion/normalizer.py new file mode 100644 index 0000000000000000000000000000000000000000..5c349efdc289ce4fc0156ccdb99f43d8916a7c37 --- /dev/null +++ b/facefusion/normalizer.py @@ -0,0 +1,33 @@ +from typing import List, Optional + +from facefusion.types import Color, Fps, Padding + + +def normalize_color(channels : Optional[List[int]]) -> Optional[Color]: + if channels and len(channels) == 1: + return tuple([ channels[0], channels[0], channels[0], 255 ]) #type:ignore[return-value] + if channels and len(channels) == 2: + return tuple([ channels[0], channels[1], channels[0], 255 ]) #type:ignore[return-value] + if channels and len(channels) == 3: + return tuple([ channels[0], channels[1], channels[2], 255 ]) #type:ignore[return-value] + if channels and len(channels) == 4: + return tuple(channels) #type:ignore[return-value] + return None + + +def normalize_space(spaces : Optional[List[int]]) -> Optional[Padding]: + if spaces and len(spaces) == 1: + return tuple([spaces[0]] * 4) #type:ignore[return-value] + if spaces and len(spaces) == 2: + return tuple([spaces[0], spaces[1], spaces[0], spaces[1]]) #type:ignore[return-value] + if spaces and len(spaces) == 3: + return tuple([spaces[0], spaces[1], spaces[2], spaces[1]]) #type:ignore[return-value] + if spaces and len(spaces) == 4: + return tuple(spaces) #type:ignore[return-value] + return None + + +def normalize_fps(fps : Optional[float]) -> Optional[Fps]: + if isinstance(fps, (int, float)): + return max(1.0, min(fps, 60.0)) + return None diff --git a/facefusion/process_manager.py b/facefusion/process_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..f2da2cf4f0ac09fcccbe6c92cb0fa70e708b5dcc --- /dev/null +++ b/facefusion/process_manager.py @@ -0,0 +1,45 @@ +from facefusion.types import ProcessState + +PROCESS_STATE : ProcessState = 'pending' + + +def get_process_state() -> ProcessState: + return PROCESS_STATE + + +def set_process_state(process_state : ProcessState) -> None: + global PROCESS_STATE + + PROCESS_STATE = process_state + + +def is_checking() -> bool: + return get_process_state() == 'checking' + + +def is_processing() -> bool: + return get_process_state() == 'processing' + + +def is_stopping() -> bool: + return get_process_state() == 'stopping' + + +def is_pending() -> bool: + return get_process_state() == 'pending' + + +def check() -> None: + set_process_state('checking') + + +def start() -> None: + set_process_state('processing') + + +def stop() -> None: + set_process_state('stopping') + + +def end() -> None: + set_process_state('pending') diff --git a/facefusion/processors/__init__.py b/facefusion/processors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/core.py b/facefusion/processors/core.py new file mode 100644 index 0000000000000000000000000000000000000000..611bd3d5504fc90da4bfded413a04ecac02f5862 --- /dev/null +++ b/facefusion/processors/core.py @@ -0,0 +1,44 @@ +import importlib +from types import ModuleType +from typing import Any, List + +from facefusion import logger, translator +from facefusion.exit_helper import hard_exit + + +PROCESSORS_METHODS =\ +[ + 'get_inference_pool', + 'clear_inference_pool', + 'register_args', + 'apply_args', + 'pre_check', + 'pre_process', + 'post_process', + 'process_frame' +] + + +def load_processor_module(processor : str) -> Any: + try: + processor_module = importlib.import_module('facefusion.processors.modules.' + processor + '.core') + for method_name in PROCESSORS_METHODS: + if not hasattr(processor_module, method_name): + raise NotImplementedError + except ModuleNotFoundError as exception: + logger.error(translator.get('processor_not_loaded').format(processor = processor), __name__) + logger.debug(exception.msg, __name__) + hard_exit(1) + except NotImplementedError: + logger.error(translator.get('processor_not_implemented').format(processor = processor), __name__) + hard_exit(1) + return processor_module + + +def get_processors_modules(processors : List[str]) -> List[ModuleType]: + processor_modules = [] + + for processor in processors: + processor_module = load_processor_module(processor) + processor_modules.append(processor_module) + return processor_modules diff --git a/facefusion/processors/live_portrait.py b/facefusion/processors/live_portrait.py new file mode 100644 index 0000000000000000000000000000000000000000..7c94f01f9e5f0a73b8018f20210a0ab5771d6396 --- /dev/null +++ b/facefusion/processors/live_portrait.py @@ -0,0 +1,101 @@ +from typing import Tuple + +import numpy +import scipy + +from facefusion.processors.types import LivePortraitExpression, LivePortraitPitch, LivePortraitRoll, LivePortraitRotation, LivePortraitYaw + +EXPRESSION_MIN = numpy.array( +[ + [ + [ -2.88067125e-02, -8.12731311e-02, -1.70541159e-03 ], + [ -4.88598682e-02, -3.32196616e-02, -1.67431499e-04 ], + [ -6.75425082e-02, -4.28681746e-02, -1.98950816e-04 ], + [ -7.23103955e-02, -3.28503326e-02, -7.31324719e-04 ], + [ -3.87073644e-02, -6.01546466e-02, -5.50269964e-04 ], + [ -6.38048723e-02, -2.23840728e-01, -7.13261834e-04 ], + [ -3.02710701e-02, -3.93195450e-02, -8.24086510e-06 ], + [ -2.95799859e-02, -5.39318882e-02, -1.74219604e-04 ], + [ -2.92359516e-02, -1.53050944e-02, -6.30460854e-05 ], + [ -5.56493877e-03, -2.34344602e-02, -1.26858242e-04 ], + [ -4.37593013e-02, -2.77768299e-02, -2.70503685e-02 ], + [ -1.76926646e-02, -1.91676542e-02, -1.15090821e-04 ], + [ -8.34268332e-03, -3.99775570e-03, -3.27481248e-05 ], + [ -3.40162888e-02, -2.81868968e-02, -1.96679524e-04 ], + [ -2.91855410e-02, -3.97511162e-02, -2.81230678e-05 ], + [ -1.50395725e-02, -2.49494594e-02, -9.42573533e-05 ], + [ -1.67938769e-02, -2.00953931e-02, -4.00750607e-04 ], + [ -1.86435618e-02, -2.48535164e-02, -2.74416432e-02 ], + [ -4.61211195e-03, -1.21660791e-02, -2.93173041e-04 ], + [ -4.10017073e-02, -7.43824020e-02, -4.42762971e-02 ], + [ -1.90370996e-02, -3.74363363e-02, -1.34740388e-02 ] + ] +]).astype(numpy.float32) +EXPRESSION_MAX = numpy.array( +[ + [ + [ 4.46682945e-02, 7.08772913e-02, 4.08344204e-04 ], + [ 2.14308221e-02, 6.15894832e-02, 4.85319615e-05 ], + [ 3.02363783e-02, 4.45043296e-02, 1.28298725e-05 ], + [ 3.05869691e-02, 3.79812494e-02, 6.57040102e-04 ], + [ 4.45670523e-02, 3.97259220e-02, 7.10966764e-04 ], + [ 9.43699256e-02, 9.85926315e-02, 2.02551950e-04 ], + [ 1.61131397e-02, 2.92906128e-02, 3.44733417e-06 ], + [ 5.23825921e-02, 1.07065082e-01, 6.61510974e-04 ], + [ 2.85718683e-03, 8.32320191e-03, 2.39314613e-04 ], + [ 2.57947259e-02, 1.60935968e-02, 2.41853559e-05 ], + [ 4.90833223e-02, 3.43903080e-02, 3.22353356e-02 ], + [ 1.44766076e-02, 3.39248963e-02, 1.42291479e-04 ], + [ 8.75749043e-04, 6.82212645e-03, 2.76097053e-05 ], + [ 1.86958015e-02, 3.84016186e-02, 7.33085908e-05 ], + [ 2.01714113e-02, 4.90544215e-02, 2.34028921e-05 ], + [ 2.46518422e-02, 3.29151377e-02, 3.48571630e-05 ], + [ 2.22457591e-02, 1.21796541e-02, 1.56396593e-04 ], + [ 1.72109623e-02, 3.01626958e-02, 1.36556877e-02 ], + [ 1.83460284e-02, 1.61141958e-02, 2.87440169e-04 ], + [ 3.57594155e-02, 1.80554688e-01, 2.75554154e-02 ], + [ 2.17450950e-02, 8.66811201e-02, 3.34241726e-02 ] + ] +]).astype(numpy.float32) + + +def limit_expression(expression : LivePortraitExpression) -> LivePortraitExpression: + return numpy.clip(expression, EXPRESSION_MIN, EXPRESSION_MAX) + + +def limit_angle(target_pitch : LivePortraitPitch, target_yaw : LivePortraitYaw, target_roll : LivePortraitRoll, output_pitch : LivePortraitPitch, output_yaw : LivePortraitYaw, output_roll : LivePortraitRoll) -> Tuple[LivePortraitPitch, LivePortraitYaw, LivePortraitRoll]: + pitch_min, pitch_max, yaw_min, yaw_max, roll_min, roll_max = calculate_euler_limits(target_pitch, target_yaw, target_roll) + output_pitch = numpy.clip(output_pitch, pitch_min, pitch_max) + output_yaw = numpy.clip(output_yaw, yaw_min, yaw_max) + output_roll = numpy.clip(output_roll, roll_min, roll_max) + return output_pitch, output_yaw, output_roll + + +def calculate_euler_limits(pitch : LivePortraitPitch, yaw : LivePortraitYaw, roll : LivePortraitRoll) -> Tuple[float, float, float, float, float, float]: + pitch_min = -30.0 + pitch_max = 30.0 + yaw_min = -60.0 + yaw_max = 60.0 + roll_min = -20.0 + roll_max = 20.0 + + if pitch < 0: + pitch_min = min(pitch, pitch_min) + else: + pitch_max = max(pitch, pitch_max) + if yaw < 0: + yaw_min = min(yaw, yaw_min) + else: + yaw_max = max(yaw, yaw_max) + if roll < 0: + roll_min = min(roll, roll_min) + else: + roll_max = max(roll, roll_max) + + return pitch_min, pitch_max, yaw_min, yaw_max, roll_min, roll_max + + +def create_rotation(pitch : LivePortraitPitch, yaw : LivePortraitYaw, roll : LivePortraitRoll) -> LivePortraitRotation: + rotation = scipy.spatial.transform.Rotation.from_euler('xyz', [ pitch, yaw, roll ], degrees = True).as_matrix() + rotation = rotation.astype(numpy.float32) + return rotation diff --git a/facefusion/processors/modules/__init__.py b/facefusion/processors/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/age_modifier/__init__.py b/facefusion/processors/modules/age_modifier/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/age_modifier/choices.py b/facefusion/processors/modules/age_modifier/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..acd35decb580e891368452a24a802ad79c075097 --- /dev/null +++ b/facefusion/processors/modules/age_modifier/choices.py @@ -0,0 +1,8 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_int_range +from facefusion.processors.modules.age_modifier.types import AgeModifierModel + +age_modifier_models : List[AgeModifierModel] = [ 'styleganex_age' ] + +age_modifier_direction_range : Sequence[int] = create_int_range(-100, 100, 1) diff --git a/facefusion/processors/modules/age_modifier/core.py b/facefusion/processors/modules/age_modifier/core.py new file mode 100644 index 0000000000000000000000000000000000000000..026308b2de7d8504e5248e46c62142e15ff351db --- /dev/null +++ b/facefusion/processors/modules/age_modifier/core.py @@ -0,0 +1,219 @@ +from argparse import ArgumentParser +from functools import lru_cache + +import cv2 +import numpy + +import facefusion.choices +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import create_int_metavar, is_macos +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.execution import has_execution_provider +from facefusion.face_analyser import scale_face +from facefusion.face_helper import merge_matrix, paste_back, scale_face_landmark_5, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_box_mask, create_occlusion_mask +from facefusion.face_selector import select_faces +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors.modules.age_modifier import choices as age_modifier_choices +from facefusion.processors.modules.age_modifier.types import AgeModifierDirection, AgeModifierInputs +from facefusion.processors.types import ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import thread_semaphore +from facefusion.types import ApplyStateItem, Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import match_frame_color, read_static_image, read_static_video_frame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'styleganex_age': + { + '__metadata__': + { + 'vendor': 'williamyang1991', + 'license': 'S-Lab-1.0', + 'year': 2023 + }, + 'hashes': + { + 'age_modifier': + { + 'url': resolve_download_url('models-3.1.0', 'styleganex_age.hash'), + 'path': resolve_relative_path('../.assets/models/styleganex_age.hash') + } + }, + 'sources': + { + 'age_modifier': + { + 'url': resolve_download_url('models-3.1.0', 'styleganex_age.onnx'), + 'path': resolve_relative_path('../.assets/models/styleganex_age.onnx') + } + }, + 'templates': + { + 'target': 'ffhq_512', + 'target_with_background': 'styleganex_384' + }, + 'sizes': + { + 'target': (256, 256), + 'target_with_background': (384, 384) + } + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('age_modifier_model') ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('age_modifier_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + model_name = state_manager.get_item('age_modifier_model') + return create_static_model_set('full').get(model_name) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--age-modifier-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'age_modifier_model', 'styleganex_age'), choices = age_modifier_choices.age_modifier_models) + group_processors.add_argument('--age-modifier-direction', help = translator.get('help.direction', __package__), type = int, default = config.get_int_value('processors', 'age_modifier_direction', '0'), choices = age_modifier_choices.age_modifier_direction_range, metavar = create_int_metavar(age_modifier_choices.age_modifier_direction_range)) + facefusion.jobs.job_store.register_step_keys([ 'age_modifier_model', 'age_modifier_direction' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('age_modifier_model', args.get('age_modifier_model')) + apply_state_item('age_modifier_direction', args.get('age_modifier_direction')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_masker.clear_inference_pool() + face_recognizer.clear_inference_pool() + + +def modify_age(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + model_templates = get_model_options().get('templates') + model_sizes = get_model_options().get('sizes') + face_landmark_5 = target_face.landmark_set.get('5/68').copy() + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5, model_templates.get('target'), model_sizes.get('target')) + extend_face_landmark_5 = scale_face_landmark_5(face_landmark_5, 0.875) + extend_vision_frame, extend_affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, extend_face_landmark_5, model_templates.get('target_with_background'), model_sizes.get('target_with_background')) + extend_vision_frame_raw = extend_vision_frame.copy() + box_mask = create_box_mask(extend_vision_frame, state_manager.get_item('face_mask_blur'), (0, 0, 0, 0)) + crop_masks =\ + [ + box_mask + ] + + if 'occlusion' in state_manager.get_item('face_mask_types'): + occlusion_mask = create_occlusion_mask(crop_vision_frame) + temp_matrix = merge_matrix([ extend_affine_matrix, cv2.invertAffineTransform(affine_matrix) ]) + occlusion_mask = cv2.warpAffine(occlusion_mask, temp_matrix, model_sizes.get('target_with_background')) + crop_masks.append(occlusion_mask) + + crop_vision_frame = prepare_vision_frame(crop_vision_frame) + extend_vision_frame = prepare_vision_frame(extend_vision_frame) + age_modifier_direction = numpy.array(numpy.interp(state_manager.get_item('age_modifier_direction'), [ -100, 100 ], [ 2.5, -2.5 ])).astype(numpy.float32) + extend_vision_frame = forward(crop_vision_frame, extend_vision_frame, age_modifier_direction) + extend_vision_frame = normalize_extend_frame(extend_vision_frame) + extend_vision_frame = match_frame_color(extend_vision_frame_raw, extend_vision_frame) + extend_affine_matrix *= (model_sizes.get('target')[0] * 4) / model_sizes.get('target_with_background')[0] + crop_mask = numpy.minimum.reduce(crop_masks).clip(0, 1) + crop_mask = cv2.resize(crop_mask, (model_sizes.get('target')[0] * 4, model_sizes.get('target')[1] * 4)) + paste_vision_frame = paste_back(temp_vision_frame, extend_vision_frame, crop_mask, extend_affine_matrix) + return paste_vision_frame + + +def forward(crop_vision_frame : VisionFrame, extend_vision_frame : VisionFrame, age_modifier_direction : AgeModifierDirection) -> VisionFrame: + age_modifier = get_inference_pool().get('age_modifier') + age_modifier_inputs = {} + + if is_macos() and has_execution_provider('coreml'): + age_modifier.set_providers([ facefusion.choices.execution_provider_set.get('cpu') ]) + + for age_modifier_input in age_modifier.get_inputs(): + if age_modifier_input.name == 'target': + age_modifier_inputs[age_modifier_input.name] = crop_vision_frame + if age_modifier_input.name == 'target_with_background': + age_modifier_inputs[age_modifier_input.name] = extend_vision_frame + if age_modifier_input.name == 'direction': + age_modifier_inputs[age_modifier_input.name] = age_modifier_direction + + with thread_semaphore(): + crop_vision_frame = age_modifier.run(None, age_modifier_inputs)[0][0] + + return crop_vision_frame + + +def prepare_vision_frame(vision_frame : VisionFrame) -> VisionFrame: + vision_frame = vision_frame[:, :, ::-1] / 255.0 + vision_frame = (vision_frame - 0.5) / 0.5 + vision_frame = numpy.expand_dims(vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) + return vision_frame + + +def normalize_extend_frame(extend_vision_frame : VisionFrame) -> VisionFrame: + model_sizes = get_model_options().get('sizes') + extend_vision_frame = numpy.clip(extend_vision_frame, -1, 1) + extend_vision_frame = (extend_vision_frame + 1) / 2 + extend_vision_frame = extend_vision_frame.transpose(1, 2, 0).clip(0, 255) + extend_vision_frame = (extend_vision_frame * 255.0) + extend_vision_frame = extend_vision_frame.astype(numpy.uint8)[:, :, ::-1] + extend_vision_frame = cv2.resize(extend_vision_frame, (model_sizes.get('target')[0] * 4, model_sizes.get('target')[1] * 4), interpolation = cv2.INTER_AREA) + return extend_vision_frame + + +def process_frame(inputs : AgeModifierInputs) -> ProcessorOutputs: + reference_vision_frame = inputs.get('reference_vision_frame') + target_vision_frame = inputs.get('target_vision_frame') + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + target_faces = select_faces(reference_vision_frame, target_vision_frame) + + if target_faces: + for target_face in target_faces: + target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) + temp_vision_frame = modify_age(target_face, temp_vision_frame) + + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/age_modifier/locales.py b/facefusion/processors/modules/age_modifier/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..f25a3d1053c1eaa7eb8d8c376cb5c201366f1763 --- /dev/null +++ b/facefusion/processors/modules/age_modifier/locales.py @@ -0,0 +1,18 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for aging the face', + 'direction': 'specify the direction in which the age should be modified' + }, + 'uis': + { + 'direction_slider': 'AGE MODIFIER DIRECTION', + 'model_dropdown': 'AGE MODIFIER MODEL' + } + } +} diff --git a/facefusion/processors/modules/age_modifier/types.py b/facefusion/processors/modules/age_modifier/types.py new file mode 100644 index 0000000000000000000000000000000000000000..a804ad38caab50e9dbbfdb78fe59060c3cc2b06a --- /dev/null +++ b/facefusion/processors/modules/age_modifier/types.py @@ -0,0 +1,17 @@ +from typing import Any, Literal, TypeAlias, TypedDict + +from numpy.typing import NDArray + +from facefusion.types import Mask, VisionFrame + +AgeModifierInputs = TypedDict('AgeModifierInputs', +{ + 'reference_vision_frame' : VisionFrame, + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +AgeModifierModel = Literal['styleganex_age'] + +AgeModifierDirection : TypeAlias = NDArray[Any] diff --git a/facefusion/processors/modules/background_remover/__init__.py b/facefusion/processors/modules/background_remover/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/background_remover/choices.py b/facefusion/processors/modules/background_remover/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..ae0380d6c5e60d220e76119aa3a3bc40f9672bd1 --- /dev/null +++ b/facefusion/processors/modules/background_remover/choices.py @@ -0,0 +1,8 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_int_range +from facefusion.processors.modules.background_remover.types import BackgroundRemoverModel + +background_remover_models : List[BackgroundRemoverModel] = [ 'ben_2', 'birefnet_general', 'birefnet_portrait', 'isnet_general', 'modnet', 'ormbg', 'rmbg_1.4', 'rmbg_2.0', 'silueta', 'u2net_cloth', 'u2net_general', 'u2net_human', 'u2netp' ] + +background_remover_color_range : Sequence[int] = create_int_range(0, 255, 1) diff --git a/facefusion/processors/modules/background_remover/core.py b/facefusion/processors/modules/background_remover/core.py new file mode 100644 index 0000000000000000000000000000000000000000..86547f3d0fe0066768580f5304fee7f30941670c --- /dev/null +++ b/facefusion/processors/modules/background_remover/core.py @@ -0,0 +1,524 @@ +from argparse import ArgumentParser +from functools import lru_cache, partial +from typing import List, Tuple + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import is_macos +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.execution import has_execution_provider +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.normalizer import normalize_color +from facefusion.processors.modules.background_remover import choices as background_remover_choices +from facefusion.processors.modules.background_remover.types import BackgroundRemoverInputs +from facefusion.processors.types import ProcessorOutputs +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 ApplyStateItem, Args, DownloadScope, ExecutionProvider, InferencePool, Mask, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import read_static_image, read_static_video_frame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'ben_2': + { + '__metadata__': + { + 'vendor': 'PramaLLC', + 'license': 'MIT', + 'year': 2025 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'ben_2.hash'), + 'path': resolve_relative_path('../.assets/models/ben_2.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'ben_2.onnx'), + 'path': resolve_relative_path('../.assets/models/ben_2.onnx') + } + }, + 'size': (1024, 1024), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'birefnet_general': + { + '__metadata__': + { + 'vendor': 'ZhengPeng7', + 'license': 'MIT', + 'year': 2024 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'birefnet_general.hash'), + 'path': resolve_relative_path('../.assets/models/birefnet_general.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'birefnet_general.onnx'), + 'path': resolve_relative_path('../.assets/models/birefnet_general.onnx') + } + }, + 'size': (1024, 1024), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'birefnet_portrait': + { + '__metadata__': + { + 'vendor': 'ZhengPeng7', + 'license': 'MIT', + 'year': 2024 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'birefnet_portrait.hash'), + 'path': resolve_relative_path('../.assets/models/birefnet_portrait.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'birefnet_portrait.onnx'), + 'path': resolve_relative_path('../.assets/models/birefnet_portrait.onnx') + } + }, + 'size': (1024, 1024), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'isnet_general': + { + '__metadata__': + { + 'vendor': 'xuebinqin', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'isnet_general.hash'), + 'path': resolve_relative_path('../.assets/models/isnet_general.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'isnet_general.onnx'), + 'path': resolve_relative_path('../.assets/models/isnet_general.onnx') + } + }, + 'size': (1024, 1024), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'modnet': + { + '__metadata__': + { + 'vendor': 'ZHKKKe', + 'license': 'Apache-2.0', + 'year': 2020 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'modnet.hash'), + 'path': resolve_relative_path('../.assets/models/modnet.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'modnet.onnx'), + 'path': resolve_relative_path('../.assets/models/modnet.onnx') + } + }, + 'size': (512, 512), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + }, + 'ormbg': + { + '__metadata__': + { + 'vendor': 'schirrmacher', + 'license': 'Apache-2.0', + 'year': 2024 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'ormbg.hash'), + 'path': resolve_relative_path('../.assets/models/ormbg.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'ormbg.onnx'), + 'path': resolve_relative_path('../.assets/models/ormbg.onnx') + } + }, + 'size': (1024, 1024), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'rmbg_1.4': + { + '__metadata__': + { + 'vendor': 'Bria', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'rmbg_1.4.hash'), + 'path': resolve_relative_path('../.assets/models/rmbg_1.4.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'rmbg_1.4.onnx'), + 'path': resolve_relative_path('../.assets/models/rmbg_1.4.onnx') + } + }, + 'size': (1024, 1024), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'rmbg_2.0': + { + '__metadata__': + { + 'vendor': 'Bria', + 'license': 'Non-Commercial', + 'year': 2024 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'rmbg_2.0.hash'), + 'path': resolve_relative_path('../.assets/models/rmbg_2.0.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'rmbg_2.0.onnx'), + 'path': resolve_relative_path('../.assets/models/rmbg_2.0.onnx') + } + }, + 'size': (1024, 1024), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + }, + 'silueta': + { + '__metadata__': + { + 'vendor': 'Kikedao', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'silueta.hash'), + 'path': resolve_relative_path('../.assets/models/silueta.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'silueta.onnx'), + 'path': resolve_relative_path('../.assets/models/silueta.onnx') + } + }, + 'size': (320, 320), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + }, + 'u2net_cloth': + { + '__metadata__': + { + 'vendor': 'levindabhi', + 'license': 'MIT', + 'year': 2021 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'u2net_cloth.hash'), + 'path': resolve_relative_path('../.assets/models/u2net_cloth.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'u2net_cloth.onnx'), + 'path': resolve_relative_path('../.assets/models/u2net_cloth.onnx') + } + }, + 'size': (768, 768), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + }, + 'u2net_general': + { + '__metadata__': + { + 'vendor': 'xuebinqin', + 'license': 'Apache-2.0', + 'year': 2020 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'u2net_general.hash'), + 'path': resolve_relative_path('../.assets/models/u2net_general.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'u2net_general.onnx'), + 'path': resolve_relative_path('../.assets/models/u2net_general.onnx') + } + }, + 'size': (320, 320), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + }, + 'u2net_human': + { + '__metadata__': + { + 'vendor': 'xuebinqin', + 'license': 'Apache-2.0', + 'year': 2021 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'u2net_human.hash'), + 'path': resolve_relative_path('../.assets/models/u2net_human.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'u2net_human.onnx'), + 'path': resolve_relative_path('../.assets/models/u2net_human.onnx') + } + }, + 'size': (320, 320), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + }, + 'u2netp': + { + '__metadata__': + { + 'vendor': 'xuebinqin', + 'license': 'Apache-2.0', + 'year': 2021 + }, + 'hashes': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'u2netp.hash'), + 'path': resolve_relative_path('../.assets/models/u2netp.hash') + } + }, + 'sources': + { + 'background_remover': + { + 'url': resolve_download_url('models-3.5.0', 'u2netp.onnx'), + 'path': resolve_relative_path('../.assets/models/u2netp.onnx') + } + }, + 'size': (320, 320), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('background_remover_model') ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('background_remover_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def resolve_execution_providers() -> List[ExecutionProvider]: + if is_macos() and has_execution_provider('coreml'): + return [ 'cpu' ] + return state_manager.get_item('execution_providers') + + +def get_model_options() -> ModelOptions: + model_name = state_manager.get_item('background_remover_model') + return create_static_model_set('full').get(model_name) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--background-remover-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'background_remover_model', 'rmbg_2.0'), choices = background_remover_choices.background_remover_models) + group_processors.add_argument('--background-remover-color', help = translator.get('help.color', __package__), type = partial(sanitize_int_range, int_range = background_remover_choices.background_remover_color_range), default = config.get_int_list('processors', 'background_remover_color', '0 0 0 0'), nargs = '+') + facefusion.jobs.job_store.register_step_keys([ 'background_remover_model', 'background_remover_color' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('background_remover_model', args.get('background_remover_model')) + apply_state_item('background_remover_color', normalize_color(args.get('background_remover_color'))) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + + +def remove_background(temp_vision_frame : VisionFrame) -> Tuple[VisionFrame, Mask]: + temp_vision_mask = forward(prepare_temp_frame(temp_vision_frame)) + temp_vision_mask = normalize_vision_mask(temp_vision_mask) + temp_vision_mask = cv2.resize(temp_vision_mask, temp_vision_frame.shape[:2][::-1]) + temp_vision_frame = apply_background_color(temp_vision_frame, temp_vision_mask) + return temp_vision_frame, temp_vision_mask + + +def forward(temp_vision_frame : VisionFrame) -> VisionFrame: + background_remover = get_inference_pool().get('background_remover') + model_name = state_manager.get_item('background_remover_model') + + with thread_semaphore(): + remove_vision_frame = background_remover.run(None, + { + 'input': temp_vision_frame + })[0] + + if model_name == 'u2net_cloth': + remove_vision_frame = numpy.argmax(remove_vision_frame, axis = 1) + + return remove_vision_frame + + +def prepare_temp_frame(temp_vision_frame : VisionFrame) -> VisionFrame: + model_size = get_model_options().get('size') + model_mean = get_model_options().get('mean') + model_standard_deviation = get_model_options().get('standard_deviation') + + temp_vision_frame = cv2.resize(temp_vision_frame, model_size) + temp_vision_frame = temp_vision_frame[:, :, ::-1] / 255.0 + temp_vision_frame = (temp_vision_frame - model_mean) / model_standard_deviation + temp_vision_frame = temp_vision_frame.transpose(2, 0, 1) + temp_vision_frame = numpy.expand_dims(temp_vision_frame, axis = 0).astype(numpy.float32) + return temp_vision_frame + + +def normalize_vision_mask(temp_vision_mask : Mask) -> Mask: + temp_vision_mask = numpy.squeeze(temp_vision_mask).clip(0, 1) * 255 + temp_vision_mask = numpy.clip(temp_vision_mask, 0, 255).astype(numpy.uint8) + return temp_vision_mask + + +def apply_background_color(temp_vision_frame : VisionFrame, temp_vision_mask : Mask) -> VisionFrame: + background_remover_color = state_manager.get_item('background_remover_color') + temp_vision_mask = temp_vision_mask.astype(numpy.float32) / 255 + temp_vision_mask = numpy.expand_dims(temp_vision_mask, axis = 2) + temp_vision_mask = (1 - temp_vision_mask) * background_remover_color[-1] / 255 + color_frame = numpy.zeros_like(temp_vision_frame) + color_frame[:, :, 0] = background_remover_color[2] + color_frame[:, :, 1] = background_remover_color[1] + color_frame[:, :, 2] = background_remover_color[0] + temp_vision_frame = temp_vision_frame * (1 - temp_vision_mask) + color_frame * temp_vision_mask + temp_vision_frame = temp_vision_frame.astype(numpy.uint8) + return temp_vision_frame + + +def process_frame(inputs : BackgroundRemoverInputs) -> ProcessorOutputs: + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_frame, temp_vision_mask = remove_background(temp_vision_frame) + temp_vision_mask = numpy.minimum.reduce([ temp_vision_mask, inputs.get('temp_vision_mask') ]) + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/background_remover/locales.py b/facefusion/processors/modules/background_remover/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..267bcb6c2fdcdc454f7da24783b70001998237a0 --- /dev/null +++ b/facefusion/processors/modules/background_remover/locales.py @@ -0,0 +1,21 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for removing the background', + 'color': 'apply red, green blue and alpha values to the background' + }, + 'uis': + { + 'model_dropdown': 'BACKGROUND REMOVER MODEL', + 'color_red_number': 'BACKGROUND COLOR RED', + 'color_green_number': 'BACKGROUND COLOR GREEN', + 'color_blue_number': 'BACKGROUND COLOR BLUE', + 'color_alpha_number': 'BACKGROUND COLOR ALPHA' + } + } +} diff --git a/facefusion/processors/modules/background_remover/types.py b/facefusion/processors/modules/background_remover/types.py new file mode 100644 index 0000000000000000000000000000000000000000..afcba01e60420e23789a35ac722be2dcb52cd3d4 --- /dev/null +++ b/facefusion/processors/modules/background_remover/types.py @@ -0,0 +1,12 @@ +from typing import Literal, TypedDict + +from facefusion.types import Mask, VisionFrame + +BackgroundRemoverInputs = TypedDict('BackgroundRemoverInputs', +{ + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +BackgroundRemoverModel = Literal['ben_2', 'birefnet_general', 'birefnet_portrait', 'isnet_general', 'modnet', 'ormbg', 'rmbg_1.4', 'rmbg_2.0', 'silueta', 'u2net_cloth', 'u2net_general', 'u2net_human', 'u2netp'] diff --git a/facefusion/processors/modules/deep_swapper/__init__.py b/facefusion/processors/modules/deep_swapper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/deep_swapper/choices.py b/facefusion/processors/modules/deep_swapper/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..c08ebeab9162b13511c428cb966ac6157cf07984 --- /dev/null +++ b/facefusion/processors/modules/deep_swapper/choices.py @@ -0,0 +1,176 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_int_range +from facefusion.filesystem import get_file_name, resolve_file_paths, resolve_relative_path +from facefusion.processors.modules.deep_swapper.types import DeepSwapperModel + +deep_swapper_models : List[DeepSwapperModel] =\ +[ + 'druuzil/adam_levine_320', + 'druuzil/adrianne_palicki_384', + 'druuzil/agnetha_falskog_224', + 'druuzil/alan_ritchson_320', + 'druuzil/alicia_vikander_320', + 'druuzil/amber_midthunder_320', + 'druuzil/andras_arato_384', + 'druuzil/andrew_tate_320', + 'druuzil/angelina_jolie_384', + 'druuzil/anne_hathaway_320', + 'druuzil/anya_chalotra_320', + 'druuzil/arnold_schwarzenegger_320', + 'druuzil/benjamin_affleck_320', + 'druuzil/benjamin_stiller_384', + 'druuzil/bradley_pitt_224', + 'druuzil/brie_larson_384', + 'druuzil/bruce_campbell_384', + 'druuzil/bryan_cranston_320', + 'druuzil/catherine_blanchett_352', + 'druuzil/christian_bale_320', + 'druuzil/christopher_hemsworth_320', + 'druuzil/christoph_waltz_384', + 'druuzil/cillian_murphy_320', + 'druuzil/cobie_smulders_256', + 'druuzil/dwayne_johnson_384', + 'druuzil/edward_norton_320', + 'druuzil/elisabeth_shue_320', + 'druuzil/elizabeth_olsen_384', + 'druuzil/elon_musk_320', + 'druuzil/emily_blunt_320', + 'druuzil/emma_stone_384', + 'druuzil/emma_watson_320', + 'druuzil/erin_moriarty_384', + 'druuzil/eva_green_320', + 'druuzil/ewan_mcgregor_320', + 'druuzil/florence_pugh_320', + 'druuzil/freya_allan_320', + 'druuzil/gary_cole_224', + 'druuzil/gigi_hadid_224', + 'druuzil/harrison_ford_384', + 'druuzil/hayden_christensen_320', + 'druuzil/heath_ledger_320', + 'druuzil/henry_cavill_448', + 'druuzil/hugh_jackman_384', + 'druuzil/idris_elba_320', + 'druuzil/jack_nicholson_320', + 'druuzil/james_carrey_384', + 'druuzil/james_mcavoy_320', + 'druuzil/james_varney_320', + 'druuzil/jason_momoa_320', + 'druuzil/jason_statham_320', + 'druuzil/jennifer_connelly_384', + 'druuzil/jimmy_donaldson_320', + 'druuzil/jordan_peterson_384', + 'druuzil/karl_urban_224', + 'druuzil/kate_beckinsale_384', + 'druuzil/laurence_fishburne_384', + 'druuzil/lili_reinhart_320', + 'druuzil/luke_evans_384', + 'druuzil/mads_mikkelsen_384', + 'druuzil/mary_winstead_320', + 'druuzil/margaret_qualley_384', + 'druuzil/melina_juergens_320', + 'druuzil/michael_fassbender_320', + 'druuzil/michael_fox_320', + 'druuzil/millie_bobby_brown_320', + 'druuzil/morgan_freeman_320', + 'druuzil/patrick_stewart_224', + 'druuzil/rachel_weisz_384', + 'druuzil/rebecca_ferguson_320', + 'druuzil/scarlett_johansson_320', + 'druuzil/shannen_doherty_384', + 'druuzil/seth_macfarlane_384', + 'druuzil/thomas_cruise_320', + 'druuzil/thomas_hanks_384', + 'druuzil/william_murray_384', + 'druuzil/zoe_saldana_384', + 'edel/emma_roberts_224', + 'edel/ivanka_trump_224', + 'edel/lize_dzjabrailova_224', + 'edel/sidney_sweeney_224', + 'edel/winona_ryder_224', + 'iperov/alexandra_daddario_224', + 'iperov/alexei_navalny_224', + 'iperov/amber_heard_224', + 'iperov/dilraba_dilmurat_224', + 'iperov/elon_musk_224', + 'iperov/emilia_clarke_224', + 'iperov/emma_watson_224', + 'iperov/erin_moriarty_224', + 'iperov/jackie_chan_224', + 'iperov/james_carrey_224', + 'iperov/jason_statham_320', + 'iperov/keanu_reeves_320', + 'iperov/margot_robbie_224', + 'iperov/natalie_dormer_224', + 'iperov/nicolas_coppola_224', + 'iperov/robert_downey_224', + 'iperov/rowan_atkinson_224', + 'iperov/ryan_reynolds_224', + 'iperov/scarlett_johansson_224', + 'iperov/sylvester_stallone_224', + 'iperov/thomas_cruise_224', + 'iperov/thomas_holland_224', + 'iperov/vin_diesel_224', + 'iperov/vladimir_putin_224', + 'jen/angelica_trae_288', + 'jen/ella_freya_224', + 'jen/emma_myers_320', + 'jen/evie_pickerill_224', + 'jen/kang_hyewon_320', + 'jen/maddie_mead_224', + 'jen/nicole_turnbull_288', + 'mats/alica_schmidt_320', + 'mats/ashley_alexiss_224', + 'mats/billie_eilish_224', + 'mats/brie_larson_224', + 'mats/cara_delevingne_224', + 'mats/carolin_kebekus_224', + 'mats/chelsea_clinton_224', + 'mats/claire_boucher_224', + 'mats/corinna_kopf_224', + 'mats/florence_pugh_224', + 'mats/hillary_clinton_224', + 'mats/jenna_fischer_224', + 'mats/kim_jisoo_320', + 'mats/mica_suarez_320', + 'mats/shailene_woodley_224', + 'mats/shraddha_kapoor_320', + 'mats/yu_jimin_352', + 'rumateus/alison_brie_224', + 'rumateus/amber_heard_224', + 'rumateus/angelina_jolie_224', + 'rumateus/aubrey_plaza_224', + 'rumateus/bridget_regan_224', + 'rumateus/cobie_smulders_224', + 'rumateus/deborah_woll_224', + 'rumateus/dua_lipa_224', + 'rumateus/emma_stone_224', + 'rumateus/hailee_steinfeld_224', + 'rumateus/hilary_duff_224', + 'rumateus/jessica_alba_224', + 'rumateus/jessica_biel_224', + 'rumateus/john_cena_224', + 'rumateus/kim_kardashian_224', + 'rumateus/kristen_bell_224', + 'rumateus/lucy_liu_224', + 'rumateus/margot_robbie_224', + 'rumateus/megan_fox_224', + 'rumateus/meghan_markle_224', + 'rumateus/millie_bobby_brown_224', + 'rumateus/natalie_portman_224', + 'rumateus/nicki_minaj_224', + 'rumateus/olivia_wilde_224', + 'rumateus/shay_mitchell_224', + 'rumateus/sophie_turner_224', + 'rumateus/taylor_swift_224' +] + +custom_model_file_paths = resolve_file_paths(resolve_relative_path('../.assets/models/custom')) + +if custom_model_file_paths: + + for model_file_path in custom_model_file_paths: + model_id = '/'.join([ 'custom', get_file_name(model_file_path) ]) + deep_swapper_models.append(model_id) + +deep_swapper_morph_range : Sequence[int] = create_int_range(0, 100, 1) diff --git a/facefusion/processors/modules/deep_swapper/core.py b/facefusion/processors/modules/deep_swapper/core.py new file mode 100644 index 0000000000000000000000000000000000000000..c443b6c25e1f2c759e9fb95ce2de471f29fecb2e --- /dev/null +++ b/facefusion/processors/modules/deep_swapper/core.py @@ -0,0 +1,424 @@ +from argparse import ArgumentParser +from functools import lru_cache +from typing import Tuple + +import cv2 +import numpy +from cv2.typing import Size + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import create_int_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url_by_provider +from facefusion.face_analyser import scale_face +from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_area_mask, create_box_mask, create_occlusion_mask, create_region_mask +from facefusion.face_selector import select_faces +from facefusion.filesystem import get_file_name, in_directory, is_image, is_video, resolve_file_paths, resolve_relative_path, same_file_extension +from facefusion.processors.modules.deep_swapper import choices as deep_swapper_choices +from facefusion.processors.modules.deep_swapper.types import DeepSwapperInputs, DeepSwapperMorph +from facefusion.processors.types import ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import thread_semaphore +from facefusion.types import ApplyStateItem, Args, DownloadScope, Face, InferencePool, Mask, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import conditional_match_frame_color, read_static_image, read_static_video_frame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + model_config = [] + + if download_scope == 'full': + model_config.extend( + [ + ('druuzil', 'adam_levine_320'), + ('druuzil', 'adrianne_palicki_384'), + ('druuzil', 'agnetha_falskog_224'), + ('druuzil', 'alan_ritchson_320'), + ('druuzil', 'alicia_vikander_320'), + ('druuzil', 'amber_midthunder_320'), + ('druuzil', 'andras_arato_384'), + ('druuzil', 'andrew_tate_320'), + ('druuzil', 'angelina_jolie_384'), + ('druuzil', 'anne_hathaway_320'), + ('druuzil', 'anya_chalotra_320'), + ('druuzil', 'arnold_schwarzenegger_320'), + ('druuzil', 'benjamin_affleck_320'), + ('druuzil', 'benjamin_stiller_384'), + ('druuzil', 'bradley_pitt_224'), + ('druuzil', 'brie_larson_384'), + ('druuzil', 'bruce_campbell_384'), + ('druuzil', 'bryan_cranston_320'), + ('druuzil', 'catherine_blanchett_352'), + ('druuzil', 'christian_bale_320'), + ('druuzil', 'christopher_hemsworth_320'), + ('druuzil', 'christoph_waltz_384'), + ('druuzil', 'cillian_murphy_320'), + ('druuzil', 'cobie_smulders_256'), + ('druuzil', 'dwayne_johnson_384'), + ('druuzil', 'edward_norton_320'), + ('druuzil', 'elisabeth_shue_320'), + ('druuzil', 'elizabeth_olsen_384'), + ('druuzil', 'elon_musk_320'), + ('druuzil', 'emily_blunt_320'), + ('druuzil', 'emma_stone_384'), + ('druuzil', 'emma_watson_320'), + ('druuzil', 'erin_moriarty_384'), + ('druuzil', 'eva_green_320'), + ('druuzil', 'ewan_mcgregor_320'), + ('druuzil', 'florence_pugh_320'), + ('druuzil', 'freya_allan_320'), + ('druuzil', 'gary_cole_224'), + ('druuzil', 'gigi_hadid_224'), + ('druuzil', 'harrison_ford_384'), + ('druuzil', 'hayden_christensen_320'), + ('druuzil', 'heath_ledger_320'), + ('druuzil', 'henry_cavill_448'), + ('druuzil', 'hugh_jackman_384'), + ('druuzil', 'idris_elba_320'), + ('druuzil', 'jack_nicholson_320'), + ('druuzil', 'james_carrey_384'), + ('druuzil', 'james_mcavoy_320'), + ('druuzil', 'james_varney_320'), + ('druuzil', 'jason_momoa_320'), + ('druuzil', 'jason_statham_320'), + ('druuzil', 'jennifer_connelly_384'), + ('druuzil', 'jimmy_donaldson_320'), + ('druuzil', 'jordan_peterson_384'), + ('druuzil', 'karl_urban_224'), + ('druuzil', 'kate_beckinsale_384'), + ('druuzil', 'laurence_fishburne_384'), + ('druuzil', 'lili_reinhart_320'), + ('druuzil', 'luke_evans_384'), + ('druuzil', 'mads_mikkelsen_384'), + ('druuzil', 'mary_winstead_320'), + ('druuzil', 'margaret_qualley_384'), + ('druuzil', 'melina_juergens_320'), + ('druuzil', 'michael_fassbender_320'), + ('druuzil', 'michael_fox_320'), + ('druuzil', 'millie_bobby_brown_320'), + ('druuzil', 'morgan_freeman_320'), + ('druuzil', 'patrick_stewart_224'), + ('druuzil', 'rachel_weisz_384'), + ('druuzil', 'rebecca_ferguson_320'), + ('druuzil', 'scarlett_johansson_320'), + ('druuzil', 'shannen_doherty_384'), + ('druuzil', 'seth_macfarlane_384'), + ('druuzil', 'thomas_cruise_320'), + ('druuzil', 'thomas_hanks_384'), + ('druuzil', 'william_murray_384'), + ('druuzil', 'zoe_saldana_384'), + ('edel', 'emma_roberts_224'), + ('edel', 'ivanka_trump_224'), + ('edel', 'lize_dzjabrailova_224'), + ('edel', 'sidney_sweeney_224'), + ('edel', 'winona_ryder_224') + ]) + if download_scope in [ 'lite', 'full' ]: + model_config.extend( + [ + ('iperov', 'alexandra_daddario_224'), + ('iperov', 'alexei_navalny_224'), + ('iperov', 'amber_heard_224'), + ('iperov', 'dilraba_dilmurat_224'), + ('iperov', 'elon_musk_224'), + ('iperov', 'emilia_clarke_224'), + ('iperov', 'emma_watson_224'), + ('iperov', 'erin_moriarty_224'), + ('iperov', 'jackie_chan_224'), + ('iperov', 'james_carrey_224'), + ('iperov', 'jason_statham_320'), + ('iperov', 'keanu_reeves_320'), + ('iperov', 'margot_robbie_224'), + ('iperov', 'natalie_dormer_224'), + ('iperov', 'nicolas_coppola_224'), + ('iperov', 'robert_downey_224'), + ('iperov', 'rowan_atkinson_224'), + ('iperov', 'ryan_reynolds_224'), + ('iperov', 'scarlett_johansson_224'), + ('iperov', 'sylvester_stallone_224'), + ('iperov', 'thomas_cruise_224'), + ('iperov', 'thomas_holland_224'), + ('iperov', 'vin_diesel_224'), + ('iperov', 'vladimir_putin_224') + ]) + if download_scope == 'full': + model_config.extend( + [ + ('jen', 'angelica_trae_288'), + ('jen', 'ella_freya_224'), + ('jen', 'emma_myers_320'), + ('jen', 'evie_pickerill_224'), + ('jen', 'kang_hyewon_320'), + ('jen', 'maddie_mead_224'), + ('jen', 'nicole_turnbull_288'), + ('mats', 'alica_schmidt_320'), + ('mats', 'ashley_alexiss_224'), + ('mats', 'billie_eilish_224'), + ('mats', 'brie_larson_224'), + ('mats', 'cara_delevingne_224'), + ('mats', 'carolin_kebekus_224'), + ('mats', 'chelsea_clinton_224'), + ('mats', 'claire_boucher_224'), + ('mats', 'corinna_kopf_224'), + ('mats', 'florence_pugh_224'), + ('mats', 'hillary_clinton_224'), + ('mats', 'jenna_fischer_224'), + ('mats', 'kim_jisoo_320'), + ('mats', 'mica_suarez_320'), + ('mats', 'shailene_woodley_224'), + ('mats', 'shraddha_kapoor_320'), + ('mats', 'yu_jimin_352'), + ('rumateus', 'alison_brie_224'), + ('rumateus', 'amber_heard_224'), + ('rumateus', 'angelina_jolie_224'), + ('rumateus', 'aubrey_plaza_224'), + ('rumateus', 'bridget_regan_224'), + ('rumateus', 'cobie_smulders_224'), + ('rumateus', 'deborah_woll_224'), + ('rumateus', 'dua_lipa_224'), + ('rumateus', 'emma_stone_224'), + ('rumateus', 'hailee_steinfeld_224'), + ('rumateus', 'hilary_duff_224'), + ('rumateus', 'jessica_alba_224'), + ('rumateus', 'jessica_biel_224'), + ('rumateus', 'john_cena_224'), + ('rumateus', 'kim_kardashian_224'), + ('rumateus', 'kristen_bell_224'), + ('rumateus', 'lucy_liu_224'), + ('rumateus', 'margot_robbie_224'), + ('rumateus', 'megan_fox_224'), + ('rumateus', 'meghan_markle_224'), + ('rumateus', 'millie_bobby_brown_224'), + ('rumateus', 'natalie_portman_224'), + ('rumateus', 'nicki_minaj_224'), + ('rumateus', 'olivia_wilde_224'), + ('rumateus', 'shay_mitchell_224'), + ('rumateus', 'sophie_turner_224'), + ('rumateus', 'taylor_swift_224') + ]) + model_set : ModelSet = {} + + for model_scope, model_name in model_config: + model_id = '/'.join([ model_scope, model_name ]) + + model_set[model_id] =\ + { + 'hashes': + { + 'deep_swapper': + { + 'url': resolve_download_url_by_provider('huggingface', 'deepfacelive-models-' + model_scope, model_name + '.hash'), + 'path': resolve_relative_path('../.assets/models/' + model_scope + '/' + model_name + '.hash') + } + }, + 'sources': + { + 'deep_swapper': + { + 'url': resolve_download_url_by_provider('huggingface', 'deepfacelive-models-' + model_scope, model_name + '.dfm'), + 'path': resolve_relative_path('../.assets/models/' + model_scope + '/' + model_name + '.dfm') + } + }, + 'template': 'dfl_whole_face' + } + + custom_model_file_paths = resolve_file_paths(resolve_relative_path('../.assets/models/custom')) + + if custom_model_file_paths: + + for model_file_path in custom_model_file_paths: + model_id = '/'.join([ 'custom', get_file_name(model_file_path) ]) + + model_set[model_id] =\ + { + 'sources': + { + 'deep_swapper': + { + 'path': resolve_relative_path(model_file_path) + } + }, + 'template': 'dfl_whole_face' + } + + return model_set + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('deep_swapper_model') ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('deep_swapper_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + model_name = state_manager.get_item('deep_swapper_model') + return create_static_model_set('full').get(model_name) + + +def get_model_size() -> Size: + deep_swapper = get_inference_pool().get('deep_swapper') + + for deep_swapper_input in deep_swapper.get_inputs(): + if deep_swapper_input.name == 'in_face:0': + return deep_swapper_input.shape[1:3] + + return 0, 0 + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--deep-swapper-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'deep_swapper_model', 'iperov/elon_musk_224'), choices = deep_swapper_choices.deep_swapper_models) + group_processors.add_argument('--deep-swapper-morph', help = translator.get('help.morph', __package__), type = int, default = config.get_int_value('processors', 'deep_swapper_morph', '100'), choices = deep_swapper_choices.deep_swapper_morph_range, metavar = create_int_metavar(deep_swapper_choices.deep_swapper_morph_range)) + facefusion.jobs.job_store.register_step_keys([ 'deep_swapper_model', 'deep_swapper_morph' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('deep_swapper_model', args.get('deep_swapper_model')) + apply_state_item('deep_swapper_morph', args.get('deep_swapper_morph')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + if model_hash_set and model_source_set: + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + return True + + +def pre_process(mode : ProcessMode) -> bool: + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_masker.clear_inference_pool() + face_recognizer.clear_inference_pool() + + +def swap_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + model_template = get_model_options().get('template') + model_size = get_model_size() + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmark_set.get('5/68'), model_template, model_size) + crop_vision_frame_raw = crop_vision_frame.copy() + box_mask = create_box_mask(crop_vision_frame, state_manager.get_item('face_mask_blur'), state_manager.get_item('face_mask_padding')) + crop_masks =\ + [ + box_mask + ] + + if 'occlusion' in state_manager.get_item('face_mask_types'): + occlusion_mask = create_occlusion_mask(crop_vision_frame) + crop_masks.append(occlusion_mask) + + crop_vision_frame = prepare_crop_frame(crop_vision_frame) + deep_swapper_morph = numpy.array([ numpy.interp(state_manager.get_item('deep_swapper_morph'), [ 0, 100 ], [ 0, 1 ]) ]).astype(numpy.float32) + crop_vision_frame, crop_source_mask, crop_target_mask = forward(crop_vision_frame, deep_swapper_morph) + crop_vision_frame = normalize_crop_frame(crop_vision_frame) + crop_vision_frame = conditional_match_frame_color(crop_vision_frame_raw, crop_vision_frame) + crop_masks.append(prepare_crop_mask(crop_source_mask, crop_target_mask)) + + if 'area' in state_manager.get_item('face_mask_types'): + face_landmark_68 = cv2.transform(target_face.landmark_set.get('68').reshape(1, -1, 2), affine_matrix).reshape(-1, 2) + area_mask = create_area_mask(crop_vision_frame, face_landmark_68, state_manager.get_item('face_mask_areas')) + crop_masks.append(area_mask) + + if 'region' in state_manager.get_item('face_mask_types'): + region_mask = create_region_mask(crop_vision_frame, state_manager.get_item('face_mask_regions')) + crop_masks.append(region_mask) + + crop_mask = numpy.minimum.reduce(crop_masks).clip(0, 1) + paste_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, crop_mask, affine_matrix) + return paste_vision_frame + + +def forward(crop_vision_frame : VisionFrame, deep_swapper_morph : DeepSwapperMorph) -> Tuple[VisionFrame, Mask, Mask]: + deep_swapper = get_inference_pool().get('deep_swapper') + deep_swapper_inputs = {} + + for deep_swapper_input in deep_swapper.get_inputs(): + if deep_swapper_input.name == 'in_face:0': + deep_swapper_inputs[deep_swapper_input.name] = crop_vision_frame + if deep_swapper_input.name == 'morph_value:0': + deep_swapper_inputs[deep_swapper_input.name] = deep_swapper_morph + + with thread_semaphore(): + crop_target_mask, crop_vision_frame, crop_source_mask = deep_swapper.run(None, deep_swapper_inputs) + + return crop_vision_frame[0], crop_source_mask[0], crop_target_mask[0] + + +def has_morph_input() -> bool: + deep_swapper = get_inference_pool().get('deep_swapper') + + for deep_swapper_input in deep_swapper.get_inputs(): + if deep_swapper_input.name == 'morph_value:0': + return True + + return False + + +def prepare_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + crop_vision_frame = cv2.addWeighted(crop_vision_frame, 1.75, cv2.GaussianBlur(crop_vision_frame, (0, 0), 2), -0.75, 0) + crop_vision_frame = crop_vision_frame / 255.0 + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0).astype(numpy.float32) + return crop_vision_frame + + +def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + crop_vision_frame = (crop_vision_frame * 255.0).clip(0, 255) + crop_vision_frame = crop_vision_frame.astype(numpy.uint8) + return crop_vision_frame + + +def prepare_crop_mask(crop_source_mask : Mask, crop_target_mask : Mask) -> Mask: + model_size = get_model_size() + blur_size = 6.25 + kernel_size = 3 + crop_mask = numpy.minimum.reduce([ crop_source_mask, crop_target_mask ]) + crop_mask = crop_mask.reshape(model_size).clip(0, 1) + crop_mask = cv2.erode(crop_mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)), iterations = 2) + crop_mask = cv2.GaussianBlur(crop_mask, (0, 0), blur_size) + return crop_mask + + +def process_frame(inputs : DeepSwapperInputs) -> ProcessorOutputs: + reference_vision_frame = inputs.get('reference_vision_frame') + target_vision_frame = inputs.get('target_vision_frame') + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + target_faces = select_faces(reference_vision_frame, target_vision_frame) + + if target_faces: + for target_face in target_faces: + target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) + temp_vision_frame = swap_face(target_face, temp_vision_frame) + + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/deep_swapper/locales.py b/facefusion/processors/modules/deep_swapper/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..25e8655321d69277ce088e10d4b91d55dccdfaa5 --- /dev/null +++ b/facefusion/processors/modules/deep_swapper/locales.py @@ -0,0 +1,18 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for swapping the face', + 'morph': 'morph between source face and target faces' + }, + 'uis': + { + 'model_dropdown': 'DEEP SWAPPER MODEL', + 'morph_slider': 'DEEP SWAPPER MORPH' + } + } +} diff --git a/facefusion/processors/modules/deep_swapper/types.py b/facefusion/processors/modules/deep_swapper/types.py new file mode 100644 index 0000000000000000000000000000000000000000..633ef63c116e4aaaf3114ca09761e401d0c0500d --- /dev/null +++ b/facefusion/processors/modules/deep_swapper/types.py @@ -0,0 +1,17 @@ +from typing import Any, TypeAlias, TypedDict + +from numpy.typing import NDArray + +from facefusion.types import Mask, VisionFrame + +DeepSwapperInputs = TypedDict('DeepSwapperInputs', +{ + 'reference_vision_frame' : VisionFrame, + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +DeepSwapperModel : TypeAlias = str + +DeepSwapperMorph : TypeAlias = NDArray[Any] diff --git a/facefusion/processors/modules/expression_restorer/__init__.py b/facefusion/processors/modules/expression_restorer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/expression_restorer/choices.py b/facefusion/processors/modules/expression_restorer/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..abca8883c652037525651ae65cb1e1a97022edd0 --- /dev/null +++ b/facefusion/processors/modules/expression_restorer/choices.py @@ -0,0 +1,10 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_int_range +from facefusion.processors.modules.expression_restorer.types import ExpressionRestorerArea, ExpressionRestorerModel + +expression_restorer_models : List[ExpressionRestorerModel] = [ 'live_portrait' ] + +expression_restorer_areas : List[ExpressionRestorerArea] = [ 'upper-face', 'lower-face' ] + +expression_restorer_factor_range : Sequence[int] = create_int_range(0, 100, 1) diff --git a/facefusion/processors/modules/expression_restorer/core.py b/facefusion/processors/modules/expression_restorer/core.py new file mode 100644 index 0000000000000000000000000000000000000000..50cf12fff2b73235d6185b71aa6564fe5b71179c --- /dev/null +++ b/facefusion/processors/modules/expression_restorer/core.py @@ -0,0 +1,270 @@ +from argparse import ArgumentParser +from functools import lru_cache +from typing import Tuple + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import create_int_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.face_analyser import scale_face +from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_box_mask, create_occlusion_mask +from facefusion.face_selector import select_faces +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors.live_portrait import create_rotation, limit_expression +from facefusion.processors.modules.expression_restorer import choices as expression_restorer_choices +from facefusion.processors.modules.expression_restorer.types import ExpressionRestorerInputs +from facefusion.processors.types import LivePortraitExpression, LivePortraitFeatureVolume, LivePortraitMotionPoints, LivePortraitPitch, LivePortraitRoll, LivePortraitScale, LivePortraitTranslation, LivePortraitYaw, ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore, thread_semaphore +from facefusion.types import ApplyStateItem, Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import read_static_image, read_static_video_frame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'live_portrait': + { + '__metadata__': + { + 'vendor': 'KwaiVGI', + 'license': 'MIT', + 'year': 2024 + }, + 'hashes': + { + 'feature_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_feature_extractor.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_feature_extractor.hash') + }, + 'motion_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_motion_extractor.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_motion_extractor.hash') + }, + 'generator': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_generator.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_generator.hash') + } + }, + 'sources': + { + 'feature_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_feature_extractor.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_feature_extractor.onnx') + }, + 'motion_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_motion_extractor.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_motion_extractor.onnx') + }, + 'generator': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_generator.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_generator.onnx') + } + }, + 'template': 'arcface_128', + 'size': (512, 512) + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('expression_restorer_model') ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('expression_restorer_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + model_name = state_manager.get_item('expression_restorer_model') + return create_static_model_set('full').get(model_name) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--expression-restorer-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'expression_restorer_model', 'live_portrait'), choices = expression_restorer_choices.expression_restorer_models) + group_processors.add_argument('--expression-restorer-factor', help = translator.get('help.factor', __package__), type = int, default = config.get_int_value('processors', 'expression_restorer_factor', '80'), choices = expression_restorer_choices.expression_restorer_factor_range, metavar = create_int_metavar(expression_restorer_choices.expression_restorer_factor_range)) + group_processors.add_argument('--expression-restorer-areas', help = translator.get('help.areas', __package__).format(choices = ', '.join(expression_restorer_choices.expression_restorer_areas)), default = config.get_str_list('processors', 'expression_restorer_areas', ' '.join(expression_restorer_choices.expression_restorer_areas)), choices = expression_restorer_choices.expression_restorer_areas, nargs = '+', metavar = 'EXPRESSION_RESTORER_AREAS') + facefusion.jobs.job_store.register_step_keys([ 'expression_restorer_model', 'expression_restorer_factor', 'expression_restorer_areas' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('expression_restorer_model', args.get('expression_restorer_model')) + apply_state_item('expression_restorer_factor', args.get('expression_restorer_factor')) + apply_state_item('expression_restorer_areas', args.get('expression_restorer_areas')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if mode == 'stream': + logger.error(translator.get('stream_not_supported') + translator.get('exclamation_mark'), __name__) + return False + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_masker.clear_inference_pool() + face_recognizer.clear_inference_pool() + + +def restore_expression(target_face : Face, target_vision_frame : VisionFrame, temp_vision_frame : VisionFrame) -> VisionFrame: + model_template = get_model_options().get('template') + model_size = get_model_options().get('size') + expression_restorer_factor = float(numpy.interp(float(state_manager.get_item('expression_restorer_factor')), [ 0, 100 ], [ 0, 1.2 ])) + target_crop_vision_frame, _ = warp_face_by_face_landmark_5(target_vision_frame, target_face.landmark_set.get('5/68'), model_template, model_size) + temp_crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmark_set.get('5/68'), model_template, model_size) + box_mask = create_box_mask(temp_crop_vision_frame, state_manager.get_item('face_mask_blur'), (0, 0, 0, 0)) + crop_masks =\ + [ + box_mask + ] + + if 'occlusion' in state_manager.get_item('face_mask_types'): + occlusion_mask = create_occlusion_mask(temp_crop_vision_frame) + crop_masks.append(occlusion_mask) + + target_crop_vision_frame = prepare_crop_frame(target_crop_vision_frame) + temp_crop_vision_frame = prepare_crop_frame(temp_crop_vision_frame) + temp_crop_vision_frame = apply_restore(target_crop_vision_frame, temp_crop_vision_frame, expression_restorer_factor) + temp_crop_vision_frame = normalize_crop_frame(temp_crop_vision_frame) + crop_mask = numpy.minimum.reduce(crop_masks).clip(0, 1) + paste_vision_frame = paste_back(temp_vision_frame, temp_crop_vision_frame, crop_mask, affine_matrix) + return paste_vision_frame + + +def apply_restore(target_crop_vision_frame : VisionFrame, temp_crop_vision_frame : VisionFrame, expression_restorer_factor : float) -> VisionFrame: + feature_volume = forward_extract_feature(temp_crop_vision_frame) + target_expression = forward_extract_motion(target_crop_vision_frame)[5] + pitch, yaw, roll, scale, translation, temp_expression, motion_points = forward_extract_motion(temp_crop_vision_frame) + rotation = create_rotation(pitch, yaw, roll) + target_expression = restrict_expression_areas(temp_expression, target_expression) + target_expression = target_expression * expression_restorer_factor + temp_expression * (1 - expression_restorer_factor) + target_expression = limit_expression(target_expression) + target_motion_points = scale * (motion_points @ rotation.T + target_expression) + translation + temp_motion_points = scale * (motion_points @ rotation.T + temp_expression) + translation + crop_vision_frame = forward_generate_frame(feature_volume, target_motion_points, temp_motion_points) + return crop_vision_frame + + +def restrict_expression_areas(temp_expression : LivePortraitExpression, target_expression : LivePortraitExpression) -> LivePortraitExpression: + expression_restorer_areas = state_manager.get_item('expression_restorer_areas') + + if 'upper-face' not in expression_restorer_areas: + target_expression[:, [1, 2, 6, 10, 11, 12, 13, 15, 16]] = temp_expression[:, [1, 2, 6, 10, 11, 12, 13, 15, 16]] + + if 'lower-face' not in expression_restorer_areas: + target_expression[:, [3, 7, 14, 17, 18, 19, 20]] = temp_expression[:, [3, 7, 14, 17, 18, 19, 20]] + + target_expression[:, [0, 4, 5, 8, 9]] = temp_expression[:, [0, 4, 5, 8, 9]] + return target_expression + + +def forward_extract_feature(crop_vision_frame : VisionFrame) -> LivePortraitFeatureVolume: + feature_extractor = get_inference_pool().get('feature_extractor') + + with conditional_thread_semaphore(): + feature_volume = feature_extractor.run(None, + { + 'input': crop_vision_frame + })[0] + + return feature_volume + + +def forward_extract_motion(crop_vision_frame : VisionFrame) -> Tuple[LivePortraitPitch, LivePortraitYaw, LivePortraitRoll, LivePortraitScale, LivePortraitTranslation, LivePortraitExpression, LivePortraitMotionPoints]: + motion_extractor = get_inference_pool().get('motion_extractor') + + with conditional_thread_semaphore(): + pitch, yaw, roll, scale, translation, expression, motion_points = motion_extractor.run(None, + { + 'input': crop_vision_frame + }) + + return pitch, yaw, roll, scale, translation, expression, motion_points + + +def forward_generate_frame(feature_volume : LivePortraitFeatureVolume, target_motion_points : LivePortraitMotionPoints, temp_motion_points : LivePortraitMotionPoints) -> VisionFrame: + generator = get_inference_pool().get('generator') + + with thread_semaphore(): + crop_vision_frame = generator.run(None, + { + 'feature_volume': feature_volume, + 'source': target_motion_points, + 'target': temp_motion_points + })[0][0] + + return crop_vision_frame + + +def prepare_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + model_size = get_model_options().get('size') + prepare_size = (model_size[0] // 2, model_size[1] // 2) + crop_vision_frame = cv2.resize(crop_vision_frame, prepare_size, interpolation = cv2.INTER_AREA) + crop_vision_frame = crop_vision_frame[:, :, ::-1] / 255.0 + crop_vision_frame = numpy.expand_dims(crop_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) + return crop_vision_frame + + +def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + crop_vision_frame = crop_vision_frame.transpose(1, 2, 0).clip(0, 1) + crop_vision_frame = crop_vision_frame * 255.0 + crop_vision_frame = crop_vision_frame.astype(numpy.uint8)[:, :, ::-1] + return crop_vision_frame + + +def process_frame(inputs : ExpressionRestorerInputs) -> ProcessorOutputs: + reference_vision_frame = inputs.get('reference_vision_frame') + target_vision_frame = inputs.get('target_vision_frame') + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + target_faces = select_faces(reference_vision_frame, target_vision_frame) + + if target_faces: + for target_face in target_faces: + target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) + temp_vision_frame = restore_expression(target_face, target_vision_frame, temp_vision_frame) + + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/expression_restorer/locales.py b/facefusion/processors/modules/expression_restorer/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..113efd35654a6a9f48e31125563fdde7e86caf63 --- /dev/null +++ b/facefusion/processors/modules/expression_restorer/locales.py @@ -0,0 +1,20 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for restoring the expression', + 'factor': 'restore factor of expression from the target face', + 'areas': 'choose the items used for the expression areas (choices: {choices})' + }, + 'uis': + { + 'model_dropdown': 'EXPRESSION RESTORER MODEL', + 'factor_slider': 'EXPRESSION RESTORER FACTOR', + 'areas_checkbox_group': 'EXPRESSION RESTORER AREAS' + } + } +} diff --git a/facefusion/processors/modules/expression_restorer/types.py b/facefusion/processors/modules/expression_restorer/types.py new file mode 100644 index 0000000000000000000000000000000000000000..04e7696bf43ee7998c802aa8ce7e626a6ef5f3a8 --- /dev/null +++ b/facefusion/processors/modules/expression_restorer/types.py @@ -0,0 +1,16 @@ +from typing import List, Literal, TypedDict + +from facefusion.types import Mask, VisionFrame + +ExpressionRestorerInputs = TypedDict('ExpressionRestorerInputs', +{ + 'reference_vision_frame' : VisionFrame, + 'source_vision_frames' : List[VisionFrame], + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +ExpressionRestorerModel = Literal['live_portrait'] + +ExpressionRestorerArea = Literal['upper-face', 'lower-face'] diff --git a/facefusion/processors/modules/face_debugger/__init__.py b/facefusion/processors/modules/face_debugger/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/face_debugger/choices.py b/facefusion/processors/modules/face_debugger/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..073ffc051a93edd5ee377723771fdeb089790f09 --- /dev/null +++ b/facefusion/processors/modules/face_debugger/choices.py @@ -0,0 +1,5 @@ +from typing import List + +from facefusion.processors.modules.face_debugger.types import FaceDebuggerItem + +face_debugger_items : List[FaceDebuggerItem] = [ 'bounding-box', 'face-landmark-5', 'face-landmark-5/68', 'face-landmark-68', 'face-landmark-68/5', 'face-mask' ] diff --git a/facefusion/processors/modules/face_debugger/core.py b/facefusion/processors/modules/face_debugger/core.py new file mode 100644 index 0000000000000000000000000000000000000000..45b4beb1e85835664598d4571fa87145eb25d86f --- /dev/null +++ b/facefusion/processors/modules/face_debugger/core.py @@ -0,0 +1,235 @@ +from argparse import ArgumentParser + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, logger, state_manager, translator, video_manager +from facefusion.face_analyser import scale_face +from facefusion.face_helper import warp_face_by_face_landmark_5 +from facefusion.face_masker import create_area_mask, create_box_mask, create_occlusion_mask, create_region_mask +from facefusion.face_selector import select_faces +from facefusion.filesystem import in_directory, is_image, is_video, same_file_extension +from facefusion.processors.modules.face_debugger import choices as face_debugger_choices +from facefusion.processors.modules.face_debugger.types import FaceDebuggerInputs +from facefusion.processors.types import ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.types import ApplyStateItem, Args, Face, InferencePool, ProcessMode, VisionFrame +from facefusion.vision import read_static_image, read_static_video_frame + + +def get_inference_pool() -> InferencePool: + pass + + +def clear_inference_pool() -> None: + pass + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--face-debugger-items', help = translator.get('help.items', __package__).format(choices = ', '.join(face_debugger_choices.face_debugger_items)), default = config.get_str_list('processors', 'face_debugger_items', 'face-landmark-5/68 face-mask'), choices = face_debugger_choices.face_debugger_items, nargs = '+', metavar = 'FACE_DEBUGGER_ITEMS') + facefusion.jobs.job_store.register_step_keys([ 'face_debugger_items' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('face_debugger_items', args.get('face_debugger_items')) + + +def pre_check() -> bool: + return True + + +def pre_process(mode : ProcessMode) -> bool: + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_masker.clear_inference_pool() + face_recognizer.clear_inference_pool() + + +def debug_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + face_debugger_items = state_manager.get_item('face_debugger_items') + + if 'bounding-box' in face_debugger_items: + temp_vision_frame = draw_bounding_box(target_face, temp_vision_frame) + + if 'face-mask' in face_debugger_items: + temp_vision_frame = draw_face_mask(target_face, temp_vision_frame) + + if 'face-landmark-5' in face_debugger_items: + temp_vision_frame = draw_face_landmark_5(target_face, temp_vision_frame) + + if 'face-landmark-5/68' in face_debugger_items: + temp_vision_frame = draw_face_landmark_5_68(target_face, temp_vision_frame) + + if 'face-landmark-68' in face_debugger_items: + temp_vision_frame = draw_face_landmark_68(target_face, temp_vision_frame) + + if 'face-landmark-68/5' in face_debugger_items: + temp_vision_frame = draw_face_landmark_68_5(target_face, temp_vision_frame) + + return temp_vision_frame + + +def draw_bounding_box(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) + box_color = 0, 0, 255 + border_color = 100, 100, 255 + bounding_box = target_face.bounding_box.astype(numpy.int32) + x1, y1, x2, y2 = bounding_box + + cv2.rectangle(temp_vision_frame, (x1, y1), (x2, y2), box_color, 2) + + if target_face.angle == 0: + cv2.line(temp_vision_frame, (x1, y1), (x2, y1), border_color, 3) + if target_face.angle == 180: + cv2.line(temp_vision_frame, (x1, y2), (x2, y2), border_color, 3) + if target_face.angle == 90: + cv2.line(temp_vision_frame, (x2, y1), (x2, y2), border_color, 3) + if target_face.angle == 270: + cv2.line(temp_vision_frame, (x1, y1), (x1, y2), border_color, 3) + + return temp_vision_frame + + +def draw_face_mask(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + crop_masks = [] + temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) + face_landmark_5 = target_face.landmark_set.get('5') + face_landmark_68 = target_face.landmark_set.get('68') + face_landmark_5_68 = target_face.landmark_set.get('5/68') + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5_68, 'arcface_128', (512, 512)) + inverse_matrix = cv2.invertAffineTransform(affine_matrix) + temp_size = temp_vision_frame.shape[:2][::-1] + mask_color = 0, 255, 0 + + if numpy.array_equal(face_landmark_5, face_landmark_5_68): + mask_color = 255, 255, 0 + + if 'box' in state_manager.get_item('face_mask_types'): + box_mask = create_box_mask(crop_vision_frame, 0, state_manager.get_item('face_mask_padding')) + crop_masks.append(box_mask) + + if 'occlusion' in state_manager.get_item('face_mask_types'): + occlusion_mask = create_occlusion_mask(crop_vision_frame) + crop_masks.append(occlusion_mask) + + if 'area' in state_manager.get_item('face_mask_types'): + face_landmark_68 = cv2.transform(face_landmark_68.reshape(1, -1, 2), affine_matrix).reshape(-1, 2) + area_mask = create_area_mask(crop_vision_frame, face_landmark_68, state_manager.get_item('face_mask_areas')) + crop_masks.append(area_mask) + + if 'region' in state_manager.get_item('face_mask_types'): + region_mask = create_region_mask(crop_vision_frame, state_manager.get_item('face_mask_regions')) + crop_masks.append(region_mask) + + crop_mask = numpy.minimum.reduce(crop_masks).clip(0, 1) + crop_mask = (crop_mask * 255).astype(numpy.uint8) + inverse_vision_frame = cv2.warpAffine(crop_mask, inverse_matrix, temp_size) + inverse_vision_frame = cv2.threshold(inverse_vision_frame, 100, 255, cv2.THRESH_BINARY)[1] + inverse_contours, _ = cv2.findContours(inverse_vision_frame, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE) + cv2.drawContours(temp_vision_frame, inverse_contours, -1, mask_color, 2) + + return temp_vision_frame + + +def draw_face_landmark_5(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) + face_landmark_5 = target_face.landmark_set.get('5') + point_color = 0, 0, 255 + + if numpy.any(face_landmark_5): + face_landmark_5 = face_landmark_5.astype(numpy.int32) + + for point in face_landmark_5: + cv2.circle(temp_vision_frame, tuple(point), 3, point_color, -1) + + return temp_vision_frame + + +def draw_face_landmark_5_68(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) + face_landmark_5 = target_face.landmark_set.get('5') + face_landmark_5_68 = target_face.landmark_set.get('5/68') + point_color = 0, 255, 0 + + if numpy.array_equal(face_landmark_5, face_landmark_5_68): + point_color = 255, 255, 0 + + if numpy.any(face_landmark_5_68): + face_landmark_5_68 = face_landmark_5_68.astype(numpy.int32) + + for point in face_landmark_5_68: + cv2.circle(temp_vision_frame, tuple(point), 3, point_color, -1) + + return temp_vision_frame + + +def draw_face_landmark_68(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) + face_landmark_68 = target_face.landmark_set.get('68') + face_landmark_68_5 = target_face.landmark_set.get('68/5') + point_color = 0, 255, 0 + + if numpy.array_equal(face_landmark_68, face_landmark_68_5): + point_color = 255, 255, 0 + + if numpy.any(face_landmark_68): + face_landmark_68 = face_landmark_68.astype(numpy.int32) + + for point in face_landmark_68: + cv2.circle(temp_vision_frame, tuple(point), 3, point_color, -1) + + return temp_vision_frame + + +def draw_face_landmark_68_5(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) + face_landmark_68_5 = target_face.landmark_set.get('68/5') + point_color = 255, 255, 0 + + if numpy.any(face_landmark_68_5): + face_landmark_68_5 = face_landmark_68_5.astype(numpy.int32) + + for point in face_landmark_68_5: + cv2.circle(temp_vision_frame, tuple(point), 3, point_color, -1) + + return temp_vision_frame + + +def process_frame(inputs : FaceDebuggerInputs) -> ProcessorOutputs: + reference_vision_frame = inputs.get('reference_vision_frame') + target_vision_frame = inputs.get('target_vision_frame') + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + target_faces = select_faces(reference_vision_frame, target_vision_frame) + + if target_faces: + for target_face in target_faces: + target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) + temp_vision_frame = debug_face(target_face, temp_vision_frame) + + return temp_vision_frame, temp_vision_mask + + diff --git a/facefusion/processors/modules/face_debugger/locales.py b/facefusion/processors/modules/face_debugger/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..0197db379dd118721bf56f286dbce76345f366fb --- /dev/null +++ b/facefusion/processors/modules/face_debugger/locales.py @@ -0,0 +1,16 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'items': 'load a single or multiple processors (choices: {choices})' + }, + 'uis': + { + 'items_checkbox_group': 'FACE DEBUGGER ITEMS' + } + } +} diff --git a/facefusion/processors/modules/face_debugger/types.py b/facefusion/processors/modules/face_debugger/types.py new file mode 100644 index 0000000000000000000000000000000000000000..c8eed1f41f9358e8e3b8e35f283e3b9b45566227 --- /dev/null +++ b/facefusion/processors/modules/face_debugger/types.py @@ -0,0 +1,13 @@ +from typing import Literal, TypedDict + +from facefusion.types import Mask, VisionFrame + +FaceDebuggerInputs = TypedDict('FaceDebuggerInputs', +{ + 'reference_vision_frame' : VisionFrame, + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +FaceDebuggerItem = Literal['bounding-box', 'face-landmark-5', 'face-landmark-5/68', 'face-landmark-68', 'face-landmark-68/5', 'face-mask'] diff --git a/facefusion/processors/modules/face_editor/__init__.py b/facefusion/processors/modules/face_editor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/face_editor/choices.py b/facefusion/processors/modules/face_editor/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..1db3aaae4c056a14d04464c0082e594f978b2948 --- /dev/null +++ b/facefusion/processors/modules/face_editor/choices.py @@ -0,0 +1,21 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_float_range +from facefusion.processors.modules.face_editor.types import FaceEditorModel + +face_editor_models : List[FaceEditorModel] = [ 'live_portrait' ] + +face_editor_eyebrow_direction_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_eye_gaze_horizontal_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_eye_gaze_vertical_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_eye_open_ratio_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_lip_open_ratio_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_mouth_grim_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_mouth_pout_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_mouth_purse_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_mouth_smile_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_mouth_position_horizontal_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_mouth_position_vertical_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_head_pitch_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_head_yaw_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) +face_editor_head_roll_range : Sequence[float] = create_float_range(-1.0, 1.0, 0.05) diff --git a/facefusion/processors/modules/face_editor/core.py b/facefusion/processors/modules/face_editor/core.py new file mode 100644 index 0000000000000000000000000000000000000000..e848b92035bc687dacab73800894aaae80dba9ba --- /dev/null +++ b/facefusion/processors/modules/face_editor/core.py @@ -0,0 +1,499 @@ +from argparse import ArgumentParser +from functools import lru_cache +from typing import Tuple + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import create_float_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.face_analyser import scale_face +from facefusion.face_helper import paste_back, scale_face_landmark_5, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_box_mask +from facefusion.face_selector import select_faces +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors.live_portrait import create_rotation, limit_angle, limit_expression +from facefusion.processors.modules.face_editor import choices as face_editor_choices +from facefusion.processors.modules.face_editor.types import FaceEditorInputs +from facefusion.processors.types import LivePortraitExpression, LivePortraitFeatureVolume, LivePortraitMotionPoints, LivePortraitPitch, LivePortraitRoll, LivePortraitRotation, LivePortraitScale, LivePortraitTranslation, LivePortraitYaw, ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore, thread_semaphore +from facefusion.types import ApplyStateItem, Args, DownloadScope, Face, FaceLandmark68, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import read_static_image, read_static_video_frame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'live_portrait': + { + '__metadata__': + { + 'vendor': 'KwaiVGI', + 'license': 'MIT', + 'year': 2024 + }, + 'hashes': + { + 'feature_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_feature_extractor.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_feature_extractor.hash') + }, + 'motion_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_motion_extractor.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_motion_extractor.hash') + }, + 'eye_retargeter': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_eye_retargeter.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_eye_retargeter.hash') + }, + 'lip_retargeter': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_lip_retargeter.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_lip_retargeter.hash') + }, + 'stitcher': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_stitcher.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_stitcher.hash') + }, + 'generator': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_generator.hash'), + 'path': resolve_relative_path('../.assets/models/live_portrait_generator.hash') + } + }, + 'sources': + { + 'feature_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_feature_extractor.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_feature_extractor.onnx') + }, + 'motion_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_motion_extractor.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_motion_extractor.onnx') + }, + 'eye_retargeter': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_eye_retargeter.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_eye_retargeter.onnx') + }, + 'lip_retargeter': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_lip_retargeter.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_lip_retargeter.onnx') + }, + 'stitcher': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_stitcher.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_stitcher.onnx') + }, + 'generator': + { + 'url': resolve_download_url('models-3.0.0', 'live_portrait_generator.onnx'), + 'path': resolve_relative_path('../.assets/models/live_portrait_generator.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('face_editor_model') ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('face_editor_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + model_name = state_manager.get_item('face_editor_model') + return create_static_model_set('full').get(model_name) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--face-editor-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'face_editor_model', 'live_portrait'), choices = face_editor_choices.face_editor_models) + group_processors.add_argument('--face-editor-eyebrow-direction', help = translator.get('help.eyebrow_direction', __package__), type = float, default = config.get_float_value('processors', 'face_editor_eyebrow_direction', '0'), choices = face_editor_choices.face_editor_eyebrow_direction_range, metavar = create_float_metavar(face_editor_choices.face_editor_eyebrow_direction_range)) + group_processors.add_argument('--face-editor-eye-gaze-horizontal', help = translator.get('help.eye_gaze_horizontal', __package__), type = float, default = config.get_float_value('processors', 'face_editor_eye_gaze_horizontal', '0'), choices = face_editor_choices.face_editor_eye_gaze_horizontal_range, metavar = create_float_metavar(face_editor_choices.face_editor_eye_gaze_horizontal_range)) + group_processors.add_argument('--face-editor-eye-gaze-vertical', help = translator.get('help.eye_gaze_vertical', __package__), type = float, default = config.get_float_value('processors', 'face_editor_eye_gaze_vertical', '0'), choices = face_editor_choices.face_editor_eye_gaze_vertical_range, metavar = create_float_metavar(face_editor_choices.face_editor_eye_gaze_vertical_range)) + group_processors.add_argument('--face-editor-eye-open-ratio', help = translator.get('help.eye_open_ratio', __package__), type = float, default = config.get_float_value('processors', 'face_editor_eye_open_ratio', '0'), choices = face_editor_choices.face_editor_eye_open_ratio_range, metavar = create_float_metavar(face_editor_choices.face_editor_eye_open_ratio_range)) + group_processors.add_argument('--face-editor-lip-open-ratio', help = translator.get('help.lip_open_ratio', __package__), type = float, default = config.get_float_value('processors', 'face_editor_lip_open_ratio', '0'), choices = face_editor_choices.face_editor_lip_open_ratio_range, metavar = create_float_metavar(face_editor_choices.face_editor_lip_open_ratio_range)) + group_processors.add_argument('--face-editor-mouth-grim', help = translator.get('help.mouth_grim', __package__), type = float, default = config.get_float_value('processors', 'face_editor_mouth_grim', '0'), choices = face_editor_choices.face_editor_mouth_grim_range, metavar = create_float_metavar(face_editor_choices.face_editor_mouth_grim_range)) + group_processors.add_argument('--face-editor-mouth-pout', help = translator.get('help.mouth_pout', __package__), type = float, default = config.get_float_value('processors', 'face_editor_mouth_pout', '0'), choices = face_editor_choices.face_editor_mouth_pout_range, metavar = create_float_metavar(face_editor_choices.face_editor_mouth_pout_range)) + group_processors.add_argument('--face-editor-mouth-purse', help = translator.get('help.mouth_purse', __package__), type = float, default = config.get_float_value('processors', 'face_editor_mouth_purse', '0'), choices = face_editor_choices.face_editor_mouth_purse_range, metavar = create_float_metavar(face_editor_choices.face_editor_mouth_purse_range)) + group_processors.add_argument('--face-editor-mouth-smile', help = translator.get('help.mouth_smile', __package__), type = float, default = config.get_float_value('processors', 'face_editor_mouth_smile', '0'), choices = face_editor_choices.face_editor_mouth_smile_range, metavar = create_float_metavar(face_editor_choices.face_editor_mouth_smile_range)) + group_processors.add_argument('--face-editor-mouth-position-horizontal', help = translator.get('help.mouth_position_horizontal', __package__), type = float, default = config.get_float_value('processors', 'face_editor_mouth_position_horizontal', '0'), choices = face_editor_choices.face_editor_mouth_position_horizontal_range, metavar = create_float_metavar(face_editor_choices.face_editor_mouth_position_horizontal_range)) + group_processors.add_argument('--face-editor-mouth-position-vertical', help = translator.get('help.mouth_position_vertical', __package__), type = float, default = config.get_float_value('processors', 'face_editor_mouth_position_vertical', '0'), choices = face_editor_choices.face_editor_mouth_position_vertical_range, metavar = create_float_metavar(face_editor_choices.face_editor_mouth_position_vertical_range)) + group_processors.add_argument('--face-editor-head-pitch', help = translator.get('help.head_pitch', __package__), type = float, default = config.get_float_value('processors', 'face_editor_head_pitch', '0'), choices = face_editor_choices.face_editor_head_pitch_range, metavar = create_float_metavar(face_editor_choices.face_editor_head_pitch_range)) + group_processors.add_argument('--face-editor-head-yaw', help = translator.get('help.head_yaw', __package__), type = float, default = config.get_float_value('processors', 'face_editor_head_yaw', '0'), choices = face_editor_choices.face_editor_head_yaw_range, metavar = create_float_metavar(face_editor_choices.face_editor_head_yaw_range)) + group_processors.add_argument('--face-editor-head-roll', help = translator.get('help.head_roll', __package__), type = float, default = config.get_float_value('processors', 'face_editor_head_roll', '0'), choices = face_editor_choices.face_editor_head_roll_range, metavar = create_float_metavar(face_editor_choices.face_editor_head_roll_range)) + facefusion.jobs.job_store.register_step_keys([ 'face_editor_model', 'face_editor_eyebrow_direction', 'face_editor_eye_gaze_horizontal', 'face_editor_eye_gaze_vertical', 'face_editor_eye_open_ratio', 'face_editor_lip_open_ratio', 'face_editor_mouth_grim', 'face_editor_mouth_pout', 'face_editor_mouth_purse', 'face_editor_mouth_smile', 'face_editor_mouth_position_horizontal', 'face_editor_mouth_position_vertical', 'face_editor_head_pitch', 'face_editor_head_yaw', 'face_editor_head_roll' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('face_editor_model', args.get('face_editor_model')) + apply_state_item('face_editor_eyebrow_direction', args.get('face_editor_eyebrow_direction')) + apply_state_item('face_editor_eye_gaze_horizontal', args.get('face_editor_eye_gaze_horizontal')) + apply_state_item('face_editor_eye_gaze_vertical', args.get('face_editor_eye_gaze_vertical')) + apply_state_item('face_editor_eye_open_ratio', args.get('face_editor_eye_open_ratio')) + apply_state_item('face_editor_lip_open_ratio', args.get('face_editor_lip_open_ratio')) + apply_state_item('face_editor_mouth_grim', args.get('face_editor_mouth_grim')) + apply_state_item('face_editor_mouth_pout', args.get('face_editor_mouth_pout')) + apply_state_item('face_editor_mouth_purse', args.get('face_editor_mouth_purse')) + apply_state_item('face_editor_mouth_smile', args.get('face_editor_mouth_smile')) + apply_state_item('face_editor_mouth_position_horizontal', args.get('face_editor_mouth_position_horizontal')) + apply_state_item('face_editor_mouth_position_vertical', args.get('face_editor_mouth_position_vertical')) + apply_state_item('face_editor_head_pitch', args.get('face_editor_head_pitch')) + apply_state_item('face_editor_head_yaw', args.get('face_editor_head_yaw')) + apply_state_item('face_editor_head_roll', args.get('face_editor_head_roll')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_masker.clear_inference_pool() + face_recognizer.clear_inference_pool() + + +def edit_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + model_template = get_model_options().get('template') + model_size = get_model_options().get('size') + face_landmark_5 = scale_face_landmark_5(target_face.landmark_set.get('5/68'), 1.5) + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5, model_template, model_size) + box_mask = create_box_mask(crop_vision_frame, state_manager.get_item('face_mask_blur'), (0, 0, 0, 0)) + crop_vision_frame = prepare_crop_frame(crop_vision_frame) + crop_vision_frame = apply_edit(crop_vision_frame, target_face.landmark_set.get('68')) + crop_vision_frame = normalize_crop_frame(crop_vision_frame) + paste_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, box_mask, affine_matrix) + return paste_vision_frame + + +def apply_edit(crop_vision_frame : VisionFrame, face_landmark_68 : FaceLandmark68) -> VisionFrame: + feature_volume = forward_extract_feature(crop_vision_frame) + pitch, yaw, roll, scale, translation, expression, motion_points = forward_extract_motion(crop_vision_frame) + rotation = create_rotation(pitch, yaw, roll) + motion_points_target = scale * (motion_points @ rotation.T + expression) + translation + expression = edit_eye_gaze(expression) + expression = edit_mouth_grim(expression) + expression = edit_mouth_position(expression) + expression = edit_mouth_pout(expression) + expression = edit_mouth_purse(expression) + expression = edit_mouth_smile(expression) + expression = edit_eyebrow_direction(expression) + expression = limit_expression(expression) + rotation = edit_head_rotation(pitch, yaw, roll) + motion_points_source = motion_points @ rotation.T + motion_points_source += expression + motion_points_source *= scale + motion_points_source += translation + motion_points_source += edit_eye_open(motion_points_target, face_landmark_68) + motion_points_source += edit_lip_open(motion_points_target, face_landmark_68) + motion_points_source = forward_stitch_motion_points(motion_points_source, motion_points_target) + crop_vision_frame = forward_generate_frame(feature_volume, motion_points_source, motion_points_target) + return crop_vision_frame + + +def forward_extract_feature(crop_vision_frame : VisionFrame) -> LivePortraitFeatureVolume: + feature_extractor = get_inference_pool().get('feature_extractor') + + with conditional_thread_semaphore(): + feature_volume = feature_extractor.run(None, + { + 'input': crop_vision_frame + })[0] + + return feature_volume + + +def forward_extract_motion(crop_vision_frame : VisionFrame) -> Tuple[LivePortraitPitch, LivePortraitYaw, LivePortraitRoll, LivePortraitScale, LivePortraitTranslation, LivePortraitExpression, LivePortraitMotionPoints]: + motion_extractor = get_inference_pool().get('motion_extractor') + + with conditional_thread_semaphore(): + pitch, yaw, roll, scale, translation, expression, motion_points = motion_extractor.run(None, + { + 'input': crop_vision_frame + }) + + return pitch, yaw, roll, scale, translation, expression, motion_points + + +def forward_retarget_eye(eye_motion_points : LivePortraitMotionPoints) -> LivePortraitMotionPoints: + eye_retargeter = get_inference_pool().get('eye_retargeter') + + with conditional_thread_semaphore(): + eye_motion_points = eye_retargeter.run(None, + { + 'input': eye_motion_points + })[0] + + return eye_motion_points + + +def forward_retarget_lip(lip_motion_points : LivePortraitMotionPoints) -> LivePortraitMotionPoints: + lip_retargeter = get_inference_pool().get('lip_retargeter') + + with conditional_thread_semaphore(): + lip_motion_points = lip_retargeter.run(None, + { + 'input': lip_motion_points + })[0] + + return lip_motion_points + + +def forward_stitch_motion_points(source_motion_points : LivePortraitMotionPoints, target_motion_points : LivePortraitMotionPoints) -> LivePortraitMotionPoints: + stitcher = get_inference_pool().get('stitcher') + + with thread_semaphore(): + motion_points = stitcher.run(None, + { + 'source': source_motion_points, + 'target': target_motion_points + })[0] + + return motion_points + + +def forward_generate_frame(feature_volume : LivePortraitFeatureVolume, source_motion_points : LivePortraitMotionPoints, target_motion_points : LivePortraitMotionPoints) -> VisionFrame: + generator = get_inference_pool().get('generator') + + with thread_semaphore(): + crop_vision_frame = generator.run(None, + { + 'feature_volume': feature_volume, + 'source': source_motion_points, + 'target': target_motion_points + })[0][0] + + return crop_vision_frame + + +def edit_eyebrow_direction(expression : LivePortraitExpression) -> LivePortraitExpression: + face_editor_eyebrow = state_manager.get_item('face_editor_eyebrow_direction') + + if face_editor_eyebrow > 0: + expression[0, 1, 1] += numpy.interp(face_editor_eyebrow, [ -1, 1 ], [ -0.015, 0.015 ]) + expression[0, 2, 1] -= numpy.interp(face_editor_eyebrow, [ -1, 1 ], [ -0.020, 0.020 ]) + else: + expression[0, 1, 0] -= numpy.interp(face_editor_eyebrow, [ -1, 1 ], [ -0.015, 0.015 ]) + expression[0, 2, 0] += numpy.interp(face_editor_eyebrow, [ -1, 1 ], [ -0.020, 0.020 ]) + expression[0, 1, 1] += numpy.interp(face_editor_eyebrow, [ -1, 1 ], [ -0.005, 0.005 ]) + expression[0, 2, 1] -= numpy.interp(face_editor_eyebrow, [ -1, 1 ], [ -0.005, 0.005 ]) + return expression + + +def edit_eye_gaze(expression : LivePortraitExpression) -> LivePortraitExpression: + face_editor_eye_gaze_horizontal = state_manager.get_item('face_editor_eye_gaze_horizontal') + face_editor_eye_gaze_vertical = state_manager.get_item('face_editor_eye_gaze_vertical') + + if face_editor_eye_gaze_horizontal > 0: + expression[0, 11, 0] += numpy.interp(face_editor_eye_gaze_horizontal, [ -1, 1 ], [ -0.015, 0.015 ]) + expression[0, 15, 0] += numpy.interp(face_editor_eye_gaze_horizontal, [ -1, 1 ], [ -0.020, 0.020 ]) + else: + expression[0, 11, 0] += numpy.interp(face_editor_eye_gaze_horizontal, [ -1, 1 ], [ -0.020, 0.020 ]) + expression[0, 15, 0] += numpy.interp(face_editor_eye_gaze_horizontal, [ -1, 1 ], [ -0.015, 0.015 ]) + expression[0, 1, 1] += numpy.interp(face_editor_eye_gaze_vertical, [ -1, 1 ], [ -0.0025, 0.0025 ]) + expression[0, 2, 1] -= numpy.interp(face_editor_eye_gaze_vertical, [ -1, 1 ], [ -0.0025, 0.0025 ]) + expression[0, 11, 1] -= numpy.interp(face_editor_eye_gaze_vertical, [ -1, 1 ], [ -0.010, 0.010 ]) + expression[0, 13, 1] -= numpy.interp(face_editor_eye_gaze_vertical, [ -1, 1 ], [ -0.005, 0.005 ]) + expression[0, 15, 1] -= numpy.interp(face_editor_eye_gaze_vertical, [ -1, 1 ], [ -0.010, 0.010 ]) + expression[0, 16, 1] -= numpy.interp(face_editor_eye_gaze_vertical, [ -1, 1 ], [ -0.005, 0.005 ]) + return expression + + +def edit_eye_open(motion_points : LivePortraitMotionPoints, face_landmark_68 : FaceLandmark68) -> LivePortraitMotionPoints: + face_editor_eye_open_ratio = state_manager.get_item('face_editor_eye_open_ratio') + left_eye_ratio = calculate_distance_ratio(face_landmark_68, 37, 40, 39, 36) + right_eye_ratio = calculate_distance_ratio(face_landmark_68, 43, 46, 45, 42) + + if face_editor_eye_open_ratio < 0: + eye_motion_points = numpy.concatenate([ motion_points.ravel(), [ left_eye_ratio, right_eye_ratio, 0.0 ] ]) + else: + eye_motion_points = numpy.concatenate([ motion_points.ravel(), [ left_eye_ratio, right_eye_ratio, 0.6 ] ]) + eye_motion_points = eye_motion_points.reshape(1, -1).astype(numpy.float32) + eye_motion_points = forward_retarget_eye(eye_motion_points) * numpy.abs(face_editor_eye_open_ratio) + eye_motion_points = eye_motion_points.reshape(-1, 21, 3) + return eye_motion_points + + +def edit_lip_open(motion_points : LivePortraitMotionPoints, face_landmark_68 : FaceLandmark68) -> LivePortraitMotionPoints: + face_editor_lip_open_ratio = state_manager.get_item('face_editor_lip_open_ratio') + lip_ratio = calculate_distance_ratio(face_landmark_68, 62, 66, 54, 48) + + if face_editor_lip_open_ratio < 0: + lip_motion_points = numpy.concatenate([ motion_points.ravel(), [ lip_ratio, 0.0 ] ]) + else: + lip_motion_points = numpy.concatenate([ motion_points.ravel(), [ lip_ratio, 1.0 ] ]) + lip_motion_points = lip_motion_points.reshape(1, -1).astype(numpy.float32) + lip_motion_points = forward_retarget_lip(lip_motion_points) * numpy.abs(face_editor_lip_open_ratio) + lip_motion_points = lip_motion_points.reshape(-1, 21, 3) + return lip_motion_points + + +def edit_mouth_grim(expression : LivePortraitExpression) -> LivePortraitExpression: + face_editor_mouth_grim = state_manager.get_item('face_editor_mouth_grim') + if face_editor_mouth_grim > 0: + expression[0, 17, 2] -= numpy.interp(face_editor_mouth_grim, [ -1, 1 ], [ -0.005, 0.005 ]) + expression[0, 19, 2] += numpy.interp(face_editor_mouth_grim, [ -1, 1 ], [ -0.01, 0.01 ]) + expression[0, 20, 1] -= numpy.interp(face_editor_mouth_grim, [ -1, 1 ], [ -0.06, 0.06 ]) + expression[0, 20, 2] -= numpy.interp(face_editor_mouth_grim, [ -1, 1 ], [ -0.03, 0.03 ]) + else: + expression[0, 19, 1] -= numpy.interp(face_editor_mouth_grim, [ -1, 1 ], [ -0.05, 0.05 ]) + expression[0, 19, 2] -= numpy.interp(face_editor_mouth_grim, [ -1, 1 ], [ -0.02, 0.02 ]) + expression[0, 20, 2] -= numpy.interp(face_editor_mouth_grim, [ -1, 1 ], [ -0.03, 0.03 ]) + return expression + + +def edit_mouth_position(expression : LivePortraitExpression) -> LivePortraitExpression: + face_editor_mouth_position_horizontal = state_manager.get_item('face_editor_mouth_position_horizontal') + face_editor_mouth_position_vertical = state_manager.get_item('face_editor_mouth_position_vertical') + expression[0, 19, 0] += numpy.interp(face_editor_mouth_position_horizontal, [ -1, 1 ], [ -0.05, 0.05 ]) + expression[0, 20, 0] += numpy.interp(face_editor_mouth_position_horizontal, [ -1, 1 ], [ -0.04, 0.04 ]) + if face_editor_mouth_position_vertical > 0: + expression[0, 19, 1] -= numpy.interp(face_editor_mouth_position_vertical, [ -1, 1 ], [ -0.04, 0.04 ]) + expression[0, 20, 1] -= numpy.interp(face_editor_mouth_position_vertical, [ -1, 1 ], [ -0.02, 0.02 ]) + else: + expression[0, 19, 1] -= numpy.interp(face_editor_mouth_position_vertical, [ -1, 1 ], [ -0.05, 0.05 ]) + expression[0, 20, 1] -= numpy.interp(face_editor_mouth_position_vertical, [ -1, 1 ], [ -0.04, 0.04 ]) + return expression + + +def edit_mouth_pout(expression : LivePortraitExpression) -> LivePortraitExpression: + face_editor_mouth_pout = state_manager.get_item('face_editor_mouth_pout') + if face_editor_mouth_pout > 0: + expression[0, 19, 1] -= numpy.interp(face_editor_mouth_pout, [ -1, 1 ], [ -0.022, 0.022 ]) + expression[0, 19, 2] += numpy.interp(face_editor_mouth_pout, [ -1, 1 ], [ -0.025, 0.025 ]) + expression[0, 20, 2] -= numpy.interp(face_editor_mouth_pout, [ -1, 1 ], [ -0.002, 0.002 ]) + else: + expression[0, 19, 1] += numpy.interp(face_editor_mouth_pout, [ -1, 1 ], [ -0.022, 0.022 ]) + expression[0, 19, 2] += numpy.interp(face_editor_mouth_pout, [ -1, 1 ], [ -0.025, 0.025 ]) + expression[0, 20, 2] -= numpy.interp(face_editor_mouth_pout, [ -1, 1 ], [ -0.002, 0.002 ]) + return expression + + +def edit_mouth_purse(expression : LivePortraitExpression) -> LivePortraitExpression: + face_editor_mouth_purse = state_manager.get_item('face_editor_mouth_purse') + if face_editor_mouth_purse > 0: + expression[0, 19, 1] -= numpy.interp(face_editor_mouth_purse, [ -1, 1 ], [ -0.04, 0.04 ]) + expression[0, 19, 2] -= numpy.interp(face_editor_mouth_purse, [ -1, 1 ], [ -0.02, 0.02 ]) + else: + expression[0, 14, 1] -= numpy.interp(face_editor_mouth_purse, [ -1, 1 ], [ -0.02, 0.02 ]) + expression[0, 17, 2] += numpy.interp(face_editor_mouth_purse, [ -1, 1 ], [ -0.01, 0.01 ]) + expression[0, 19, 2] -= numpy.interp(face_editor_mouth_purse, [ -1, 1 ], [ -0.015, 0.015 ]) + expression[0, 20, 2] -= numpy.interp(face_editor_mouth_purse, [ -1, 1 ], [ -0.002, 0.002 ]) + return expression + + +def edit_mouth_smile(expression : LivePortraitExpression) -> LivePortraitExpression: + face_editor_mouth_smile = state_manager.get_item('face_editor_mouth_smile') + if face_editor_mouth_smile > 0: + expression[0, 20, 1] -= numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.015, 0.015 ]) + expression[0, 14, 1] -= numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.025, 0.025 ]) + expression[0, 17, 1] += numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.01, 0.01 ]) + expression[0, 17, 2] += numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.004, 0.004 ]) + expression[0, 3, 1] -= numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.0045, 0.0045 ]) + expression[0, 7, 1] -= numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.0045, 0.0045 ]) + else: + expression[0, 14, 1] -= numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.02, 0.02 ]) + expression[0, 17, 1] += numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.003, 0.003 ]) + expression[0, 19, 1] += numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.02, 0.02 ]) + expression[0, 19, 2] -= numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.005, 0.005 ]) + expression[0, 20, 2] += numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.01, 0.01 ]) + expression[0, 3, 1] += numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.0045, 0.0045 ]) + expression[0, 7, 1] += numpy.interp(face_editor_mouth_smile, [ -1, 1 ], [ -0.0045, 0.0045 ]) + return expression + + +def edit_head_rotation(pitch : LivePortraitPitch, yaw : LivePortraitYaw, roll : LivePortraitRoll) -> LivePortraitRotation: + face_editor_head_pitch = state_manager.get_item('face_editor_head_pitch') + face_editor_head_yaw = state_manager.get_item('face_editor_head_yaw') + face_editor_head_roll = state_manager.get_item('face_editor_head_roll') + edit_pitch = pitch + float(numpy.interp(face_editor_head_pitch, [ -1, 1 ], [ 20, -20 ])) + edit_yaw = yaw + float(numpy.interp(face_editor_head_yaw, [ -1, 1 ], [ 60, -60 ])) + edit_roll = roll + float(numpy.interp(face_editor_head_roll, [ -1, 1 ], [ -15, 15 ])) + edit_pitch, edit_yaw, edit_roll = limit_angle(pitch, yaw, roll, edit_pitch, edit_yaw, edit_roll) + rotation = create_rotation(edit_pitch, edit_yaw, edit_roll) + return rotation + + +def calculate_distance_ratio(face_landmark_68 : FaceLandmark68, top_index : int, bottom_index : int, left_index : int, right_index : int) -> float: + vertical_direction = face_landmark_68[top_index] - face_landmark_68[bottom_index] + horizontal_direction = face_landmark_68[left_index] - face_landmark_68[right_index] + distance_ratio = float(numpy.linalg.norm(vertical_direction) / (numpy.linalg.norm(horizontal_direction) + 1e-6)) + return distance_ratio + + +def prepare_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + model_size = get_model_options().get('size') + prepare_size = (model_size[0] // 2, model_size[1] // 2) + crop_vision_frame = cv2.resize(crop_vision_frame, prepare_size, interpolation = cv2.INTER_AREA) + crop_vision_frame = crop_vision_frame[:, :, ::-1] / 255.0 + crop_vision_frame = numpy.expand_dims(crop_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) + return crop_vision_frame + + +def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + crop_vision_frame = crop_vision_frame.transpose(1, 2, 0).clip(0, 1) + crop_vision_frame = (crop_vision_frame * 255.0) + crop_vision_frame = crop_vision_frame.astype(numpy.uint8)[:, :, ::-1] + return crop_vision_frame + + +def process_frame(inputs : FaceEditorInputs) -> ProcessorOutputs: + reference_vision_frame = inputs.get('reference_vision_frame') + target_vision_frame = inputs.get('target_vision_frame') + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + target_faces = select_faces(reference_vision_frame, target_vision_frame) + + if target_faces: + for target_face in target_faces: + target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) + temp_vision_frame = edit_face(target_face, temp_vision_frame) + + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/face_editor/locales.py b/facefusion/processors/modules/face_editor/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..a951702b6072222b3a3ed73bbf64753be328d907 --- /dev/null +++ b/facefusion/processors/modules/face_editor/locales.py @@ -0,0 +1,44 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for editing the face', + 'eyebrow_direction': 'specify the eyebrow direction', + 'eye_gaze_horizontal': 'specify the horizontal eye gaze', + 'eye_gaze_vertical': 'specify the vertical eye gaze', + 'eye_open_ratio': 'specify the ratio of eye opening', + 'lip_open_ratio': 'specify the ratio of lip opening', + 'mouth_grim': 'specify the mouth grim', + 'mouth_pout': 'specify the mouth pout', + 'mouth_purse': 'specify the mouth purse', + 'mouth_smile': 'specify the mouth smile', + 'mouth_position_horizontal': 'specify the horizontal mouth position', + 'mouth_position_vertical': 'specify the vertical mouth position', + 'head_pitch': 'specify the head pitch', + 'head_yaw': 'specify the head yaw', + 'head_roll': 'specify the head roll' + }, + 'uis': + { + 'eyebrow_direction_slider': 'FACE EDITOR EYEBROW DIRECTION', + 'eye_gaze_horizontal_slider': 'FACE EDITOR EYE GAZE HORIZONTAL', + 'eye_gaze_vertical_slider': 'FACE EDITOR EYE GAZE VERTICAL', + 'eye_open_ratio_slider': 'FACE EDITOR EYE OPEN RATIO', + 'head_pitch_slider': 'FACE EDITOR HEAD PITCH', + 'head_roll_slider': 'FACE EDITOR HEAD ROLL', + 'head_yaw_slider': 'FACE EDITOR HEAD YAW', + 'lip_open_ratio_slider': 'FACE EDITOR LIP OPEN RATIO', + 'model_dropdown': 'FACE EDITOR MODEL', + 'mouth_grim_slider': 'FACE EDITOR MOUTH GRIM', + 'mouth_position_horizontal_slider': 'FACE EDITOR MOUTH POSITION HORIZONTAL', + 'mouth_position_vertical_slider': 'FACE EDITOR MOUTH POSITION VERTICAL', + 'mouth_pout_slider': 'FACE EDITOR MOUTH POUT', + 'mouth_purse_slider': 'FACE EDITOR MOUTH PURSE', + 'mouth_smile_slider': 'FACE EDITOR MOUTH SMILE' + } + } +} diff --git a/facefusion/processors/modules/face_editor/types.py b/facefusion/processors/modules/face_editor/types.py new file mode 100644 index 0000000000000000000000000000000000000000..012a8507155863f48adb35a8a65630f0e505ddf9 --- /dev/null +++ b/facefusion/processors/modules/face_editor/types.py @@ -0,0 +1,13 @@ +from typing import Literal, TypedDict + +from facefusion.types import Mask, VisionFrame + +FaceEditorInputs = TypedDict('FaceEditorInputs', +{ + 'reference_vision_frame' : VisionFrame, + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +FaceEditorModel = Literal['live_portrait'] diff --git a/facefusion/processors/modules/face_enhancer/__init__.py b/facefusion/processors/modules/face_enhancer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/face_enhancer/choices.py b/facefusion/processors/modules/face_enhancer/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..c4e3fdd8d0627501fd61efe39275f9770eae99a6 --- /dev/null +++ b/facefusion/processors/modules/face_enhancer/choices.py @@ -0,0 +1,10 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_float_range, create_int_range +from facefusion.processors.modules.face_enhancer.types import FaceEnhancerModel + +face_enhancer_models : List[FaceEnhancerModel] = [ 'codeformer', 'gfpgan_1.2', 'gfpgan_1.3', 'gfpgan_1.4', 'gpen_bfr_256', 'gpen_bfr_512', 'gpen_bfr_1024', 'gpen_bfr_2048', 'restoreformer_plus_plus' ] + +face_enhancer_blend_range : Sequence[int] = create_int_range(0, 100, 1) + +face_enhancer_weight_range : Sequence[float] = create_float_range(0.0, 1.0, 0.05) diff --git a/facefusion/processors/modules/face_enhancer/core.py b/facefusion/processors/modules/face_enhancer/core.py new file mode 100644 index 0000000000000000000000000000000000000000..f638d3b1f600f2ee643057126231eff9fe073ded --- /dev/null +++ b/facefusion/processors/modules/face_enhancer/core.py @@ -0,0 +1,426 @@ +from argparse import ArgumentParser +from functools import lru_cache + +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import create_float_metavar, create_int_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.face_analyser import scale_face +from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_box_mask, create_occlusion_mask +from facefusion.face_selector import select_faces +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors.modules.face_enhancer import choices as face_enhancer_choices +from facefusion.processors.modules.face_enhancer.types import FaceEnhancerInputs, FaceEnhancerWeight +from facefusion.processors.types import ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import thread_semaphore +from facefusion.types import ApplyStateItem, Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import blend_frame, read_static_image, read_static_video_frame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'codeformer': + { + '__metadata__': + { + 'vendor': 'sczhou', + 'license': 'S-Lab-1.0', + 'year': 2022 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'codeformer.hash'), + 'path': resolve_relative_path('../.assets/models/codeformer.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'codeformer.onnx'), + 'path': resolve_relative_path('../.assets/models/codeformer.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + }, + 'gfpgan_1.2': + { + '__metadata__': + { + 'vendor': 'TencentARC', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gfpgan_1.2.hash'), + 'path': resolve_relative_path('../.assets/models/gfpgan_1.2.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gfpgan_1.2.onnx'), + 'path': resolve_relative_path('../.assets/models/gfpgan_1.2.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + }, + 'gfpgan_1.3': + { + '__metadata__': + { + 'vendor': 'TencentARC', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gfpgan_1.3.hash'), + 'path': resolve_relative_path('../.assets/models/gfpgan_1.3.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gfpgan_1.3.onnx'), + 'path': resolve_relative_path('../.assets/models/gfpgan_1.3.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + }, + 'gfpgan_1.4': + { + '__metadata__': + { + 'vendor': 'TencentARC', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gfpgan_1.4.hash'), + 'path': resolve_relative_path('../.assets/models/gfpgan_1.4.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gfpgan_1.4.onnx'), + 'path': resolve_relative_path('../.assets/models/gfpgan_1.4.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + }, + 'gpen_bfr_256': + { + '__metadata__': + { + 'vendor': 'yangxy', + 'license': 'Non-Commercial', + 'year': 2021 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gpen_bfr_256.hash'), + 'path': resolve_relative_path('../.assets/models/gpen_bfr_256.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gpen_bfr_256.onnx'), + 'path': resolve_relative_path('../.assets/models/gpen_bfr_256.onnx') + } + }, + 'template': 'arcface_128', + 'size': (256, 256) + }, + 'gpen_bfr_512': + { + '__metadata__': + { + 'vendor': 'yangxy', + 'license': 'Non-Commercial', + 'year': 2021 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gpen_bfr_512.hash'), + 'path': resolve_relative_path('../.assets/models/gpen_bfr_512.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gpen_bfr_512.onnx'), + 'path': resolve_relative_path('../.assets/models/gpen_bfr_512.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + }, + 'gpen_bfr_1024': + { + '__metadata__': + { + 'vendor': 'yangxy', + 'license': 'Non-Commercial', + 'year': 2021 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gpen_bfr_1024.hash'), + 'path': resolve_relative_path('../.assets/models/gpen_bfr_1024.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gpen_bfr_1024.onnx'), + 'path': resolve_relative_path('../.assets/models/gpen_bfr_1024.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (1024, 1024) + }, + 'gpen_bfr_2048': + { + '__metadata__': + { + 'vendor': 'yangxy', + 'license': 'Non-Commercial', + 'year': 2021 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gpen_bfr_2048.hash'), + 'path': resolve_relative_path('../.assets/models/gpen_bfr_2048.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'gpen_bfr_2048.onnx'), + 'path': resolve_relative_path('../.assets/models/gpen_bfr_2048.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (2048, 2048) + }, + 'restoreformer_plus_plus': + { + '__metadata__': + { + 'vendor': 'wzhouxiff', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'restoreformer_plus_plus.hash'), + 'path': resolve_relative_path('../.assets/models/restoreformer_plus_plus.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'restoreformer_plus_plus.onnx'), + 'path': resolve_relative_path('../.assets/models/restoreformer_plus_plus.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('face_enhancer_model') ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('face_enhancer_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + model_name = state_manager.get_item('face_enhancer_model') + return create_static_model_set('full').get(model_name) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--face-enhancer-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'face_enhancer_model', 'gfpgan_1.4'), choices = face_enhancer_choices.face_enhancer_models) + group_processors.add_argument('--face-enhancer-blend', help = translator.get('help.blend', __package__), type = int, default = config.get_int_value('processors', 'face_enhancer_blend', '80'), choices = face_enhancer_choices.face_enhancer_blend_range, metavar = create_int_metavar(face_enhancer_choices.face_enhancer_blend_range)) + group_processors.add_argument('--face-enhancer-weight', help = translator.get('help.weight', __package__), type = float, default = config.get_float_value('processors', 'face_enhancer_weight', '0.5'), choices = face_enhancer_choices.face_enhancer_weight_range, metavar = create_float_metavar(face_enhancer_choices.face_enhancer_weight_range)) + facefusion.jobs.job_store.register_step_keys([ 'face_enhancer_model', 'face_enhancer_blend', 'face_enhancer_weight' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('face_enhancer_model', args.get('face_enhancer_model')) + apply_state_item('face_enhancer_blend', args.get('face_enhancer_blend')) + apply_state_item('face_enhancer_weight', args.get('face_enhancer_weight')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_masker.clear_inference_pool() + face_recognizer.clear_inference_pool() + + +def enhance_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + model_template = get_model_options().get('template') + model_size = get_model_options().get('size') + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmark_set.get('5/68'), model_template, model_size) + box_mask = create_box_mask(crop_vision_frame, state_manager.get_item('face_mask_blur'), (0, 0, 0, 0)) + crop_masks =\ + [ + box_mask + ] + + if 'occlusion' in state_manager.get_item('face_mask_types'): + occlusion_mask = create_occlusion_mask(crop_vision_frame) + crop_masks.append(occlusion_mask) + + crop_vision_frame = prepare_crop_frame(crop_vision_frame) + face_enhancer_weight = numpy.array([ state_manager.get_item('face_enhancer_weight') ]).astype(numpy.double) + crop_vision_frame = forward(crop_vision_frame, face_enhancer_weight) + crop_vision_frame = normalize_crop_frame(crop_vision_frame) + crop_mask = numpy.minimum.reduce(crop_masks).clip(0, 1) + paste_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, crop_mask, affine_matrix) + temp_vision_frame = blend_paste_frame(temp_vision_frame, paste_vision_frame) + return temp_vision_frame + + +def forward(crop_vision_frame : VisionFrame, face_enhancer_weight : FaceEnhancerWeight) -> VisionFrame: + face_enhancer = get_inference_pool().get('face_enhancer') + face_enhancer_inputs = {} + + for face_enhancer_input in face_enhancer.get_inputs(): + if face_enhancer_input.name == 'input': + face_enhancer_inputs[face_enhancer_input.name] = crop_vision_frame + if face_enhancer_input.name == 'weight': + face_enhancer_inputs[face_enhancer_input.name] = face_enhancer_weight + + with thread_semaphore(): + crop_vision_frame = face_enhancer.run(None, face_enhancer_inputs)[0][0] + + return crop_vision_frame + + +def has_weight_input() -> bool: + face_enhancer = get_inference_pool().get('face_enhancer') + + for deep_swapper_input in face_enhancer.get_inputs(): + if deep_swapper_input.name == 'weight': + return True + + return False + + +def prepare_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + crop_vision_frame = crop_vision_frame[:, :, ::-1] / 255.0 + crop_vision_frame = (crop_vision_frame - 0.5) / 0.5 + crop_vision_frame = numpy.expand_dims(crop_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) + return crop_vision_frame + + +def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + crop_vision_frame = numpy.clip(crop_vision_frame, -1, 1) + crop_vision_frame = (crop_vision_frame + 1) / 2 + crop_vision_frame = crop_vision_frame.transpose(1, 2, 0) + crop_vision_frame = (crop_vision_frame * 255.0).round() + crop_vision_frame = crop_vision_frame.astype(numpy.uint8)[:, :, ::-1] + return crop_vision_frame + + +def blend_paste_frame(temp_vision_frame : VisionFrame, paste_vision_frame : VisionFrame) -> VisionFrame: + face_enhancer_blend = 1 - (state_manager.get_item('face_enhancer_blend') / 100) + temp_vision_frame = blend_frame(temp_vision_frame, paste_vision_frame, 1 - face_enhancer_blend) + return temp_vision_frame + + +def process_frame(inputs : FaceEnhancerInputs) -> ProcessorOutputs: + reference_vision_frame = inputs.get('reference_vision_frame') + target_vision_frame = inputs.get('target_vision_frame') + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + target_faces = select_faces(reference_vision_frame, target_vision_frame) + + if target_faces: + for target_face in target_faces: + target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) + temp_vision_frame = enhance_face(target_face, temp_vision_frame) + + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/face_enhancer/locales.py b/facefusion/processors/modules/face_enhancer/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..b76da6c03d611700466a2e6bc4b9b90748ae75ca --- /dev/null +++ b/facefusion/processors/modules/face_enhancer/locales.py @@ -0,0 +1,20 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for enhancing the face', + 'blend': 'blend the enhanced into the previous face', + 'weight': 'specify the degree of weight applied to the face' + }, + 'uis': + { + 'blend_slider': 'FACE ENHANCER BLEND', + 'model_dropdown': 'FACE ENHANCER MODEL', + 'weight_slider': 'FACE ENHANCER WEIGHT' + } + } +} diff --git a/facefusion/processors/modules/face_enhancer/types.py b/facefusion/processors/modules/face_enhancer/types.py new file mode 100644 index 0000000000000000000000000000000000000000..2dbb8335ab0b6f906f7cfe4686d425562c54e31c --- /dev/null +++ b/facefusion/processors/modules/face_enhancer/types.py @@ -0,0 +1,17 @@ +from typing import Any, Literal, TypeAlias, TypedDict + +from numpy.typing import NDArray + +from facefusion.types import Mask, VisionFrame + +FaceEnhancerInputs = TypedDict('FaceEnhancerInputs', +{ + 'reference_vision_frame' : VisionFrame, + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +FaceEnhancerModel = Literal['codeformer', 'gfpgan_1.2', 'gfpgan_1.3', 'gfpgan_1.4', 'gpen_bfr_256', 'gpen_bfr_512', 'gpen_bfr_1024', 'gpen_bfr_2048', 'restoreformer_plus_plus'] + +FaceEnhancerWeight : TypeAlias = NDArray[Any] diff --git a/facefusion/processors/modules/face_swapper/__init__.py b/facefusion/processors/modules/face_swapper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/face_swapper/choices.py b/facefusion/processors/modules/face_swapper/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..9155b71e0a6481c0f5b9227516de7e0bb377b0b1 --- /dev/null +++ b/facefusion/processors/modules/face_swapper/choices.py @@ -0,0 +1,25 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_float_range +from facefusion.processors.modules.face_swapper.types import FaceSwapperModel, FaceSwapperSet, FaceSwapperWeight + +face_swapper_set : FaceSwapperSet =\ +{ + 'blendswap_256': [ '256x256', '384x384', '512x512', '768x768', '1024x1024' ], + 'ghost_1_256': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'ghost_2_256': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'ghost_3_256': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'hififace_unofficial_256': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'hyperswap_1a_256': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'hyperswap_1b_256': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'hyperswap_1c_256': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'inswapper_128': [ '128x128', '256x256', '384x384', '512x512', '768x768', '1024x1024' ], + 'inswapper_128_fp16': [ '128x128', '256x256', '384x384', '512x512', '768x768', '1024x1024' ], + 'simswap_256': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'simswap_unofficial_512': [ '512x512', '768x768', '1024x1024' ], + 'uniface_256': [ '256x256', '512x512', '768x768', '1024x1024' ] +} + +face_swapper_models : List[FaceSwapperModel] = list(face_swapper_set.keys()) + +face_swapper_weight_range : Sequence[FaceSwapperWeight] = create_float_range(0.0, 1.0, 0.05) diff --git a/facefusion/processors/modules/face_swapper/core.py b/facefusion/processors/modules/face_swapper/core.py new file mode 100644 index 0000000000000000000000000000000000000000..6778cb8f245eef32217fad4501dd2ebf4780293e --- /dev/null +++ b/facefusion/processors/modules/face_swapper/core.py @@ -0,0 +1,774 @@ +from argparse import ArgumentParser +from functools import lru_cache +from typing import List, Optional, Tuple + +import cv2 +import numpy + +import facefusion.choices +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import get_first, is_macos +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.execution import has_execution_provider +from facefusion.face_analyser import get_average_face, get_many_faces, get_one_face, scale_face +from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_area_mask, create_box_mask, create_occlusion_mask, create_region_mask +from facefusion.face_selector import select_faces, sort_faces_by_order +from facefusion.filesystem import filter_image_paths, has_image, in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.model_helper import get_static_model_initializer +from facefusion.processors.modules.face_swapper import choices as face_swapper_choices +from facefusion.processors.modules.face_swapper.types import FaceSwapperInputs +from facefusion.processors.pixel_boost import explode_pixel_boost, implode_pixel_boost +from facefusion.processors.types import ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.types import ApplyStateItem, Args, DownloadScope, Embedding, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import read_static_image, read_static_images, read_static_video_frame, unpack_resolution + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'blendswap_256': + { + '__metadata__': + { + 'vendor': 'mapooon', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'blendswap_256.hash'), + 'path': resolve_relative_path('../.assets/models/blendswap_256.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'blendswap_256.onnx'), + 'path': resolve_relative_path('../.assets/models/blendswap_256.onnx') + } + }, + 'type': 'blendswap', + 'template': 'ffhq_512', + 'size': (256, 256), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'ghost_1_256': + { + '__metadata__': + { + 'vendor': 'ai-forever', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'ghost_1_256.hash'), + 'path': resolve_relative_path('../.assets/models/ghost_1_256.hash') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_ghost.hash'), + 'path': resolve_relative_path('../.assets/models/crossface_ghost.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'ghost_1_256.onnx'), + 'path': resolve_relative_path('../.assets/models/ghost_1_256.onnx') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_ghost.onnx'), + 'path': resolve_relative_path('../.assets/models/crossface_ghost.onnx') + } + }, + 'type': 'ghost', + 'template': 'arcface_112_v1', + 'size': (256, 256), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + }, + 'ghost_2_256': + { + '__metadata__': + { + 'vendor': 'ai-forever', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'ghost_2_256.hash'), + 'path': resolve_relative_path('../.assets/models/ghost_2_256.hash') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_ghost.hash'), + 'path': resolve_relative_path('../.assets/models/crossface_ghost.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'ghost_2_256.onnx'), + 'path': resolve_relative_path('../.assets/models/ghost_2_256.onnx') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_ghost.onnx'), + 'path': resolve_relative_path('../.assets/models/crossface_ghost.onnx') + } + }, + 'type': 'ghost', + 'template': 'arcface_112_v1', + 'size': (256, 256), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + }, + 'ghost_3_256': + { + '__metadata__': + { + 'vendor': 'ai-forever', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'ghost_3_256.hash'), + 'path': resolve_relative_path('../.assets/models/ghost_3_256.hash') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_ghost.hash'), + 'path': resolve_relative_path('../.assets/models/crossface_ghost.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'ghost_3_256.onnx'), + 'path': resolve_relative_path('../.assets/models/ghost_3_256.onnx') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_ghost.onnx'), + 'path': resolve_relative_path('../.assets/models/crossface_ghost.onnx') + } + }, + 'type': 'ghost', + 'template': 'arcface_112_v1', + 'size': (256, 256), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + }, + 'hififace_unofficial_256': + { + '__metadata__': + { + 'vendor': 'GuijiAI', + 'license': 'Unknown', + 'year': 2021 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.1.0', 'hififace_unofficial_256.hash'), + 'path': resolve_relative_path('../.assets/models/hififace_unofficial_256.hash') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_hififace.hash'), + 'path': resolve_relative_path('../.assets/models/crossface_hififace.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.1.0', 'hififace_unofficial_256.onnx'), + 'path': resolve_relative_path('../.assets/models/hififace_unofficial_256.onnx') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_hififace.onnx'), + 'path': resolve_relative_path('../.assets/models/crossface_hififace.onnx') + } + }, + 'type': 'hififace', + 'template': 'mtcnn_512', + 'size': (256, 256), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + }, + 'hyperswap_1a_256': + { + '__metadata__': + { + 'vendor': 'FaceFusion', + 'license': 'ResearchRAIL', + 'year': 2025 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.3.0', 'hyperswap_1a_256.hash'), + 'path': resolve_relative_path('../.assets/models/hyperswap_1a_256.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.3.0', 'hyperswap_1a_256.onnx'), + 'path': resolve_relative_path('../.assets/models/hyperswap_1a_256.onnx') + } + }, + 'type': 'hyperswap', + 'template': 'arcface_128', + 'size': (256, 256), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + }, + 'hyperswap_1b_256': + { + '__metadata__': + { + 'vendor': 'FaceFusion', + 'license': 'ResearchRAIL', + 'year': 2025 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.3.0', 'hyperswap_1b_256.hash'), + 'path': resolve_relative_path('../.assets/models/hyperswap_1b_256.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.3.0', 'hyperswap_1b_256.onnx'), + 'path': resolve_relative_path('../.assets/models/hyperswap_1b_256.onnx') + } + }, + 'type': 'hyperswap', + 'template': 'arcface_128', + 'size': (256, 256), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + }, + 'hyperswap_1c_256': + { + '__metadata__': + { + 'vendor': 'FaceFusion', + 'license': 'ResearchRAIL', + 'year': 2025 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.3.0', 'hyperswap_1c_256.hash'), + 'path': resolve_relative_path('../.assets/models/hyperswap_1c_256.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.3.0', 'hyperswap_1c_256.onnx'), + 'path': resolve_relative_path('../.assets/models/hyperswap_1c_256.onnx') + } + }, + 'type': 'hyperswap', + 'template': 'arcface_128', + 'size': (256, 256), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + }, + 'inswapper_128': + { + '__metadata__': + { + 'vendor': 'InsightFace', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'inswapper_128.hash'), + 'path': resolve_relative_path('../.assets/models/inswapper_128.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'inswapper_128.onnx'), + 'path': resolve_relative_path('../.assets/models/inswapper_128.onnx') + } + }, + 'type': 'inswapper', + 'template': 'arcface_128', + 'size': (128, 128), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'inswapper_128_fp16': + { + '__metadata__': + { + 'vendor': 'InsightFace', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'inswapper_128_fp16.hash'), + 'path': resolve_relative_path('../.assets/models/inswapper_128_fp16.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'inswapper_128_fp16.onnx'), + 'path': resolve_relative_path('../.assets/models/inswapper_128_fp16.onnx') + } + }, + 'type': 'inswapper', + 'template': 'arcface_128', + 'size': (128, 128), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'simswap_256': + { + '__metadata__': + { + 'vendor': 'neuralchen', + 'license': 'Non-Commercial', + 'year': 2020 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'simswap_256.hash'), + 'path': resolve_relative_path('../.assets/models/simswap_256.hash') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_simswap.hash'), + 'path': resolve_relative_path('../.assets/models/crossface_simswap.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'simswap_256.onnx'), + 'path': resolve_relative_path('../.assets/models/simswap_256.onnx') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_simswap.onnx'), + 'path': resolve_relative_path('../.assets/models/crossface_simswap.onnx') + } + }, + 'type': 'simswap', + 'template': 'arcface_112_v1', + 'size': (256, 256), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + }, + 'simswap_unofficial_512': + { + '__metadata__': + { + 'vendor': 'neuralchen', + 'license': 'Non-Commercial', + 'year': 2020 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'simswap_unofficial_512.hash'), + 'path': resolve_relative_path('../.assets/models/simswap_unofficial_512.hash') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_simswap.hash'), + 'path': resolve_relative_path('../.assets/models/crossface_simswap.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'simswap_unofficial_512.onnx'), + 'path': resolve_relative_path('../.assets/models/simswap_unofficial_512.onnx') + }, + 'embedding_converter': + { + 'url': resolve_download_url('models-3.4.0', 'crossface_simswap.onnx'), + 'path': resolve_relative_path('../.assets/models/crossface_simswap.onnx') + } + }, + 'type': 'simswap', + 'template': 'arcface_112_v1', + 'size': (512, 512), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'uniface_256': + { + '__metadata__': + { + 'vendor': 'xc-csc101', + 'license': 'Unknown', + 'year': 2022 + }, + 'hashes': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'uniface_256.hash'), + 'path': resolve_relative_path('../.assets/models/uniface_256.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': resolve_download_url('models-3.0.0', 'uniface_256.onnx'), + 'path': resolve_relative_path('../.assets/models/uniface_256.onnx') + } + }, + 'type': 'uniface', + 'template': 'ffhq_512', + 'size': (256, 256), + 'mean': [ 0.5, 0.5, 0.5 ], + 'standard_deviation': [ 0.5, 0.5, 0.5 ] + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ get_model_name() ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ get_model_name() ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + model_name = get_model_name() + return create_static_model_set('full').get(model_name) + + +def get_model_name() -> str: + model_name = state_manager.get_item('face_swapper_model') + + if is_macos() and has_execution_provider('coreml') and model_name == 'inswapper_128_fp16': + return 'inswapper_128' + return model_name + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--face-swapper-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'face_swapper_model', 'hyperswap_1a_256'), choices = face_swapper_choices.face_swapper_models) + known_args, _ = program.parse_known_args() + face_swapper_pixel_boost_choices = face_swapper_choices.face_swapper_set.get(known_args.face_swapper_model) + group_processors.add_argument('--face-swapper-pixel-boost', help = translator.get('help.pixel_boost', __package__), default = config.get_str_value('processors', 'face_swapper_pixel_boost', get_first(face_swapper_pixel_boost_choices)), choices = face_swapper_pixel_boost_choices) + group_processors.add_argument('--face-swapper-weight', help = translator.get('help.weight', __package__), type = float, default = config.get_float_value('processors', 'face_swapper_weight', '0.5'), choices = face_swapper_choices.face_swapper_weight_range) + facefusion.jobs.job_store.register_step_keys([ 'face_swapper_model', 'face_swapper_pixel_boost', 'face_swapper_weight' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('face_swapper_model', args.get('face_swapper_model')) + apply_state_item('face_swapper_pixel_boost', args.get('face_swapper_pixel_boost')) + apply_state_item('face_swapper_weight', args.get('face_swapper_weight')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if not has_image(state_manager.get_item('source_paths')): + logger.error(translator.get('choose_image_source') + translator.get('exclamation_mark'), __name__) + return False + + source_image_paths = filter_image_paths(state_manager.get_item('source_paths')) + source_frames = read_static_images(source_image_paths) + source_faces = get_many_faces(source_frames) + + if not get_one_face(source_faces): + logger.error(translator.get('no_source_face_detected') + translator.get('exclamation_mark'), __name__) + return False + + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + get_static_model_initializer.cache_clear() + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_masker.clear_inference_pool() + face_recognizer.clear_inference_pool() + + +def swap_face(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + model_template = get_model_options().get('template') + model_size = get_model_options().get('size') + pixel_boost_size = unpack_resolution(state_manager.get_item('face_swapper_pixel_boost')) + pixel_boost_total = pixel_boost_size[0] // model_size[0] + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmark_set.get('5/68'), model_template, pixel_boost_size) + temp_vision_frames = [] + crop_masks = [] + + if 'box' in state_manager.get_item('face_mask_types'): + box_mask = create_box_mask(crop_vision_frame, state_manager.get_item('face_mask_blur'), state_manager.get_item('face_mask_padding')) + crop_masks.append(box_mask) + + if 'occlusion' in state_manager.get_item('face_mask_types'): + occlusion_mask = create_occlusion_mask(crop_vision_frame) + crop_masks.append(occlusion_mask) + + pixel_boost_vision_frames = implode_pixel_boost(crop_vision_frame, pixel_boost_total, model_size) + for pixel_boost_vision_frame in pixel_boost_vision_frames: + pixel_boost_vision_frame = prepare_crop_frame(pixel_boost_vision_frame) + pixel_boost_vision_frame = forward_swap_face(source_face, target_face, pixel_boost_vision_frame) + pixel_boost_vision_frame = normalize_crop_frame(pixel_boost_vision_frame) + temp_vision_frames.append(pixel_boost_vision_frame) + crop_vision_frame = explode_pixel_boost(temp_vision_frames, pixel_boost_total, model_size, pixel_boost_size) + + if 'area' in state_manager.get_item('face_mask_types'): + face_landmark_68 = cv2.transform(target_face.landmark_set.get('68').reshape(1, -1, 2), affine_matrix).reshape(-1, 2) + area_mask = create_area_mask(crop_vision_frame, face_landmark_68, state_manager.get_item('face_mask_areas')) + crop_masks.append(area_mask) + + if 'region' in state_manager.get_item('face_mask_types'): + region_mask = create_region_mask(crop_vision_frame, state_manager.get_item('face_mask_regions')) + crop_masks.append(region_mask) + + crop_mask = numpy.minimum.reduce(crop_masks).clip(0, 1) + paste_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, crop_mask, affine_matrix) + return paste_vision_frame + + +def forward_swap_face(source_face : Face, target_face : Face, crop_vision_frame : VisionFrame) -> VisionFrame: + face_swapper = get_inference_pool().get('face_swapper') + model_type = get_model_options().get('type') + face_swapper_inputs = {} + + if is_macos() and has_execution_provider('coreml') and model_type in [ 'ghost', 'uniface' ]: + face_swapper.set_providers([ facefusion.choices.execution_provider_set.get('cpu') ]) + + for face_swapper_input in face_swapper.get_inputs(): + if face_swapper_input.name == 'source': + if model_type in [ 'blendswap', 'uniface' ]: + face_swapper_inputs[face_swapper_input.name] = prepare_source_frame(source_face) + else: + source_embedding = prepare_source_embedding(source_face) + source_embedding = balance_source_embedding(source_embedding, target_face.embedding) + face_swapper_inputs[face_swapper_input.name] = source_embedding + if face_swapper_input.name == 'target': + face_swapper_inputs[face_swapper_input.name] = crop_vision_frame + + with conditional_thread_semaphore(): + crop_vision_frame = face_swapper.run(None, face_swapper_inputs)[0][0] + + return crop_vision_frame + + +def forward_convert_embedding(face_embedding : Embedding) -> Embedding: + embedding_converter = get_inference_pool().get('embedding_converter') + + with conditional_thread_semaphore(): + face_embedding = embedding_converter.run(None, + { + 'input': face_embedding + })[0] + + return face_embedding + + +def prepare_source_frame(source_face : Face) -> VisionFrame: + model_type = get_model_options().get('type') + source_vision_frame = read_static_image(get_first(state_manager.get_item('source_paths'))) + + if model_type == 'blendswap': + source_vision_frame, _ = warp_face_by_face_landmark_5(source_vision_frame, source_face.landmark_set.get('5/68'), 'arcface_112_v2', (112, 112)) + + if model_type == 'uniface': + source_vision_frame, _ = warp_face_by_face_landmark_5(source_vision_frame, source_face.landmark_set.get('5/68'), 'ffhq_512', (256, 256)) + + source_vision_frame = source_vision_frame[:, :, ::-1] / 255.0 + source_vision_frame = source_vision_frame.transpose(2, 0, 1) + source_vision_frame = numpy.expand_dims(source_vision_frame, axis = 0).astype(numpy.float32) + return source_vision_frame + + +def prepare_source_embedding(source_face : Face) -> Embedding: + model_type = get_model_options().get('type') + + if model_type == 'ghost': + source_embedding = source_face.embedding.reshape(-1, 512) + source_embedding, _ = convert_source_embedding(source_embedding) + source_embedding = source_embedding.reshape(1, -1) + return source_embedding + + if model_type == 'hyperswap': + source_embedding = source_face.embedding_norm.reshape((1, -1)) + return source_embedding + + if model_type == 'inswapper': + model_path = get_model_options().get('sources').get('face_swapper').get('path') + model_initializer = get_static_model_initializer(model_path) + source_embedding = source_face.embedding.reshape((1, -1)) + source_embedding = numpy.dot(source_embedding, model_initializer) / numpy.linalg.norm(source_embedding) + return source_embedding + + source_embedding = source_face.embedding.reshape(-1, 512) + _, source_embedding_norm = convert_source_embedding(source_embedding) + source_embedding = source_embedding_norm.reshape(1, -1) + return source_embedding + + +def balance_source_embedding(source_embedding : Embedding, target_embedding : Embedding) -> Embedding: + model_type = get_model_options().get('type') + face_swapper_weight = state_manager.get_item('face_swapper_weight') + face_swapper_weight = numpy.interp(face_swapper_weight, [ 0, 1 ], [ 0.35, -0.35 ]).astype(numpy.float32) + + if model_type in [ 'hififace', 'hyperswap', 'inswapper', 'simswap' ]: + target_embedding = target_embedding / numpy.linalg.norm(target_embedding) + + source_embedding = source_embedding.reshape(1, -1) + target_embedding = target_embedding.reshape(1, -1) + source_embedding = source_embedding * (1 - face_swapper_weight) + target_embedding * face_swapper_weight + return source_embedding + + +def convert_source_embedding(source_embedding : Embedding) -> Tuple[Embedding, Embedding]: + source_embedding = forward_convert_embedding(source_embedding) + source_embedding = source_embedding.ravel() + source_embedding_norm = source_embedding / numpy.linalg.norm(source_embedding) + return source_embedding, source_embedding_norm + + +def prepare_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + model_mean = get_model_options().get('mean') + model_standard_deviation = get_model_options().get('standard_deviation') + + crop_vision_frame = crop_vision_frame[:, :, ::-1] / 255.0 + crop_vision_frame = (crop_vision_frame - model_mean) / model_standard_deviation + crop_vision_frame = crop_vision_frame.transpose(2, 0, 1) + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0).astype(numpy.float32) + return crop_vision_frame + + +def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + model_type = get_model_options().get('type') + model_mean = get_model_options().get('mean') + model_standard_deviation = get_model_options().get('standard_deviation') + + crop_vision_frame = crop_vision_frame.transpose(1, 2, 0) + + if model_type in [ 'ghost', 'hififace', 'hyperswap', 'uniface' ]: + crop_vision_frame = crop_vision_frame * model_standard_deviation + model_mean + + crop_vision_frame = crop_vision_frame.clip(0, 1) + crop_vision_frame = crop_vision_frame[:, :, ::-1] * 255 + return crop_vision_frame + + +def extract_source_face(source_vision_frames : List[VisionFrame]) -> Optional[Face]: + source_faces = [] + + if source_vision_frames: + for source_vision_frame in source_vision_frames: + temp_faces = get_many_faces([source_vision_frame]) + temp_faces = sort_faces_by_order(temp_faces, 'large-small') + + if temp_faces: + source_faces.append(get_first(temp_faces)) + + return get_average_face(source_faces) + + +def process_frame(inputs : FaceSwapperInputs) -> ProcessorOutputs: + reference_vision_frame = inputs.get('reference_vision_frame') + source_vision_frames = inputs.get('source_vision_frames') + target_vision_frame = inputs.get('target_vision_frame') + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + source_face = extract_source_face(source_vision_frames) + target_faces = select_faces(reference_vision_frame, target_vision_frame) + + if source_face and target_faces: + for target_face in target_faces: + target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) + temp_vision_frame = swap_face(source_face, target_face, temp_vision_frame) + + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/face_swapper/locales.py b/facefusion/processors/modules/face_swapper/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..d8af5ee7f385af361df3aa12c01a6947f13185f2 --- /dev/null +++ b/facefusion/processors/modules/face_swapper/locales.py @@ -0,0 +1,20 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for swapping the face', + 'pixel_boost': 'choose the pixel boost resolution for the face swapper', + 'weight': 'specify the degree of weight applied to the face' + }, + 'uis': + { + 'model_dropdown': 'FACE SWAPPER MODEL', + 'pixel_boost_dropdown': 'FACE SWAPPER PIXEL BOOST', + 'weight_slider': 'FACE SWAPPER WEIGHT' + } + } +} diff --git a/facefusion/processors/modules/face_swapper/types.py b/facefusion/processors/modules/face_swapper/types.py new file mode 100644 index 0000000000000000000000000000000000000000..afe593f5b6f834b3bbf47e5770f4320fec6475c7 --- /dev/null +++ b/facefusion/processors/modules/face_swapper/types.py @@ -0,0 +1,18 @@ +from typing import Dict, List, Literal, TypeAlias, TypedDict + +from facefusion.types import Mask, VisionFrame + +FaceSwapperInputs = TypedDict('FaceSwapperInputs', +{ + 'reference_vision_frame' : VisionFrame, + 'source_vision_frames' : List[VisionFrame], + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +FaceSwapperModel = Literal['blendswap_256', 'ghost_1_256', 'ghost_2_256', 'ghost_3_256', 'hififace_unofficial_256', 'hyperswap_1a_256', 'hyperswap_1b_256', 'hyperswap_1c_256', 'inswapper_128', 'inswapper_128_fp16', 'simswap_256', 'simswap_unofficial_512', 'uniface_256'] + +FaceSwapperWeight : TypeAlias = float + +FaceSwapperSet : TypeAlias = Dict[FaceSwapperModel, List[str]] diff --git a/facefusion/processors/modules/frame_colorizer/__init__.py b/facefusion/processors/modules/frame_colorizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/frame_colorizer/choices.py b/facefusion/processors/modules/frame_colorizer/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..5ce08a786dbf239e39a8aafc57fea7756d765c2d --- /dev/null +++ b/facefusion/processors/modules/frame_colorizer/choices.py @@ -0,0 +1,10 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_int_range +from facefusion.processors.modules.frame_colorizer.types import FrameColorizerModel + +frame_colorizer_models : List[FrameColorizerModel] = [ 'ddcolor', 'ddcolor_artistic', 'deoldify', 'deoldify_artistic', 'deoldify_stable' ] + +frame_colorizer_sizes : List[str] = [ '192x192', '256x256', '384x384', '512x512' ] + +frame_colorizer_blend_range : Sequence[int] = create_int_range(0, 100, 1) diff --git a/facefusion/processors/modules/frame_colorizer/core.py b/facefusion/processors/modules/frame_colorizer/core.py new file mode 100644 index 0000000000000000000000000000000000000000..0218e39aeba66b2f81a4689cb0f3b980de345c67 --- /dev/null +++ b/facefusion/processors/modules/frame_colorizer/core.py @@ -0,0 +1,299 @@ +from argparse import ArgumentParser +from functools import lru_cache +from typing import List + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import create_int_metavar, is_macos +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.execution import has_execution_provider +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors.modules.frame_colorizer import choices as frame_colorizer_choices +from facefusion.processors.modules.frame_colorizer.types import FrameColorizerInputs +from facefusion.processors.types import ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import thread_semaphore +from facefusion.types import ApplyStateItem, Args, DownloadScope, ExecutionProvider, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import blend_frame, read_static_image, read_static_video_frame, unpack_resolution + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'ddcolor': + { + '__metadata__': + { + 'vendor': 'piddnad', + 'license': 'Apache-2.0', + 'year': 2023 + }, + 'hashes': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'ddcolor.hash'), + 'path': resolve_relative_path('../.assets/models/ddcolor.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'ddcolor.onnx'), + 'path': resolve_relative_path('../.assets/models/ddcolor.onnx') + } + }, + 'type': 'ddcolor' + }, + 'ddcolor_artistic': + { + '__metadata__': + { + 'vendor': 'piddnad', + 'license': 'Apache-2.0', + 'year': 2023 + }, + 'hashes': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'ddcolor_artistic.hash'), + 'path': resolve_relative_path('../.assets/models/ddcolor_artistic.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'ddcolor_artistic.onnx'), + 'path': resolve_relative_path('../.assets/models/ddcolor_artistic.onnx') + } + }, + 'type': 'ddcolor' + }, + 'deoldify': + { + '__metadata__': + { + 'vendor': 'jantic', + 'license': 'MIT', + 'year': 2022 + }, + 'hashes': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'deoldify.hash'), + 'path': resolve_relative_path('../.assets/models/deoldify.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'deoldify.onnx'), + 'path': resolve_relative_path('../.assets/models/deoldify.onnx') + } + }, + 'type': 'deoldify' + }, + 'deoldify_artistic': + { + '__metadata__': + { + 'vendor': 'jantic', + 'license': 'MIT', + 'year': 2022 + }, + 'hashes': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'deoldify_artistic.hash'), + 'path': resolve_relative_path('../.assets/models/deoldify_artistic.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'deoldify_artistic.onnx'), + 'path': resolve_relative_path('../.assets/models/deoldify_artistic.onnx') + } + }, + 'type': 'deoldify' + }, + 'deoldify_stable': + { + '__metadata__': + { + 'vendor': 'jantic', + 'license': 'MIT', + 'year': 2022 + }, + 'hashes': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'deoldify_stable.hash'), + 'path': resolve_relative_path('../.assets/models/deoldify_stable.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': resolve_download_url('models-3.0.0', 'deoldify_stable.onnx'), + 'path': resolve_relative_path('../.assets/models/deoldify_stable.onnx') + } + }, + 'type': 'deoldify' + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('frame_colorizer_model') ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('frame_colorizer_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def resolve_execution_providers() -> List[ExecutionProvider]: + if is_macos() and has_execution_provider('coreml'): + return [ 'cpu' ] + return state_manager.get_item('execution_providers') + + +def get_model_options() -> ModelOptions: + model_name = state_manager.get_item('frame_colorizer_model') + return create_static_model_set('full').get(model_name) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--frame-colorizer-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'frame_colorizer_model', 'ddcolor'), choices = frame_colorizer_choices.frame_colorizer_models) + group_processors.add_argument('--frame-colorizer-size', help = translator.get('help.size', __package__), type = str, default = config.get_str_value('processors', 'frame_colorizer_size', '256x256'), choices = frame_colorizer_choices.frame_colorizer_sizes) + group_processors.add_argument('--frame-colorizer-blend', help = translator.get('help.blend', __package__), type = int, default = config.get_int_value('processors', 'frame_colorizer_blend', '100'), choices = frame_colorizer_choices.frame_colorizer_blend_range, metavar = create_int_metavar(frame_colorizer_choices.frame_colorizer_blend_range)) + facefusion.jobs.job_store.register_step_keys([ 'frame_colorizer_model', 'frame_colorizer_blend', 'frame_colorizer_size' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('frame_colorizer_model', args.get('frame_colorizer_model')) + apply_state_item('frame_colorizer_blend', args.get('frame_colorizer_blend')) + apply_state_item('frame_colorizer_size', args.get('frame_colorizer_size')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + + +def colorize_frame(temp_vision_frame : VisionFrame) -> VisionFrame: + color_vision_frame = prepare_temp_frame(temp_vision_frame) + color_vision_frame = forward(color_vision_frame) + color_vision_frame = merge_color_frame(temp_vision_frame, color_vision_frame) + color_vision_frame = blend_color_frame(temp_vision_frame, color_vision_frame) + return color_vision_frame + + +def forward(color_vision_frame : VisionFrame) -> VisionFrame: + frame_colorizer = get_inference_pool().get('frame_colorizer') + + with thread_semaphore(): + color_vision_frame = frame_colorizer.run(None, + { + 'input': color_vision_frame + })[0][0] + + return color_vision_frame + + +def prepare_temp_frame(temp_vision_frame : VisionFrame) -> VisionFrame: + model_size = unpack_resolution(state_manager.get_item('frame_colorizer_size')) + model_type = get_model_options().get('type') + temp_vision_frame = cv2.cvtColor(temp_vision_frame, cv2.COLOR_BGR2GRAY) + temp_vision_frame = cv2.cvtColor(temp_vision_frame, cv2.COLOR_GRAY2RGB) + + if model_type == 'ddcolor': + temp_vision_frame = (temp_vision_frame / 255.0).astype(numpy.float32) #type:ignore[operator] + temp_vision_frame = cv2.cvtColor(temp_vision_frame, cv2.COLOR_RGB2LAB)[:, :, :1] + temp_vision_frame = numpy.concatenate((temp_vision_frame, numpy.zeros_like(temp_vision_frame), numpy.zeros_like(temp_vision_frame)), axis = -1) + temp_vision_frame = cv2.cvtColor(temp_vision_frame, cv2.COLOR_LAB2RGB) + + temp_vision_frame = cv2.resize(temp_vision_frame, model_size) + temp_vision_frame = temp_vision_frame.transpose((2, 0, 1)) + temp_vision_frame = numpy.expand_dims(temp_vision_frame, axis = 0).astype(numpy.float32) + return temp_vision_frame + + +def merge_color_frame(temp_vision_frame : VisionFrame, color_vision_frame : VisionFrame) -> VisionFrame: + model_type = get_model_options().get('type') + color_vision_frame = color_vision_frame.transpose(1, 2, 0) + color_vision_frame = cv2.resize(color_vision_frame, (temp_vision_frame.shape[1], temp_vision_frame.shape[0])) + + if model_type == 'ddcolor': + temp_vision_frame = (temp_vision_frame / 255.0).astype(numpy.float32) + temp_vision_frame = cv2.cvtColor(temp_vision_frame, cv2.COLOR_BGR2LAB)[:, :, :1] + color_vision_frame = numpy.concatenate((temp_vision_frame, color_vision_frame), axis = -1) + color_vision_frame = cv2.cvtColor(color_vision_frame, cv2.COLOR_LAB2BGR) + color_vision_frame = (color_vision_frame * 255.0).round().astype(numpy.uint8) #type:ignore[operator] + + if model_type == 'deoldify': + temp_blue_channel, _, _ = cv2.split(temp_vision_frame) + color_vision_frame = cv2.cvtColor(color_vision_frame, cv2.COLOR_BGR2RGB).astype(numpy.uint8) + color_vision_frame = cv2.cvtColor(color_vision_frame, cv2.COLOR_BGR2LAB) + _, color_green_channel, color_red_channel = cv2.split(color_vision_frame) + color_vision_frame = cv2.merge((temp_blue_channel, color_green_channel, color_red_channel)) + color_vision_frame = cv2.cvtColor(color_vision_frame, cv2.COLOR_LAB2BGR) + return color_vision_frame + + +def blend_color_frame(temp_vision_frame : VisionFrame, color_vision_frame : VisionFrame) -> VisionFrame: + frame_colorizer_blend = 1 - (state_manager.get_item('frame_colorizer_blend') / 100) + temp_vision_frame = blend_frame(temp_vision_frame, color_vision_frame, 1 - frame_colorizer_blend) + return temp_vision_frame + + +def process_frame(inputs : FrameColorizerInputs) -> ProcessorOutputs: + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + temp_vision_frame = colorize_frame(temp_vision_frame) + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/frame_colorizer/locales.py b/facefusion/processors/modules/frame_colorizer/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..f451825141f245763fd036fa87a4db2657dbe2c9 --- /dev/null +++ b/facefusion/processors/modules/frame_colorizer/locales.py @@ -0,0 +1,20 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for colorizing the frame', + 'size': 'specify the frame size provided to the frame colorizer', + 'blend': 'blend the colorized into the previous frame' + }, + 'uis': + { + 'blend_slider': 'FRAME COLORIZER BLEND', + 'model_dropdown': 'FRAME COLORIZER MODEL', + 'size_dropdown': 'FRAME COLORIZER SIZE' + } + } +} diff --git a/facefusion/processors/modules/frame_colorizer/types.py b/facefusion/processors/modules/frame_colorizer/types.py new file mode 100644 index 0000000000000000000000000000000000000000..05ae0a3e2792dc9fe89839ed5ec0642e262534a6 --- /dev/null +++ b/facefusion/processors/modules/frame_colorizer/types.py @@ -0,0 +1,12 @@ +from typing import Literal, TypedDict + +from facefusion.types import Mask, VisionFrame + +FrameColorizerInputs = TypedDict('FrameColorizerInputs', +{ + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +FrameColorizerModel = Literal['ddcolor', 'ddcolor_artistic', 'deoldify', 'deoldify_artistic', 'deoldify_stable'] diff --git a/facefusion/processors/modules/frame_enhancer/__init__.py b/facefusion/processors/modules/frame_enhancer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/frame_enhancer/choices.py b/facefusion/processors/modules/frame_enhancer/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..d6cea51e9f83634321427d60d7a79df2a3f39b37 --- /dev/null +++ b/facefusion/processors/modules/frame_enhancer/choices.py @@ -0,0 +1,8 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_int_range +from facefusion.processors.modules.frame_enhancer.types import FrameEnhancerModel + +frame_enhancer_models : List[FrameEnhancerModel] = [ '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' ] + +frame_enhancer_blend_range : Sequence[int] = create_int_range(0, 100, 1) diff --git a/facefusion/processors/modules/frame_enhancer/core.py b/facefusion/processors/modules/frame_enhancer/core.py new file mode 100644 index 0000000000000000000000000000000000000000..d2f9dbad2887c7eace117a66010709623bd72bcd --- /dev/null +++ b/facefusion/processors/modules/frame_enhancer/core.py @@ -0,0 +1,669 @@ +from argparse import ArgumentParser +from functools import lru_cache + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, inference_manager, logger, state_manager, translator, video_manager +from facefusion.common_helper import create_int_metavar, is_macos +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.execution import has_execution_provider +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors.modules.frame_enhancer import choices as frame_enhancer_choices +from facefusion.processors.modules.frame_enhancer.types import FrameEnhancerInputs +from facefusion.processors.types import ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.types import ApplyStateItem, Args, DownloadScope, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import blend_frame, create_tile_frames, merge_tile_frames, read_static_image, read_static_video_frame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'clear_reality_x4': + { + '__metadata__': + { + 'vendor': 'Kim2091', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'clear_reality_x4.hash'), + 'path': resolve_relative_path('../.assets/models/clear_reality_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'clear_reality_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/clear_reality_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'face_dat_x4': + { + '__metadata__': + { + 'vendor': 'Helaman', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.5.0', 'face_dat_x4.hash'), + 'path': resolve_relative_path('../.assets/models/face_dat_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.5.0', 'face_dat_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/face_dat_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'nomos8k_sc_x4': + { + '__metadata__': + { + 'vendor': 'Phhofm', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'nomos8k_sc_x4.hash'), + 'path': resolve_relative_path('../.assets/models/nomos8k_sc_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'nomos8k_sc_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/nomos8k_sc_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'real_esrgan_x2': + { + '__metadata__': + { + 'vendor': 'xinntao', + 'license': 'BSD-3-Clause', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x2.hash'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x2.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x2.onnx'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x2.onnx') + } + }, + 'size': (256, 16, 8), + 'scale': 2 + }, + 'real_esrgan_x2_fp16': + { + '__metadata__': + { + 'vendor': 'xinntao', + 'license': 'BSD-3-Clause', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x2_fp16.hash'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x2_fp16.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x2_fp16.onnx'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x2_fp16.onnx') + } + }, + 'size': (256, 16, 8), + 'scale': 2 + }, + 'real_esrgan_x4': + { + '__metadata__': + { + 'vendor': 'xinntao', + 'license': 'BSD-3-Clause', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x4.hash'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x4.onnx') + } + }, + 'size': (256, 16, 8), + 'scale': 4 + }, + 'real_esrgan_x4_fp16': + { + '__metadata__': + { + 'vendor': 'xinntao', + 'license': 'BSD-3-Clause', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x4_fp16.hash'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x4_fp16.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x4_fp16.onnx'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x4_fp16.onnx') + } + }, + 'size': (256, 16, 8), + 'scale': 4 + }, + 'real_esrgan_x8': + { + '__metadata__': + { + 'vendor': 'xinntao', + 'license': 'BSD-3-Clause', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x8.hash'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x8.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x8.onnx'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x8.onnx') + } + }, + 'size': (256, 16, 8), + 'scale': 8 + }, + 'real_esrgan_x8_fp16': + { + '__metadata__': + { + 'vendor': 'xinntao', + 'license': 'BSD-3-Clause', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x8_fp16.hash'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x8_fp16.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_esrgan_x8_fp16.onnx'), + 'path': resolve_relative_path('../.assets/models/real_esrgan_x8_fp16.onnx') + } + }, + 'size': (256, 16, 8), + 'scale': 8 + }, + 'real_hatgan_x4': + { + '__metadata__': + { + 'vendor': 'XPixelGroup', + 'license': 'Apache-2.0', + 'year': 2023 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_hatgan_x4.hash'), + 'path': resolve_relative_path('../.assets/models/real_hatgan_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'real_hatgan_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/real_hatgan_x4.onnx') + } + }, + 'size': (256, 16, 8), + 'scale': 4 + }, + 'real_web_photo_x4': + { + '__metadata__': + { + 'vendor': 'Helaman', + 'license': 'Non-Commercial', + 'year': 2024 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'real_web_photo_x4.hash'), + 'path': resolve_relative_path('../.assets/models/real_web_photo_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'real_web_photo_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/real_web_photo_x4.onnx') + } + }, + 'size': (64, 4, 2), + 'scale': 4 + }, + 'realistic_rescaler_x4': + { + '__metadata__': + { + 'vendor': 'Mutin Choler', + 'license': 'WTFPL', + 'year': 2023 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'realistic_rescaler_x4.hash'), + 'path': resolve_relative_path('../.assets/models/realistic_rescaler_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'realistic_rescaler_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/realistic_rescaler_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'remacri_x4': + { + '__metadata__': + { + 'vendor': 'FoolhardyVEVO', + 'license': 'Non-Commercial', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'remacri_x4.hash'), + 'path': resolve_relative_path('../.assets/models/remacri_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'remacri_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/remacri_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'siax_x4': + { + '__metadata__': + { + 'vendor': 'NMKD', + 'license': 'WTFPL', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'siax_x4.hash'), + 'path': resolve_relative_path('../.assets/models/siax_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'siax_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/siax_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'span_kendata_x4': + { + '__metadata__': + { + 'vendor': 'terrainer', + 'license': 'Non-Commercial', + 'year': 2024 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'span_kendata_x4.hash'), + 'path': resolve_relative_path('../.assets/models/span_kendata_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'span_kendata_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/span_kendata_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'swin2_sr_x4': + { + '__metadata__': + { + 'vendor': 'mv-lab', + 'license': 'Apache-2.0', + 'year': 2022 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'swin2_sr_x4.hash'), + 'path': resolve_relative_path('../.assets/models/swin2_sr_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.1.0', 'swin2_sr_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/swin2_sr_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'tghq_face_x8': + { + '__metadata__': + { + 'vendor': 'TorrentGuy', + 'license': 'GPL-3.0', + 'year': 2019 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.5.0', 'tghq_face_x8.hash'), + 'path': resolve_relative_path('../.assets/models/tghq_face_x8.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.5.0', 'tghq_face_x8.onnx'), + 'path': resolve_relative_path('../.assets/models/tghq_face_x8.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 8 + }, + 'ultra_sharp_x4': + { + '__metadata__': + { + 'vendor': 'Kim2091', + 'license': 'Non-Commercial', + 'year': 2021 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'ultra_sharp_x4.hash'), + 'path': resolve_relative_path('../.assets/models/ultra_sharp_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.0.0', 'ultra_sharp_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/ultra_sharp_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'ultra_sharp_2_x4': + { + '__metadata__': + { + 'vendor': 'Kim2091', + 'license': 'Non-Commercial', + 'year': 2025 + }, + 'hashes': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.3.0', 'ultra_sharp_2_x4.hash'), + 'path': resolve_relative_path('../.assets/models/ultra_sharp_2_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': resolve_download_url('models-3.3.0', 'ultra_sharp_2_x4.onnx'), + 'path': resolve_relative_path('../.assets/models/ultra_sharp_2_x4.onnx') + } + }, + 'size': (1024, 64, 32), + 'scale': 4 + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ get_frame_enhancer_model() ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ get_frame_enhancer_model() ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + model_name = get_frame_enhancer_model() + return create_static_model_set('full').get(model_name) + + +def get_frame_enhancer_model() -> str: + frame_enhancer_model = state_manager.get_item('frame_enhancer_model') + + if is_macos() and has_execution_provider('coreml'): + if frame_enhancer_model == 'real_esrgan_x2_fp16': + return 'real_esrgan_x2' + if frame_enhancer_model == 'real_esrgan_x4_fp16': + return 'real_esrgan_x4' + if frame_enhancer_model == 'real_esrgan_x8_fp16': + return 'real_esrgan_x8' + return frame_enhancer_model + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--frame-enhancer-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'frame_enhancer_model', 'span_kendata_x4'), choices = frame_enhancer_choices.frame_enhancer_models) + group_processors.add_argument('--frame-enhancer-blend', help = translator.get('help.blend', __package__), type = int, default = config.get_int_value('processors', 'frame_enhancer_blend', '80'), choices = frame_enhancer_choices.frame_enhancer_blend_range, metavar = create_int_metavar(frame_enhancer_choices.frame_enhancer_blend_range)) + facefusion.jobs.job_store.register_step_keys([ 'frame_enhancer_model', 'frame_enhancer_blend' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('frame_enhancer_model', args.get('frame_enhancer_model')) + apply_state_item('frame_enhancer_blend', args.get('frame_enhancer_blend')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): + logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not same_file_extension(state_manager.get_item('target_path'), state_manager.get_item('output_path')): + logger.error(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + + +def enhance_frame(temp_vision_frame : VisionFrame) -> VisionFrame: + model_size = get_model_options().get('size') + model_scale = get_model_options().get('scale') + temp_height, temp_width = temp_vision_frame.shape[:2] + tile_vision_frames, pad_width, pad_height = create_tile_frames(temp_vision_frame, model_size) + + for index, tile_vision_frame in enumerate(tile_vision_frames): + tile_vision_frame = prepare_tile_frame(tile_vision_frame) + tile_vision_frame = forward(tile_vision_frame) + tile_vision_frames[index] = normalize_tile_frame(tile_vision_frame) + + merge_vision_frame = merge_tile_frames(tile_vision_frames, temp_width * model_scale, temp_height * model_scale, pad_width * model_scale, pad_height * model_scale, (model_size[0] * model_scale, model_size[1] * model_scale, model_size[2] * model_scale)) + temp_vision_frame = blend_merge_frame(temp_vision_frame, merge_vision_frame) + return temp_vision_frame + + +def forward(tile_vision_frame : VisionFrame) -> VisionFrame: + frame_enhancer = get_inference_pool().get('frame_enhancer') + + with conditional_thread_semaphore(): + tile_vision_frame = frame_enhancer.run(None, + { + 'input': tile_vision_frame + })[0] + + return tile_vision_frame + + +def prepare_tile_frame(tile_vision_frame : VisionFrame) -> VisionFrame: + tile_vision_frame = numpy.expand_dims(tile_vision_frame[:, :, ::-1], axis = 0) + tile_vision_frame = tile_vision_frame.transpose(0, 3, 1, 2) + tile_vision_frame = tile_vision_frame.astype(numpy.float32) / 255.0 + return tile_vision_frame + + +def normalize_tile_frame(tile_vision_frame : VisionFrame) -> VisionFrame: + tile_vision_frame = tile_vision_frame.transpose(0, 2, 3, 1).squeeze(0) * 255 + tile_vision_frame = tile_vision_frame.clip(0, 255).astype(numpy.uint8)[:, :, ::-1] + return tile_vision_frame + + +def blend_merge_frame(temp_vision_frame : VisionFrame, merge_vision_frame : VisionFrame) -> VisionFrame: + frame_enhancer_blend = 1 - (state_manager.get_item('frame_enhancer_blend') / 100) + temp_vision_frame = cv2.resize(temp_vision_frame, (merge_vision_frame.shape[1], merge_vision_frame.shape[0])) + temp_vision_frame = blend_frame(temp_vision_frame, merge_vision_frame, 1 - frame_enhancer_blend) + return temp_vision_frame + + +def process_frame(inputs : FrameEnhancerInputs) -> ProcessorOutputs: + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + temp_vision_frame = enhance_frame(temp_vision_frame) + temp_vision_mask = cv2.resize(temp_vision_mask, temp_vision_frame.shape[:2][::-1]) + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/frame_enhancer/locales.py b/facefusion/processors/modules/frame_enhancer/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..11e3a6e7555ee83494d2aefe7fad53348103d779 --- /dev/null +++ b/facefusion/processors/modules/frame_enhancer/locales.py @@ -0,0 +1,18 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for enhancing the frame', + 'blend': 'blend the enhanced into the previous frame' + }, + 'uis': + { + 'blend_slider': 'FRAME ENHANCER BLEND', + 'model_dropdown': 'FRAME ENHANCER MODEL' + } + } +} diff --git a/facefusion/processors/modules/frame_enhancer/types.py b/facefusion/processors/modules/frame_enhancer/types.py new file mode 100644 index 0000000000000000000000000000000000000000..1b09cae483d000f83a1ca3db082d5497da05532c --- /dev/null +++ b/facefusion/processors/modules/frame_enhancer/types.py @@ -0,0 +1,12 @@ +from typing import Literal, TypedDict + +from facefusion.types import Mask, VisionFrame + +FrameEnhancerInputs = TypedDict('FrameEnhancerInputs', +{ + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + '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'] diff --git a/facefusion/processors/modules/lip_syncer/__init__.py b/facefusion/processors/modules/lip_syncer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/processors/modules/lip_syncer/choices.py b/facefusion/processors/modules/lip_syncer/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..69bb2fb51d3cf6ad0b02b43ee9a7bc6d362cf10c --- /dev/null +++ b/facefusion/processors/modules/lip_syncer/choices.py @@ -0,0 +1,8 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_float_range +from facefusion.processors.modules.lip_syncer.types import LipSyncerModel + +lip_syncer_models : List[LipSyncerModel] = [ 'edtalk_256', 'wav2lip_96', 'wav2lip_gan_96' ] + +lip_syncer_weight_range : Sequence[float] = create_float_range(0.0, 1.0, 0.05) diff --git a/facefusion/processors/modules/lip_syncer/core.py b/facefusion/processors/modules/lip_syncer/core.py new file mode 100644 index 0000000000000000000000000000000000000000..b559053d47934da2648e5423e1f022326ff1e437 --- /dev/null +++ b/facefusion/processors/modules/lip_syncer/core.py @@ -0,0 +1,296 @@ +from argparse import ArgumentParser +from functools import lru_cache + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager, voice_extractor +from facefusion.audio import read_static_voice +from facefusion.common_helper import create_float_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url +from facefusion.face_analyser import scale_face +from facefusion.face_helper import create_bounding_box, paste_back, warp_face_by_bounding_box, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_area_mask, create_box_mask, create_occlusion_mask +from facefusion.face_selector import select_faces +from facefusion.filesystem import has_audio, resolve_relative_path +from facefusion.processors.modules.lip_syncer import choices as lip_syncer_choices +from facefusion.processors.modules.lip_syncer.types import LipSyncerInputs, LipSyncerWeight +from facefusion.processors.types import ProcessorOutputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.types import ApplyStateItem, Args, AudioFrame, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import read_static_image, read_static_video_frame + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'edtalk_256': + { + '__metadata__': + { + 'vendor': 'tanshuai0219', + 'license': 'Apache-2.0', + 'year': 2024 + }, + 'hashes': + { + 'lip_syncer': + { + 'url': resolve_download_url('models-3.3.0', 'edtalk_256.hash'), + 'path': resolve_relative_path('../.assets/models/edtalk_256.hash') + } + }, + 'sources': + { + 'lip_syncer': + { + 'url': resolve_download_url('models-3.3.0', 'edtalk_256.onnx'), + 'path': resolve_relative_path('../.assets/models/edtalk_256.onnx') + } + }, + 'type': 'edtalk', + 'size': (256, 256) + }, + 'wav2lip_96': + { + '__metadata__': + { + 'vendor': 'Rudrabha', + 'license': 'Non-Commercial', + 'year': 2020 + }, + 'hashes': + { + 'lip_syncer': + { + 'url': resolve_download_url('models-3.0.0', 'wav2lip_96.hash'), + 'path': resolve_relative_path('../.assets/models/wav2lip_96.hash') + } + }, + 'sources': + { + 'lip_syncer': + { + 'url': resolve_download_url('models-3.0.0', 'wav2lip_96.onnx'), + 'path': resolve_relative_path('../.assets/models/wav2lip_96.onnx') + } + }, + 'type': 'wav2lip', + 'size': (96, 96) + }, + 'wav2lip_gan_96': + { + '__metadata__': + { + 'vendor': 'Rudrabha', + 'license': 'Non-Commercial', + 'year': 2020 + }, + 'hashes': + { + 'lip_syncer': + { + 'url': resolve_download_url('models-3.0.0', 'wav2lip_gan_96.hash'), + 'path': resolve_relative_path('../.assets/models/wav2lip_gan_96.hash') + } + }, + 'sources': + { + 'lip_syncer': + { + 'url': resolve_download_url('models-3.0.0', 'wav2lip_gan_96.onnx'), + 'path': resolve_relative_path('../.assets/models/wav2lip_gan_96.onnx') + } + }, + 'type': 'wav2lip', + 'size': (96, 96) + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('lip_syncer_model') ] + model_source_set = get_model_options().get('sources') + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('lip_syncer_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def get_model_options() -> ModelOptions: + model_name = state_manager.get_item('lip_syncer_model') + return create_static_model_set('full').get(model_name) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--lip-syncer-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'lip_syncer_model', 'wav2lip_gan_96'), choices = lip_syncer_choices.lip_syncer_models) + group_processors.add_argument('--lip-syncer-weight', help = translator.get('help.weight', __package__), type = float, default = config.get_float_value('processors', 'lip_syncer_weight', '0.5'), choices = lip_syncer_choices.lip_syncer_weight_range, metavar = create_float_metavar(lip_syncer_choices.lip_syncer_weight_range)) + facefusion.jobs.job_store.register_step_keys([ 'lip_syncer_model', 'lip_syncer_weight' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('lip_syncer_model', args.get('lip_syncer_model')) + apply_state_item('lip_syncer_weight', args.get('lip_syncer_weight')) + + +def pre_check() -> bool: + model_hash_set = get_model_options().get('hashes') + model_source_set = get_model_options().get('sources') + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def pre_process(mode : ProcessMode) -> bool: + if not has_audio(state_manager.get_item('source_paths')): + logger.error(translator.get('choose_audio_source') + translator.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_video_frame.cache_clear() + read_static_voice.cache_clear() + video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': + content_analyser.clear_inference_pool() + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_masker.clear_inference_pool() + face_recognizer.clear_inference_pool() + voice_extractor.clear_inference_pool() + + +def sync_lip(target_face : Face, source_voice_frame : AudioFrame, temp_vision_frame : VisionFrame) -> VisionFrame: + model_type = get_model_options().get('type') + model_size = get_model_options().get('size') + source_voice_frame = prepare_audio_frame(source_voice_frame) + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmark_set.get('5/68'), 'ffhq_512', (512, 512)) + crop_masks = [] + + if 'occlusion' in state_manager.get_item('face_mask_types'): + occlusion_mask = create_occlusion_mask(crop_vision_frame) + crop_masks.append(occlusion_mask) + + if model_type == 'edtalk': + lip_syncer_weight = numpy.array([ state_manager.get_item('lip_syncer_weight') ]).astype(numpy.float32) + box_mask = create_box_mask(crop_vision_frame, state_manager.get_item('face_mask_blur'), state_manager.get_item('face_mask_padding')) + crop_masks.append(box_mask) + crop_vision_frame = prepare_crop_frame(crop_vision_frame) + crop_vision_frame = forward_edtalk(source_voice_frame, crop_vision_frame, lip_syncer_weight) + crop_vision_frame = normalize_crop_frame(crop_vision_frame) + + if model_type == 'wav2lip': + face_landmark_68 = cv2.transform(target_face.landmark_set.get('68').reshape(1, -1, 2), affine_matrix).reshape(-1, 2) + area_mask = create_area_mask(crop_vision_frame, face_landmark_68, [ 'lower-face' ]) + crop_masks.append(area_mask) + bounding_box = create_bounding_box(face_landmark_68) + area_vision_frame, area_matrix = warp_face_by_bounding_box(crop_vision_frame, bounding_box, model_size) + area_vision_frame = prepare_crop_frame(area_vision_frame) + area_vision_frame = forward_wav2lip(source_voice_frame, area_vision_frame) + area_vision_frame = normalize_crop_frame(area_vision_frame) + crop_vision_frame = cv2.warpAffine(area_vision_frame, cv2.invertAffineTransform(area_matrix), (512, 512), borderMode = cv2.BORDER_REPLICATE) + + crop_mask = numpy.minimum.reduce(crop_masks) + paste_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, crop_mask, affine_matrix) + return paste_vision_frame + + +def forward_edtalk(temp_audio_frame : AudioFrame, crop_vision_frame : VisionFrame, lip_syncer_weight : LipSyncerWeight) -> VisionFrame: + lip_syncer = get_inference_pool().get('lip_syncer') + + with conditional_thread_semaphore(): + crop_vision_frame = lip_syncer.run(None, + { + 'source': temp_audio_frame, + 'target': crop_vision_frame, + 'weight': lip_syncer_weight + })[0] + + return crop_vision_frame + + +def forward_wav2lip(temp_audio_frame : AudioFrame, area_vision_frame : VisionFrame) -> VisionFrame: + lip_syncer = get_inference_pool().get('lip_syncer') + + with conditional_thread_semaphore(): + area_vision_frame = lip_syncer.run(None, + { + 'source': temp_audio_frame, + 'target': area_vision_frame + })[0] + + return area_vision_frame + + +def prepare_audio_frame(temp_audio_frame : AudioFrame) -> AudioFrame: + model_type = get_model_options().get('type') + temp_audio_frame = numpy.maximum(numpy.exp(-5 * numpy.log(10)), temp_audio_frame) + temp_audio_frame = numpy.log10(temp_audio_frame) * 1.6 + 3.2 + temp_audio_frame = temp_audio_frame.clip(-4, 4).astype(numpy.float32) + + if model_type == 'wav2lip': + temp_audio_frame = temp_audio_frame * state_manager.get_item('lip_syncer_weight') * 2.0 + + temp_audio_frame = numpy.expand_dims(temp_audio_frame, axis = (0, 1)) + return temp_audio_frame + + +def prepare_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + model_type = get_model_options().get('type') + model_size = get_model_options().get('size') + + if model_type == 'edtalk': + crop_vision_frame = cv2.resize(crop_vision_frame, model_size, interpolation = cv2.INTER_AREA) + crop_vision_frame = crop_vision_frame[:, :, ::-1] / 255.0 + crop_vision_frame = numpy.expand_dims(crop_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) + + if model_type == 'wav2lip': + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0) + prepare_vision_frame = crop_vision_frame.copy() + prepare_vision_frame[:, model_size[0] // 2:] = 0 + crop_vision_frame = numpy.concatenate((prepare_vision_frame, crop_vision_frame), axis = 3) + crop_vision_frame = crop_vision_frame.transpose(0, 3, 1, 2).astype(numpy.float32) / 255.0 + + return crop_vision_frame + + +def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + model_type = get_model_options().get('type') + crop_vision_frame = crop_vision_frame[0].transpose(1, 2, 0) + crop_vision_frame = crop_vision_frame.clip(0, 1) * 255 + crop_vision_frame = crop_vision_frame.astype(numpy.uint8) + + if model_type == 'edtalk': + crop_vision_frame = crop_vision_frame[:, :, ::-1] + crop_vision_frame = cv2.resize(crop_vision_frame, (512, 512), interpolation = cv2.INTER_CUBIC) + + return crop_vision_frame + + +def process_frame(inputs : LipSyncerInputs) -> ProcessorOutputs: + reference_vision_frame = inputs.get('reference_vision_frame') + source_voice_frame = inputs.get('source_voice_frame') + target_vision_frame = inputs.get('target_vision_frame') + temp_vision_frame = inputs.get('temp_vision_frame') + temp_vision_mask = inputs.get('temp_vision_mask') + target_faces = select_faces(reference_vision_frame, target_vision_frame) + + if target_faces: + for target_face in target_faces: + target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) + temp_vision_frame = sync_lip(target_face, source_voice_frame, temp_vision_frame) + + return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/lip_syncer/locales.py b/facefusion/processors/modules/lip_syncer/locales.py new file mode 100644 index 0000000000000000000000000000000000000000..75c4d0b4fb3b456a74a1d311129335cb068a28b0 --- /dev/null +++ b/facefusion/processors/modules/lip_syncer/locales.py @@ -0,0 +1,18 @@ +from facefusion.types import Locales + +LOCALES : Locales =\ +{ + 'en': + { + 'help': + { + 'model': 'choose the model responsible for syncing the lips', + 'weight': 'specify the degree of weight applied to the lips' + }, + 'uis': + { + 'model_dropdown': 'LIP SYNCER MODEL', + 'weight_slider': 'LIP SYNCER WEIGHT' + } + } +} diff --git a/facefusion/processors/modules/lip_syncer/types.py b/facefusion/processors/modules/lip_syncer/types.py new file mode 100644 index 0000000000000000000000000000000000000000..88fab696af7af862a5c28c444a700ca908a45576 --- /dev/null +++ b/facefusion/processors/modules/lip_syncer/types.py @@ -0,0 +1,18 @@ +from typing import Any, Literal, TypeAlias, TypedDict + +from numpy.typing import NDArray + +from facefusion.types import AudioFrame, Mask, VisionFrame + +LipSyncerInputs = TypedDict('LipSyncerInputs', +{ + 'reference_vision_frame' : VisionFrame, + 'source_voice_frame' : AudioFrame, + 'target_vision_frame' : VisionFrame, + 'temp_vision_frame' : VisionFrame, + 'temp_vision_mask' : Mask +}) + +LipSyncerModel = Literal['edtalk_256', 'wav2lip_96', 'wav2lip_gan_96'] + +LipSyncerWeight : TypeAlias = NDArray[Any] diff --git a/facefusion/processors/pixel_boost.py b/facefusion/processors/pixel_boost.py new file mode 100644 index 0000000000000000000000000000000000000000..d31c6ae07b7bde9d23a4402a9c8236ea484ac7e5 --- /dev/null +++ b/facefusion/processors/pixel_boost.py @@ -0,0 +1,18 @@ +from typing import List + +import numpy +from cv2.typing import Size + +from facefusion.types import VisionFrame + + +def implode_pixel_boost(crop_vision_frame : VisionFrame, pixel_boost_total : int, model_size : Size) -> VisionFrame: + pixel_boost_vision_frame = crop_vision_frame.reshape(model_size[0], pixel_boost_total, model_size[1], pixel_boost_total, 3) + pixel_boost_vision_frame = pixel_boost_vision_frame.transpose(1, 3, 0, 2, 4).reshape(pixel_boost_total ** 2, model_size[0], model_size[1], 3) + return pixel_boost_vision_frame + + +def explode_pixel_boost(temp_vision_frames : List[VisionFrame], pixel_boost_total : int, model_size : Size, pixel_boost_size : Size) -> VisionFrame: + crop_vision_frame = numpy.stack(temp_vision_frames).reshape(pixel_boost_total, pixel_boost_total, model_size[0], model_size[1], 3) + crop_vision_frame = crop_vision_frame.transpose(2, 0, 3, 1, 4).reshape(pixel_boost_size[0], pixel_boost_size[1], 3) + return crop_vision_frame diff --git a/facefusion/processors/types.py b/facefusion/processors/types.py new file mode 100644 index 0000000000000000000000000000000000000000..b2601f8a83109e5279a361c734f2ae98d4a07591 --- /dev/null +++ b/facefusion/processors/types.py @@ -0,0 +1,20 @@ +from typing import Any, Dict, Tuple, TypeAlias + +from numpy.typing import NDArray + +from facefusion.types import AppContext, Mask, VisionFrame + +LivePortraitPitch : TypeAlias = float +LivePortraitYaw : TypeAlias = float +LivePortraitRoll : TypeAlias = float +LivePortraitExpression : TypeAlias = NDArray[Any] +LivePortraitFeatureVolume : TypeAlias = NDArray[Any] +LivePortraitMotionPoints : TypeAlias = NDArray[Any] +LivePortraitRotation : TypeAlias = NDArray[Any] +LivePortraitScale : TypeAlias = NDArray[Any] +LivePortraitTranslation : TypeAlias = NDArray[Any] + +ProcessorStateKey = str +ProcessorState : TypeAlias = Dict[ProcessorStateKey, Any] +ProcessorStateSet : TypeAlias = Dict[AppContext, ProcessorState] +ProcessorOutputs : TypeAlias = Tuple[VisionFrame, Mask] diff --git a/facefusion/program.py b/facefusion/program.py new file mode 100644 index 0000000000000000000000000000000000000000..7fa1009cf4b40b757a3014a3606923a734b51d92 --- /dev/null +++ b/facefusion/program.py @@ -0,0 +1,327 @@ +import tempfile +from argparse import ArgumentParser, HelpFormatter +from functools import partial + +import facefusion.choices +from facefusion import config, metadata, state_manager, translator +from facefusion.common_helper import create_float_metavar, create_int_metavar, get_first, get_last +from facefusion.execution import get_available_execution_providers +from facefusion.ffmpeg import get_available_encoder_set +from facefusion.filesystem import get_file_name, resolve_file_paths +from facefusion.jobs import job_store +from facefusion.processors.core import get_processors_modules +from facefusion.sanitizer import sanitize_int_range + + +def create_help_formatter_small(prog : str) -> HelpFormatter: + return HelpFormatter(prog, max_help_position = 50) + + +def create_help_formatter_large(prog : str) -> HelpFormatter: + return HelpFormatter(prog, max_help_position = 300) + + +def create_config_path_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_paths = program.add_argument_group('paths') + group_paths.add_argument('--config-path', help = translator.get('help.config_path'), default = 'facefusion.ini') + job_store.register_job_keys([ 'config_path' ]) + apply_config_path(program) + return program + + +def create_temp_path_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_paths = program.add_argument_group('paths') + group_paths.add_argument('--temp-path', help = translator.get('help.temp_path'), default = config.get_str_value('paths', 'temp_path', tempfile.gettempdir())) + job_store.register_job_keys([ 'temp_path' ]) + return program + + +def create_jobs_path_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_paths = program.add_argument_group('paths') + group_paths.add_argument('--jobs-path', help = translator.get('help.jobs_path'), default = config.get_str_value('paths', 'jobs_path', '.jobs')) + job_store.register_job_keys([ 'jobs_path' ]) + return program + + +def create_source_paths_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_paths = program.add_argument_group('paths') + group_paths.add_argument('-s', '--source-paths', help = translator.get('help.source_paths'), default = config.get_str_list('paths', 'source_paths'), nargs = '+') + job_store.register_step_keys([ 'source_paths' ]) + return program + + +def create_target_path_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_paths = program.add_argument_group('paths') + group_paths.add_argument('-t', '--target-path', help = translator.get('help.target_path'), default = config.get_str_value('paths', 'target_path')) + job_store.register_step_keys([ 'target_path' ]) + return program + + +def create_output_path_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_paths = program.add_argument_group('paths') + group_paths.add_argument('-o', '--output-path', help = translator.get('help.output_path'), default = config.get_str_value('paths', 'output_path')) + job_store.register_step_keys([ 'output_path' ]) + return program + + +def create_source_pattern_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_patterns = program.add_argument_group('patterns') + group_patterns.add_argument('-s', '--source-pattern', help = translator.get('help.source_pattern'), default = config.get_str_value('patterns', 'source_pattern')) + job_store.register_job_keys([ 'source_pattern' ]) + return program + + +def create_target_pattern_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_patterns = program.add_argument_group('patterns') + group_patterns.add_argument('-t', '--target-pattern', help = translator.get('help.target_pattern'), default = config.get_str_value('patterns', 'target_pattern')) + job_store.register_job_keys([ 'target_pattern' ]) + return program + + +def create_output_pattern_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_patterns = program.add_argument_group('patterns') + group_patterns.add_argument('-o', '--output-pattern', help = translator.get('help.output_pattern'), default = config.get_str_value('patterns', 'output_pattern')) + job_store.register_job_keys([ 'output_pattern' ]) + return program + + +def create_face_detector_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_face_detector = program.add_argument_group('face detector') + group_face_detector.add_argument('--face-detector-model', help = translator.get('help.face_detector_model'), default = config.get_str_value('face_detector', 'face_detector_model', 'yolo_face'), choices = facefusion.choices.face_detector_models) + known_args, _ = program.parse_known_args() + face_detector_size_choices = facefusion.choices.face_detector_set.get(known_args.face_detector_model) + group_face_detector.add_argument('--face-detector-size', help = translator.get('help.face_detector_size'), default = config.get_str_value('face_detector', 'face_detector_size', get_last(face_detector_size_choices)), choices = face_detector_size_choices) + group_face_detector.add_argument('--face-detector-margin', help = translator.get('help.face_detector_margin'), type = partial(sanitize_int_range, int_range = facefusion.choices.face_detector_margin_range), default = config.get_int_list('face_detector', 'face_detector_margin', '0 0 0 0'), nargs = '+') + group_face_detector.add_argument('--face-detector-angles', help = translator.get('help.face_detector_angles'), type = int, default = config.get_int_list('face_detector', 'face_detector_angles', '0'), choices = facefusion.choices.face_detector_angles, nargs = '+', metavar = 'FACE_DETECTOR_ANGLES') + group_face_detector.add_argument('--face-detector-score', help = translator.get('help.face_detector_score'), type = float, default = config.get_float_value('face_detector', 'face_detector_score', '0.5'), choices = facefusion.choices.face_detector_score_range, metavar = create_float_metavar(facefusion.choices.face_detector_score_range)) + job_store.register_step_keys([ 'face_detector_model', 'face_detector_size', 'face_detector_margin', 'face_detector_angles', 'face_detector_score' ]) + return program + + +def create_face_landmarker_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_face_landmarker = program.add_argument_group('face landmarker') + group_face_landmarker.add_argument('--face-landmarker-model', help = translator.get('help.face_landmarker_model'), default = config.get_str_value('face_landmarker', 'face_landmarker_model', '2dfan4'), choices = facefusion.choices.face_landmarker_models) + group_face_landmarker.add_argument('--face-landmarker-score', help = translator.get('help.face_landmarker_score'), type = float, default = config.get_float_value('face_landmarker', 'face_landmarker_score', '0.5'), choices = facefusion.choices.face_landmarker_score_range, metavar = create_float_metavar(facefusion.choices.face_landmarker_score_range)) + job_store.register_step_keys([ 'face_landmarker_model', 'face_landmarker_score' ]) + return program + + +def create_face_selector_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_face_selector = program.add_argument_group('face selector') + group_face_selector.add_argument('--face-selector-mode', help = translator.get('help.face_selector_mode'), default = config.get_str_value('face_selector', 'face_selector_mode', 'reference'), choices = facefusion.choices.face_selector_modes) + group_face_selector.add_argument('--face-selector-order', help = translator.get('help.face_selector_order'), default = config.get_str_value('face_selector', 'face_selector_order', 'large-small'), choices = facefusion.choices.face_selector_orders) + group_face_selector.add_argument('--face-selector-age-start', help = translator.get('help.face_selector_age_start'), type = int, default = config.get_int_value('face_selector', 'face_selector_age_start'), choices = facefusion.choices.face_selector_age_range, metavar = create_int_metavar(facefusion.choices.face_selector_age_range)) + group_face_selector.add_argument('--face-selector-age-end', help = translator.get('help.face_selector_age_end'), type = int, default = config.get_int_value('face_selector', 'face_selector_age_end'), choices = facefusion.choices.face_selector_age_range, metavar = create_int_metavar(facefusion.choices.face_selector_age_range)) + group_face_selector.add_argument('--face-selector-gender', help = translator.get('help.face_selector_gender'), default = config.get_str_value('face_selector', 'face_selector_gender'), choices = facefusion.choices.face_selector_genders) + group_face_selector.add_argument('--face-selector-race', help = translator.get('help.face_selector_race'), default = config.get_str_value('face_selector', 'face_selector_race'), choices = facefusion.choices.face_selector_races) + group_face_selector.add_argument('--reference-face-position', help = translator.get('help.reference_face_position'), type = int, default = config.get_int_value('face_selector', 'reference_face_position', '0')) + group_face_selector.add_argument('--reference-face-distance', help = translator.get('help.reference_face_distance'), type = float, default = config.get_float_value('face_selector', 'reference_face_distance', '0.3'), choices = facefusion.choices.reference_face_distance_range, metavar = create_float_metavar(facefusion.choices.reference_face_distance_range)) + group_face_selector.add_argument('--reference-frame-number', help = translator.get('help.reference_frame_number'), type = int, default = config.get_int_value('face_selector', 'reference_frame_number', '0')) + job_store.register_step_keys([ 'face_selector_mode', 'face_selector_order', 'face_selector_gender', 'face_selector_race', 'face_selector_age_start', 'face_selector_age_end', 'reference_face_position', 'reference_face_distance', 'reference_frame_number' ]) + return program + + +def create_face_masker_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_face_masker = program.add_argument_group('face masker') + group_face_masker.add_argument('--face-occluder-model', help = translator.get('help.face_occluder_model'), default = config.get_str_value('face_masker', 'face_occluder_model', 'xseg_1'), choices = facefusion.choices.face_occluder_models) + group_face_masker.add_argument('--face-parser-model', help = translator.get('help.face_parser_model'), default = config.get_str_value('face_masker', 'face_parser_model', 'bisenet_resnet_34'), choices = facefusion.choices.face_parser_models) + group_face_masker.add_argument('--face-mask-types', help = translator.get('help.face_mask_types').format(choices = ', '.join(facefusion.choices.face_mask_types)), default = config.get_str_list('face_masker', 'face_mask_types', 'box'), choices = facefusion.choices.face_mask_types, nargs = '+', metavar = 'FACE_MASK_TYPES') + group_face_masker.add_argument('--face-mask-areas', help = translator.get('help.face_mask_areas').format(choices = ', '.join(facefusion.choices.face_mask_areas)), default = config.get_str_list('face_masker', 'face_mask_areas', ' '.join(facefusion.choices.face_mask_areas)), choices = facefusion.choices.face_mask_areas, nargs = '+', metavar = 'FACE_MASK_AREAS') + group_face_masker.add_argument('--face-mask-regions', help = translator.get('help.face_mask_regions').format(choices = ', '.join(facefusion.choices.face_mask_regions)), default = config.get_str_list('face_masker', 'face_mask_regions', ' '.join(facefusion.choices.face_mask_regions)), choices = facefusion.choices.face_mask_regions, nargs = '+', metavar = 'FACE_MASK_REGIONS') + group_face_masker.add_argument('--face-mask-blur', help = translator.get('help.face_mask_blur'), type = float, default = config.get_float_value('face_masker', 'face_mask_blur', '0.3'), choices = facefusion.choices.face_mask_blur_range, metavar = create_float_metavar(facefusion.choices.face_mask_blur_range)) + group_face_masker.add_argument('--face-mask-padding', help = translator.get('help.face_mask_padding'), type = partial(sanitize_int_range, int_range = facefusion.choices.face_mask_padding_range), default = config.get_int_list('face_masker', 'face_mask_padding', '0 0 0 0'), nargs = '+') + job_store.register_step_keys([ 'face_occluder_model', 'face_parser_model', 'face_mask_types', 'face_mask_areas', 'face_mask_regions', 'face_mask_blur', 'face_mask_padding' ]) + return program + + +def create_voice_extractor_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_voice_extractor = program.add_argument_group('voice extractor') + group_voice_extractor.add_argument('--voice-extractor-model', help = translator.get('help.voice_extractor_model'), default = config.get_str_value('voice_extractor', 'voice_extractor_model', 'kim_vocal_2'), choices = facefusion.choices.voice_extractor_models) + job_store.register_step_keys([ 'voice_extractor_model' ]) + return program + + +def create_frame_extraction_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_frame_extraction = program.add_argument_group('frame extraction') + group_frame_extraction.add_argument('--trim-frame-start', help = translator.get('help.trim_frame_start'), type = int, default = facefusion.config.get_int_value('frame_extraction', 'trim_frame_start')) + group_frame_extraction.add_argument('--trim-frame-end', help = translator.get('help.trim_frame_end'), type = int, default = facefusion.config.get_int_value('frame_extraction', 'trim_frame_end')) + group_frame_extraction.add_argument('--temp-frame-format', help = translator.get('help.temp_frame_format'), default = config.get_str_value('frame_extraction', 'temp_frame_format', 'png'), choices = facefusion.choices.temp_frame_formats) + group_frame_extraction.add_argument('--keep-temp', help = translator.get('help.keep_temp'), action = 'store_true', default = config.get_bool_value('frame_extraction', 'keep_temp')) + job_store.register_step_keys([ 'trim_frame_start', 'trim_frame_end', 'temp_frame_format', 'keep_temp' ]) + return program + + +def create_output_creation_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + available_encoder_set = get_available_encoder_set() + group_output_creation = program.add_argument_group('output creation') + group_output_creation.add_argument('--output-image-quality', help = translator.get('help.output_image_quality'), type = int, default = config.get_int_value('output_creation', 'output_image_quality', '80'), choices = facefusion.choices.output_image_quality_range, metavar = create_int_metavar(facefusion.choices.output_image_quality_range)) + group_output_creation.add_argument('--output-image-scale', help = translator.get('help.output_image_scale'), type = float, default = config.get_float_value('output_creation', 'output_image_scale', '1.0'), choices = facefusion.choices.output_image_scale_range) + group_output_creation.add_argument('--output-audio-encoder', help = translator.get('help.output_audio_encoder'), default = config.get_str_value('output_creation', 'output_audio_encoder', get_first(available_encoder_set.get('audio'))), choices = available_encoder_set.get('audio')) + group_output_creation.add_argument('--output-audio-quality', help = translator.get('help.output_audio_quality'), type = int, default = config.get_int_value('output_creation', 'output_audio_quality', '80'), choices = facefusion.choices.output_audio_quality_range, metavar = create_int_metavar(facefusion.choices.output_audio_quality_range)) + group_output_creation.add_argument('--output-audio-volume', help = translator.get('help.output_audio_volume'), type = int, default = config.get_int_value('output_creation', 'output_audio_volume', '100'), choices = facefusion.choices.output_audio_volume_range, metavar = create_int_metavar(facefusion.choices.output_audio_volume_range)) + group_output_creation.add_argument('--output-video-encoder', help = translator.get('help.output_video_encoder'), default = config.get_str_value('output_creation', 'output_video_encoder', get_first(available_encoder_set.get('video'))), choices = available_encoder_set.get('video')) + group_output_creation.add_argument('--output-video-preset', help = translator.get('help.output_video_preset'), default = config.get_str_value('output_creation', 'output_video_preset', 'veryfast'), choices = facefusion.choices.output_video_presets) + group_output_creation.add_argument('--output-video-quality', help = translator.get('help.output_video_quality'), type = int, default = config.get_int_value('output_creation', 'output_video_quality', '80'), choices = facefusion.choices.output_video_quality_range, metavar = create_int_metavar(facefusion.choices.output_video_quality_range)) + group_output_creation.add_argument('--output-video-scale', help = translator.get('help.output_video_scale'), type = float, default = config.get_float_value('output_creation', 'output_video_scale', '1.0'), choices = facefusion.choices.output_video_scale_range) + group_output_creation.add_argument('--output-video-fps', help = translator.get('help.output_video_fps'), type = float, default = config.get_float_value('output_creation', 'output_video_fps')) + job_store.register_step_keys([ 'output_image_quality', 'output_image_scale', 'output_audio_encoder', 'output_audio_quality', 'output_audio_volume', 'output_video_encoder', 'output_video_preset', 'output_video_quality', 'output_video_scale', 'output_video_fps' ]) + return program + + +def create_processors_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + available_processors = [ get_file_name(file_path) for file_path in resolve_file_paths('facefusion/processors/modules') ] + group_processors = program.add_argument_group('processors') + group_processors.add_argument('--processors', help = translator.get('help.processors').format(choices = ', '.join(available_processors)), default = config.get_str_list('processors', 'processors', 'face_swapper'), nargs = '+') + job_store.register_step_keys([ 'processors' ]) + for processor_module in get_processors_modules(available_processors): + processor_module.register_args(program) + return program + + +def create_uis_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + available_ui_layouts = [ get_file_name(file_path) for file_path in resolve_file_paths('facefusion/uis/layouts') ] + group_uis = program.add_argument_group('uis') + group_uis.add_argument('--open-browser', help = translator.get('help.open_browser'), action = 'store_true', default = config.get_bool_value('uis', 'open_browser')) + group_uis.add_argument('--ui-layouts', help = translator.get('help.ui_layouts').format(choices = ', '.join(available_ui_layouts)), default = config.get_str_list('uis', 'ui_layouts', 'default'), nargs = '+') + group_uis.add_argument('--ui-workflow', help = translator.get('help.ui_workflow'), default = config.get_str_value('uis', 'ui_workflow', 'instant_runner'), choices = facefusion.choices.ui_workflows) + return program + + +def create_download_providers_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_download = program.add_argument_group('download') + group_download.add_argument('--download-providers', help = translator.get('help.download_providers').format(choices = ', '.join(facefusion.choices.download_providers)), default = config.get_str_list('download', 'download_providers', ' '.join(facefusion.choices.download_providers)), choices = facefusion.choices.download_providers, nargs = '+', metavar = 'DOWNLOAD_PROVIDERS') + job_store.register_job_keys([ 'download_providers' ]) + return program + + +def create_download_scope_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_download = program.add_argument_group('download') + group_download.add_argument('--download-scope', help = translator.get('help.download_scope'), default = config.get_str_value('download', 'download_scope', 'lite'), choices = facefusion.choices.download_scopes) + job_store.register_job_keys([ 'download_scope' ]) + return program + + +def create_benchmark_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_benchmark = program.add_argument_group('benchmark') + group_benchmark.add_argument('--benchmark-mode', help = translator.get('help.benchmark_mode'), default = config.get_str_value('benchmark', 'benchmark_mode', 'warm'), choices = facefusion.choices.benchmark_modes) + group_benchmark.add_argument('--benchmark-resolutions', help = translator.get('help.benchmark_resolutions'), default = config.get_str_list('benchmark', 'benchmark_resolutions', get_first(facefusion.choices.benchmark_resolutions)), choices = facefusion.choices.benchmark_resolutions, nargs = '+') + group_benchmark.add_argument('--benchmark-cycle-count', help = translator.get('help.benchmark_cycle_count'), type = int, default = config.get_int_value('benchmark', 'benchmark_cycle_count', '5'), choices = facefusion.choices.benchmark_cycle_count_range) + return program + + +def create_execution_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + available_execution_providers = get_available_execution_providers() + group_execution = program.add_argument_group('execution') + group_execution.add_argument('--execution-device-ids', help = translator.get('help.execution_device_ids'), type = int, default = config.get_int_list('execution', 'execution_device_ids', '0'), nargs = '+', metavar = 'EXECUTION_DEVICE_IDS') + group_execution.add_argument('--execution-providers', help = translator.get('help.execution_providers').format(choices = ', '.join(available_execution_providers)), default = config.get_str_list('execution', 'execution_providers', get_first(available_execution_providers)), choices = available_execution_providers, nargs = '+', metavar = 'EXECUTION_PROVIDERS') + group_execution.add_argument('--execution-thread-count', help = translator.get('help.execution_thread_count'), type = int, default = config.get_int_value('execution', 'execution_thread_count', '8'), choices = facefusion.choices.execution_thread_count_range, metavar = create_int_metavar(facefusion.choices.execution_thread_count_range)) + job_store.register_job_keys([ 'execution_device_ids', 'execution_providers', 'execution_thread_count' ]) + return program + + +def create_memory_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_memory = program.add_argument_group('memory') + group_memory.add_argument('--video-memory-strategy', help = translator.get('help.video_memory_strategy'), default = config.get_str_value('memory', 'video_memory_strategy', 'strict'), choices = facefusion.choices.video_memory_strategies) + group_memory.add_argument('--system-memory-limit', help = translator.get('help.system_memory_limit'), type = int, default = config.get_int_value('memory', 'system_memory_limit', '0'), choices = facefusion.choices.system_memory_limit_range, metavar = create_int_metavar(facefusion.choices.system_memory_limit_range)) + job_store.register_job_keys([ 'video_memory_strategy', 'system_memory_limit' ]) + return program + + +def create_log_level_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_misc = program.add_argument_group('misc') + group_misc.add_argument('--log-level', help = translator.get('help.log_level'), default = config.get_str_value('misc', 'log_level', 'info'), choices = facefusion.choices.log_levels) + job_store.register_job_keys([ 'log_level' ]) + return program + + +def create_halt_on_error_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_misc = program.add_argument_group('misc') + group_misc.add_argument('--halt-on-error', help = translator.get('help.halt_on_error'), action = 'store_true', default = config.get_bool_value('misc', 'halt_on_error')) + job_store.register_job_keys([ 'halt_on_error' ]) + return program + + +def create_job_id_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + program.add_argument('job_id', help = translator.get('help.job_id')) + return program + + +def create_job_status_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + program.add_argument('job_status', help = translator.get('help.job_status'), choices = facefusion.choices.job_statuses) + return program + + +def create_step_index_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + program.add_argument('step_index', help = translator.get('help.step_index'), type = int) + return program + + +def collect_step_program() -> ArgumentParser: + return ArgumentParser(parents = [ create_face_detector_program(), create_face_landmarker_program(), create_face_selector_program(), create_face_masker_program(), create_voice_extractor_program(), create_frame_extraction_program(), create_output_creation_program(), create_processors_program() ], add_help = False) + + +def collect_job_program() -> ArgumentParser: + return ArgumentParser(parents = [ create_execution_program(), create_download_providers_program(), create_memory_program(), create_log_level_program() ], add_help = False) + + +def create_program() -> ArgumentParser: + program = ArgumentParser(formatter_class = create_help_formatter_large, add_help = False) + program._positionals.title = 'commands' + program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version') + sub_program = program.add_subparsers(dest = 'command') + # general + sub_program.add_parser('run', help = translator.get('help.run'), parents = [ create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), create_source_paths_program(), create_target_path_program(), create_output_path_program(), collect_step_program(), create_uis_program(), create_benchmark_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('headless-run', help = translator.get('help.headless_run'), parents = [ create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), create_source_paths_program(), create_target_path_program(), create_output_path_program(), collect_step_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('batch-run', help = translator.get('help.batch_run'), parents = [ create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), create_source_pattern_program(), create_target_pattern_program(), create_output_pattern_program(), collect_step_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('force-download', help = translator.get('help.force_download'), parents = [ create_download_providers_program(), create_download_scope_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('benchmark', help = translator.get('help.benchmark'), parents = [ create_temp_path_program(), collect_step_program(), create_benchmark_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + # job manager + sub_program.add_parser('job-list', help = translator.get('help.job_list'), parents = [ create_job_status_program(), create_jobs_path_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-create', help = translator.get('help.job_create'), parents = [ create_job_id_program(), create_jobs_path_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-submit', help = translator.get('help.job_submit'), parents = [ create_job_id_program(), create_jobs_path_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-submit-all', help = translator.get('help.job_submit_all'), parents = [ create_jobs_path_program(), create_log_level_program(), create_halt_on_error_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-delete', help = translator.get('help.job_delete'), parents = [ create_job_id_program(), create_jobs_path_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-delete-all', help = translator.get('help.job_delete_all'), parents = [ create_jobs_path_program(), create_log_level_program(), create_halt_on_error_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-add-step', help = translator.get('help.job_add_step'), parents = [ create_job_id_program(), create_config_path_program(), create_jobs_path_program(), create_source_paths_program(), create_target_path_program(), create_output_path_program(), collect_step_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-remix-step', help = translator.get('help.job_remix_step'), parents = [ create_job_id_program(), create_step_index_program(), create_config_path_program(), create_jobs_path_program(), create_source_paths_program(), create_output_path_program(), collect_step_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-insert-step', help = translator.get('help.job_insert_step'), parents = [ create_job_id_program(), create_step_index_program(), create_config_path_program(), create_jobs_path_program(), create_source_paths_program(), create_target_path_program(), create_output_path_program(), collect_step_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-remove-step', help = translator.get('help.job_remove_step'), parents = [ create_job_id_program(), create_step_index_program(), create_jobs_path_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + # job runner + sub_program.add_parser('job-run', help = translator.get('help.job_run'), parents = [ create_job_id_program(), create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-run-all', help = translator.get('help.job_run_all'), parents = [ create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), collect_job_program(), create_halt_on_error_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-retry', help = translator.get('help.job_retry'), parents = [ create_job_id_program(), create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-retry-all', help = translator.get('help.job_retry_all'), parents = [ create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), collect_job_program(), create_halt_on_error_program() ], formatter_class = create_help_formatter_large) + return ArgumentParser(parents = [ program ], formatter_class = create_help_formatter_small) + + +def apply_config_path(program : ArgumentParser) -> None: + known_args, _ = program.parse_known_args() + state_manager.init_item('config_path', known_args.config_path) diff --git a/facefusion/program_helper.py b/facefusion/program_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..688354bb9441ec807f885a3ddec2cd86124cb50a --- /dev/null +++ b/facefusion/program_helper.py @@ -0,0 +1,31 @@ +from argparse import ArgumentParser, _ArgumentGroup, _SubParsersAction +from typing import Optional + + +def find_argument_group(program : ArgumentParser, group_name : str) -> Optional[_ArgumentGroup]: + for group in program._action_groups: + if group.title == group_name: + return group + return None + + +def validate_args(program : ArgumentParser) -> bool: + if validate_actions(program): + for action in program._actions: + if isinstance(action, _SubParsersAction): + for _, sub_program in action._name_parser_map.items(): + if not validate_args(sub_program): + return False + return True + return False + + +def validate_actions(program : ArgumentParser) -> bool: + for action in program._actions: + if action.default and action.choices: + if isinstance(action.default, list): + if any(default not in action.choices for default in action.default): + return False + elif action.default not in action.choices: + return False + return True diff --git a/facefusion/sanitizer.py b/facefusion/sanitizer.py new file mode 100644 index 0000000000000000000000000000000000000000..bda78f0698af4a09d7e580b80704ef0ae952cb39 --- /dev/null +++ b/facefusion/sanitizer.py @@ -0,0 +1,7 @@ +from typing import Sequence + + +def sanitize_int_range(value : int, int_range : Sequence[int]) -> int: + if value in int_range: + return value + return int_range[0] diff --git a/facefusion/state_manager.py b/facefusion/state_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..e7f0c8030760d1d01cb9fb13817fdf9ba7eb5b9a --- /dev/null +++ b/facefusion/state_manager.py @@ -0,0 +1,42 @@ +from typing import Any, Union + +from facefusion.app_context import detect_app_context +from facefusion.processors.types import ProcessorState, ProcessorStateKey, ProcessorStateSet +from facefusion.types import State, StateKey, StateSet + +STATE_SET : Union[StateSet, ProcessorStateSet] =\ +{ + 'cli': {}, #type:ignore[assignment] + 'ui': {} #type:ignore[assignment] +} + + +def get_state() -> Union[State, ProcessorState]: + app_context = detect_app_context() + return STATE_SET.get(app_context) + + +def sync_state() -> None: + STATE_SET['cli'] = STATE_SET.get('ui') #type:ignore[assignment] + + +def init_item(key : Union[StateKey, ProcessorStateKey], value : Any) -> None: + STATE_SET['cli'][key] = value #type:ignore[literal-required] + STATE_SET['ui'][key] = value #type:ignore[literal-required] + + +def get_item(key : Union[StateKey, ProcessorStateKey]) -> Any: + return get_state().get(key) #type:ignore[literal-required] + + +def set_item(key : Union[StateKey, ProcessorStateKey], value : Any) -> None: + app_context = detect_app_context() + STATE_SET[app_context][key] = value #type:ignore[literal-required] + + +def sync_item(key : Union[StateKey, ProcessorStateKey]) -> None: + STATE_SET['cli'][key] = STATE_SET.get('ui').get(key) #type:ignore[literal-required] + + +def clear_item(key : Union[StateKey, ProcessorStateKey]) -> None: + set_item(key, None) diff --git a/facefusion/streamer.py b/facefusion/streamer.py new file mode 100644 index 0000000000000000000000000000000000000000..b8c5d02a3f21cd7f4cd6bd1edc4418f6ffbfc5f0 --- /dev/null +++ b/facefusion/streamer.py @@ -0,0 +1,100 @@ +import os +import subprocess +from collections import deque +from concurrent.futures import ThreadPoolExecutor +from typing import Deque, Iterator + +import cv2 +import numpy +from tqdm import tqdm + +from facefusion import ffmpeg_builder, logger, state_manager, translator +from facefusion.audio import create_empty_audio_frame +from facefusion.content_analyser import analyse_stream +from facefusion.ffmpeg import open_ffmpeg +from facefusion.filesystem import is_directory +from facefusion.processors.core import get_processors_modules +from facefusion.types import Fps, StreamMode, VisionFrame +from facefusion.vision import extract_vision_mask, read_static_images + + +def multi_process_capture(camera_capture : cv2.VideoCapture, camera_fps : Fps) -> Iterator[VisionFrame]: + capture_deque : Deque[VisionFrame] = deque() + + 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: + futures = [] + + while camera_capture and camera_capture.isOpened(): + _, capture_frame = camera_capture.read() + if analyse_stream(capture_frame, camera_fps): + camera_capture.release() + + if numpy.any(capture_frame): + future = executor.submit(process_stream_frame, capture_frame) + futures.append(future) + + for future_done in [ future for future in futures if future.done() ]: + capture_frame = future_done.result() + capture_deque.append(capture_frame) + futures.remove(future_done) + + while capture_deque: + progress.update() + 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')) + source_audio_frame = create_empty_audio_frame() + source_voice_frame = create_empty_audio_frame() + temp_vision_frame = target_vision_frame.copy() + temp_vision_mask = extract_vision_mask(temp_vision_frame) + + for processor_module in get_processors_modules(state_manager.get_item('processors')): + logger.disable() + if processor_module.pre_process('stream'): + logger.enable() + temp_vision_frame, temp_vision_mask = processor_module.process_frame( + { + 'source_vision_frames': source_vision_frames, + 'source_audio_frame': source_audio_frame, + 'source_voice_frame': source_voice_frame, + 'target_vision_frame': target_vision_frame, + 'temp_vision_frame': temp_vision_frame, + 'temp_vision_mask': temp_vision_mask + }) + logger.enable() + + return temp_vision_frame + + +def open_stream(stream_mode : StreamMode, stream_resolution : str, stream_fps : Fps) -> subprocess.Popen[bytes]: + commands = ffmpeg_builder.chain( + ffmpeg_builder.capture_video(), + ffmpeg_builder.set_media_resolution(stream_resolution), + ffmpeg_builder.set_input_fps(stream_fps) + ) + + if stream_mode == 'udp': + commands.extend(ffmpeg_builder.set_input('-')) + commands.extend(ffmpeg_builder.set_stream_mode('udp')) + commands.extend(ffmpeg_builder.set_stream_quality(2000)) + commands.extend(ffmpeg_builder.set_output('udp://localhost:27000?pkt_size=1316')) + + if stream_mode == 'v4l2': + device_directory_path = '/sys/devices/virtual/video4linux' + commands.extend(ffmpeg_builder.set_input('-')) + commands.extend(ffmpeg_builder.set_stream_mode('v4l2')) + + if is_directory(device_directory_path): + device_names = os.listdir(device_directory_path) + + for device_name in device_names: + device_path = '/dev/' + device_name + commands.extend(ffmpeg_builder.set_output(device_path)) + + else: + logger.error(translator.get('stream_not_loaded').format(stream_mode = stream_mode), __name__) + + return open_ffmpeg(commands) diff --git a/facefusion/temp_helper.py b/facefusion/temp_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..95f9f5445f168f5e2f20fe06a929735ad6e0ac17 --- /dev/null +++ b/facefusion/temp_helper.py @@ -0,0 +1,43 @@ +import os +from typing import List + +from facefusion import state_manager +from facefusion.filesystem import create_directory, get_file_extension, get_file_name, move_file, remove_directory, resolve_file_pattern + + +def get_temp_file_path(file_path : str) -> str: + temp_directory_path = get_temp_directory_path(file_path) + temp_file_extension = get_file_extension(file_path) + return os.path.join(temp_directory_path, 'temp' + temp_file_extension) + + +def move_temp_file(file_path : str, move_path : str) -> bool: + temp_file_path = get_temp_file_path(file_path) + return move_file(temp_file_path, move_path) + + +def resolve_temp_frame_paths(target_path : str) -> List[str]: + temp_frames_pattern = get_temp_frames_pattern(target_path, '*') + return resolve_file_pattern(temp_frames_pattern) + + +def get_temp_frames_pattern(target_path : str, temp_frame_prefix : str) -> str: + temp_directory_path = get_temp_directory_path(target_path) + return os.path.join(temp_directory_path, temp_frame_prefix + '.' + state_manager.get_item('temp_frame_format')) + + +def get_temp_directory_path(file_path : str) -> str: + temp_file_name = get_file_name(file_path) + return os.path.join(state_manager.get_item('temp_path'), 'facefusion', temp_file_name) + + +def create_temp_directory(file_path : str) -> bool: + temp_directory_path = get_temp_directory_path(file_path) + return create_directory(temp_directory_path) + + +def clear_temp_directory(file_path : str) -> bool: + if not state_manager.get_item('keep_temp'): + temp_directory_path = get_temp_directory_path(file_path) + return remove_directory(temp_directory_path) + return True diff --git a/facefusion/thread_helper.py b/facefusion/thread_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..84816c6008d83f28fdb1853504e7c25440d3a361 --- /dev/null +++ b/facefusion/thread_helper.py @@ -0,0 +1,24 @@ +import threading +from contextlib import nullcontext +from typing import ContextManager, Union + +from facefusion.common_helper import is_linux, is_windows +from facefusion.execution import has_execution_provider + +THREAD_LOCK : threading.Lock = threading.Lock() +THREAD_SEMAPHORE : threading.Semaphore = threading.Semaphore() +NULL_CONTEXT : ContextManager[None] = nullcontext() + + +def thread_lock() -> threading.Lock: + return THREAD_LOCK + + +def thread_semaphore() -> threading.Semaphore: + return THREAD_SEMAPHORE + + +def conditional_thread_semaphore() -> Union[threading.Semaphore, ContextManager[None]]: + if is_windows() and has_execution_provider('directml') or is_linux() and has_execution_provider('migraphx') or is_linux() and has_execution_provider('rocm'): + return THREAD_SEMAPHORE + return NULL_CONTEXT diff --git a/facefusion/time_helper.py b/facefusion/time_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..e4134182d7386dfa1989f36313365a8284aa4589 --- /dev/null +++ b/facefusion/time_helper.py @@ -0,0 +1,33 @@ +from datetime import datetime, timedelta +from time import time +from typing import Optional, Tuple + +from facefusion import translator + + +def get_current_date_time() -> datetime: + return datetime.now().astimezone() + + +def calculate_end_time(start_time : float) -> float: + return round(time() - start_time, 2) + + +def split_time_delta(time_delta : timedelta) -> Tuple[int, int, int, int]: + days, hours = divmod(time_delta.total_seconds(), 86400) + hours, minutes = divmod(hours, 3600) + minutes, seconds = divmod(minutes, 60) + return int(days), int(hours), int(minutes), int(seconds) + + +def describe_time_ago(date_time : datetime) -> Optional[str]: + time_ago = datetime.now(date_time.tzinfo) - date_time + days, hours, minutes, _ = split_time_delta(time_ago) + + if timedelta(days = 1) < time_ago: + return translator.get('time_ago_days').format(days = days, hours = hours, minutes = minutes) + if timedelta(hours = 1) < time_ago: + return translator.get('time_ago_hours').format(hours = hours, minutes = minutes) + if timedelta(minutes = 1) < time_ago: + return translator.get('time_ago_minutes').format(minutes = minutes) + return translator.get('time_ago_now') diff --git a/facefusion/translator.py b/facefusion/translator.py new file mode 100644 index 0000000000000000000000000000000000000000..d05a5427381237c51824e906db98e06fac9c94f1 --- /dev/null +++ b/facefusion/translator.py @@ -0,0 +1,35 @@ +import importlib +from typing import Optional + +from facefusion.types import Language, LocalePoolSet, Locales + +LOCALE_POOL_SET : LocalePoolSet = {} +CURRENT_LANGUAGE : Language = 'en' + + +def __autoload__(module_name : str) -> None: + try: + __locales__ = importlib.import_module(module_name + '.locales') + load(__locales__.LOCALES, module_name) + except ImportError: + pass + + +def load(__locales__ : Locales, module_name : str) -> None: + LOCALE_POOL_SET[module_name] = __locales__ + + +def get(notation : str, module_name : str = 'facefusion') -> Optional[str]: + if module_name not in LOCALE_POOL_SET: + __autoload__(module_name) + + current = LOCALE_POOL_SET.get(module_name).get(CURRENT_LANGUAGE) + + for fragment in notation.split('.'): + if fragment in current: + current = current.get(fragment) + + if isinstance(current, str): + return current + + return None diff --git a/facefusion/types.py b/facefusion/types.py new file mode 100644 index 0000000000000000000000000000000000000000..eadb524b5c8be042e43d6c29f968cd804c725ba1 --- /dev/null +++ b/facefusion/types.py @@ -0,0 +1,401 @@ +from collections import namedtuple +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, TypeAlias, TypedDict + +import cv2 +import numpy +from numpy.typing import NDArray +from onnxruntime import InferenceSession + +Scale : TypeAlias = float +Score : TypeAlias = float +Angle : TypeAlias = int + +Detection : TypeAlias = NDArray[Any] +Prediction : TypeAlias = NDArray[Any] + +BoundingBox : TypeAlias = NDArray[Any] +FaceLandmark5 : TypeAlias = NDArray[Any] +FaceLandmark68 : TypeAlias = NDArray[Any] +FaceLandmarkSet = TypedDict('FaceLandmarkSet', +{ + '5' : FaceLandmark5, #type:ignore[valid-type] + '5/68' : FaceLandmark5, #type:ignore[valid-type] + '68' : FaceLandmark68, #type:ignore[valid-type] + '68/5' : FaceLandmark68 #type:ignore[valid-type] +}) +FaceScoreSet = TypedDict('FaceScoreSet', +{ + 'detector' : Score, + 'landmarker' : Score +}) +Embedding : TypeAlias = NDArray[numpy.float64] +Gender = Literal['female', 'male'] +Age : TypeAlias = range +Race = Literal['white', 'black', 'latino', 'asian', 'indian', 'arabic'] +Face = namedtuple('Face', +[ + 'bounding_box', + 'score_set', + 'landmark_set', + 'angle', + 'embedding', + 'embedding_norm', + 'gender', + 'age', + 'race' +]) +FaceSet : TypeAlias = Dict[str, List[Face]] +FaceStore = TypedDict('FaceStore', +{ + 'static_faces' : FaceSet +}) + +Language = Literal['en'] +Locales : TypeAlias = Dict[Language, Dict[str, Any]] +LocalePoolSet : TypeAlias = Dict[str, Locales] + +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'] +VisionFrame : TypeAlias = NDArray[Any] +Mask : TypeAlias = NDArray[Any] +Points : TypeAlias = NDArray[Any] +Distance : TypeAlias = NDArray[Any] +Matrix : TypeAlias = NDArray[Any] +Anchors : TypeAlias = NDArray[Any] +Translation : TypeAlias = NDArray[Any] + +AudioBuffer : TypeAlias = bytes +Audio : TypeAlias = NDArray[Any] +AudioChunk : TypeAlias = NDArray[Any] +AudioFrame : TypeAlias = NDArray[Any] +Spectrogram : TypeAlias = NDArray[Any] +Mel : TypeAlias = NDArray[Any] +MelFilterBank : TypeAlias = NDArray[Any] +Voice : TypeAlias = NDArray[Any] +VoiceChunk : TypeAlias = NDArray[Any] + +Fps : TypeAlias = float +Duration : TypeAlias = float +Color : TypeAlias = Tuple[int, int, int, int] +Padding : TypeAlias = Tuple[int, int, int, int] +Margin : TypeAlias = Tuple[int, int, int, int] +Orientation = Literal['landscape', 'portrait'] +Resolution : TypeAlias = Tuple[int, int] + +ProcessState = Literal['checking', 'processing', 'stopping', 'pending'] +Args : TypeAlias = Dict[str, Any] +UpdateProgress : TypeAlias = Callable[[int], None] +ProcessStep : TypeAlias = Callable[[str, int, Args], bool] + +Content : TypeAlias = Dict[str, Any] + +Command : TypeAlias = str +CommandSet : TypeAlias = Dict[str, List[Command]] + +WarpTemplate = Literal['arcface_112_v1', 'arcface_112_v2', 'arcface_128', 'dfl_whole_face', 'ffhq_512', 'mtcnn_512', 'styleganex_384'] +WarpTemplateSet : TypeAlias = Dict[WarpTemplate, NDArray[Any]] +ProcessMode = Literal['output', 'preview', 'stream'] + +ErrorCode = Literal[0, 1, 2, 3, 4] +LogLevel = Literal['error', 'warn', 'info', 'debug'] +LogLevelSet : TypeAlias = Dict[LogLevel, int] + +TableHeader : TypeAlias = str +TableContent : TypeAlias = Any + +FaceDetectorModel = Literal['many', 'retinaface', 'scrfd', 'yolo_face', 'yunet'] +FaceLandmarkerModel = Literal['many', '2dfan4', 'peppa_wutz'] +FaceDetectorSet : TypeAlias = Dict[FaceDetectorModel, List[str]] +FaceSelectorMode = Literal['many', 'one', 'reference'] +FaceSelectorOrder = Literal['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best'] +FaceOccluderModel = Literal['many', 'xseg_1', 'xseg_2', 'xseg_3'] +FaceParserModel = Literal['bisenet_resnet_18', 'bisenet_resnet_34'] +FaceMaskType = Literal['box', 'occlusion', 'area', 'region'] +FaceMaskArea = Literal['upper-face', 'lower-face', 'mouth'] +FaceMaskRegion = Literal['skin', 'left-eyebrow', 'right-eyebrow', 'left-eye', 'right-eye', 'glasses', 'nose', 'mouth', 'upper-lip', 'lower-lip'] +FaceMaskRegionSet : TypeAlias = Dict[FaceMaskRegion, int] +FaceMaskAreaSet : TypeAlias = Dict[FaceMaskArea, List[int]] + +VoiceExtractorModel = Literal['kim_vocal_1', 'kim_vocal_2', 'uvr_mdxnet'] + +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'] +AudioTypeSet : TypeAlias = Dict[AudioFormat, str] +ImageTypeSet : TypeAlias = Dict[ImageFormat, str] +VideoTypeSet : TypeAlias = Dict[VideoFormat, str] + +AudioEncoder = Literal['flac', 'aac', 'libmp3lame', 'libopus', 'libvorbis', 'pcm_s16le', 'pcm_s32le'] +VideoEncoder = Literal['libx264', 'libx264rgb', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc', 'h264_amf', 'hevc_amf', 'h264_qsv', 'hevc_qsv', 'h264_videotoolbox', 'hevc_videotoolbox', 'rawvideo'] +EncoderSet = TypedDict('EncoderSet', +{ + 'audio' : List[AudioEncoder], + 'video' : List[VideoEncoder] +}) +VideoPreset = Literal['ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow'] + +BenchmarkMode = Literal['warm', 'cold'] +BenchmarkResolution = Literal['240p', '360p', '540p', '720p', '1080p', '1440p', '2160p'] +BenchmarkSet : TypeAlias = Dict[BenchmarkResolution, str] +BenchmarkCycleSet = TypedDict('BenchmarkCycleSet', +{ + 'target_path' : str, + 'cycle_count' : int, + 'average_run' : float, + 'fastest_run' : float, + 'slowest_run' : float, + 'relative_fps' : float +}) + +WebcamMode = Literal['inline', 'udp', 'v4l2'] +StreamMode = Literal['udp', 'v4l2'] + +ModelOptions : TypeAlias = Dict[str, Any] +ModelSet : TypeAlias = Dict[str, ModelOptions] +ModelInitializer : TypeAlias = NDArray[Any] + +ExecutionProvider = Literal['cpu', 'coreml', 'cuda', 'directml', 'openvino', 'migraphx', 'rocm', 'tensorrt'] +ExecutionProviderValue = Literal['CPUExecutionProvider', 'CoreMLExecutionProvider', 'CUDAExecutionProvider', 'DmlExecutionProvider', 'OpenVINOExecutionProvider', 'MIGraphXExecutionProvider', 'ROCMExecutionProvider', 'TensorrtExecutionProvider'] +ExecutionProviderSet : TypeAlias = Dict[ExecutionProvider, ExecutionProviderValue] +InferenceSessionProvider : TypeAlias = Any +ValueAndUnit = TypedDict('ValueAndUnit', +{ + 'value' : int, + 'unit' : str +}) +ExecutionDeviceFramework = TypedDict('ExecutionDeviceFramework', +{ + 'name' : str, + 'version' : str +}) +ExecutionDeviceProduct = TypedDict('ExecutionDeviceProduct', +{ + 'vendor' : str, + 'name' : str +}) +ExecutionDeviceVideoMemory = TypedDict('ExecutionDeviceVideoMemory', +{ + 'total' : Optional[ValueAndUnit], + 'free' : Optional[ValueAndUnit] +}) +ExecutionDeviceTemperature = TypedDict('ExecutionDeviceTemperature', +{ + 'gpu' : Optional[ValueAndUnit], + 'memory' : Optional[ValueAndUnit] +}) +ExecutionDeviceUtilization = TypedDict('ExecutionDeviceUtilization', +{ + 'gpu' : Optional[ValueAndUnit], + 'memory' : Optional[ValueAndUnit] +}) +ExecutionDevice = TypedDict('ExecutionDevice', +{ + 'driver_version' : str, + 'framework' : ExecutionDeviceFramework, + 'product' : ExecutionDeviceProduct, + 'video_memory' : ExecutionDeviceVideoMemory, + 'temperature' : ExecutionDeviceTemperature, + 'utilization' : ExecutionDeviceUtilization +}) + +DownloadProvider = Literal['github', 'huggingface'] +DownloadProviderValue = TypedDict('DownloadProviderValue', +{ + 'urls' : List[str], + 'path' : str +}) +DownloadProviderSet : TypeAlias = Dict[DownloadProvider, DownloadProviderValue] +DownloadScope = Literal['lite', 'full'] +Download = TypedDict('Download', +{ + 'url' : str, + 'path' : str +}) +DownloadSet : TypeAlias = Dict[str, Download] + +VideoMemoryStrategy = Literal['strict', 'moderate', 'tolerant'] +AppContext = Literal['cli', 'ui'] + +InferencePool : TypeAlias = Dict[str, InferenceSession] +InferencePoolSet : TypeAlias = Dict[AppContext, Dict[str, InferencePool]] + +UiWorkflow = Literal['instant_runner', 'job_runner', 'job_manager'] + +JobStore = TypedDict('JobStore', +{ + 'job_keys' : List[str], + 'step_keys' : List[str] +}) +JobOutputSet : TypeAlias = Dict[str, List[str]] +JobStatus = Literal['drafted', 'queued', 'completed', 'failed'] +JobStepStatus = Literal['drafted', 'queued', 'started', 'completed', 'failed'] +JobStep = TypedDict('JobStep', +{ + 'args' : Args, + 'status' : JobStepStatus +}) +Job = TypedDict('Job', +{ + 'version' : str, + 'date_created' : str, + 'date_updated' : Optional[str], + 'steps' : List[JobStep] +}) +JobSet : TypeAlias = Dict[str, Job] + +StateKey = Literal\ +[ + 'command', + 'config_path', + 'temp_path', + 'jobs_path', + 'source_paths', + 'target_path', + 'output_path', + 'source_pattern', + 'target_pattern', + 'output_pattern', + 'download_providers', + 'download_scope', + 'benchmark_mode', + 'benchmark_resolutions', + 'benchmark_cycle_count', + 'face_detector_model', + 'face_detector_size', + 'face_detector_margin', + 'face_detector_angles', + 'face_detector_score', + 'face_landmarker_model', + 'face_landmarker_score', + 'face_selector_mode', + 'face_selector_order', + 'face_selector_gender', + 'face_selector_race', + 'face_selector_age_start', + 'face_selector_age_end', + 'reference_face_position', + 'reference_face_distance', + 'reference_frame_number', + 'face_occluder_model', + 'face_parser_model', + 'face_mask_types', + 'face_mask_areas', + 'face_mask_regions', + 'face_mask_blur', + 'face_mask_padding', + 'voice_extractor_model', + 'trim_frame_start', + 'trim_frame_end', + 'temp_frame_format', + 'keep_temp', + 'output_image_quality', + 'output_image_scale', + 'output_audio_encoder', + 'output_audio_quality', + 'output_audio_volume', + 'output_video_encoder', + 'output_video_preset', + 'output_video_quality', + 'output_video_scale', + 'output_video_fps', + 'processors', + 'open_browser', + 'ui_layouts', + 'ui_workflow', + 'execution_device_ids', + 'execution_providers', + 'execution_thread_count', + 'video_memory_strategy', + 'system_memory_limit', + 'log_level', + 'halt_on_error', + 'job_id', + 'job_status', + 'step_index' +] +State = TypedDict('State', +{ + 'command' : str, + 'config_path' : str, + 'temp_path' : str, + 'jobs_path' : str, + 'source_paths' : List[str], + 'target_path' : str, + 'output_path' : str, + 'source_pattern' : str, + 'target_pattern' : str, + 'output_pattern' : str, + 'download_providers' : List[DownloadProvider], + 'download_scope' : DownloadScope, + 'benchmark_mode' : BenchmarkMode, + 'benchmark_resolutions' : List[BenchmarkResolution], + 'benchmark_cycle_count' : int, + 'face_detector_model' : FaceDetectorModel, + 'face_detector_size' : str, + 'face_detector_margin': Margin, + 'face_detector_angles' : List[Angle], + 'face_detector_score' : Score, + 'face_landmarker_model' : FaceLandmarkerModel, + 'face_landmarker_score' : Score, + 'face_selector_mode' : FaceSelectorMode, + 'face_selector_order' : FaceSelectorOrder, + 'face_selector_race' : Race, + 'face_selector_gender' : Gender, + 'face_selector_age_start' : int, + 'face_selector_age_end' : int, + 'reference_face_position' : int, + 'reference_face_distance' : float, + 'reference_frame_number' : int, + 'face_occluder_model' : FaceOccluderModel, + 'face_parser_model' : FaceParserModel, + 'face_mask_types' : List[FaceMaskType], + 'face_mask_areas' : List[FaceMaskArea], + 'face_mask_regions' : List[FaceMaskRegion], + 'face_mask_blur' : float, + 'face_mask_padding' : Padding, + 'voice_extractor_model': VoiceExtractorModel, + 'trim_frame_start' : int, + 'trim_frame_end' : int, + 'temp_frame_format' : TempFrameFormat, + 'keep_temp' : bool, + 'output_image_quality' : int, + 'output_image_scale' : Scale, + 'output_audio_encoder' : AudioEncoder, + 'output_audio_quality' : int, + 'output_audio_volume' : int, + 'output_video_encoder' : VideoEncoder, + 'output_video_preset' : VideoPreset, + 'output_video_quality' : int, + 'output_video_scale' : Scale, + 'output_video_fps' : float, + 'processors' : List[str], + 'open_browser' : bool, + 'ui_layouts' : List[str], + 'ui_workflow' : UiWorkflow, + 'execution_device_ids' : List[int], + 'execution_providers' : List[ExecutionProvider], + 'execution_thread_count' : int, + 'video_memory_strategy' : VideoMemoryStrategy, + 'system_memory_limit' : int, + 'log_level' : LogLevel, + 'halt_on_error' : bool, + 'job_id' : str, + 'job_status' : JobStatus, + 'step_index' : int +}) +ApplyStateItem : TypeAlias = Callable[[Any, Any], None] +StateSet : TypeAlias = Dict[AppContext, State] + diff --git a/facefusion/uis/__init__.py b/facefusion/uis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/uis/assets/overrides.css b/facefusion/uis/assets/overrides.css new file mode 100644 index 0000000000000000000000000000000000000000..acd344257987bec471868f953cf382d8e09f4fea --- /dev/null +++ b/facefusion/uis/assets/overrides.css @@ -0,0 +1,163 @@ +:root:root:root:root .gradio-container +{ + overflow: unset; +} + +:root:root:root:root main +{ + max-width: 110em; +} + +:root:root:root:root .tab-like-container input[type="number"] +{ + border-radius: unset; + text-align: center; + order: 1; + padding: unset +} + +:root:root:root:root input[type="number"] +{ + appearance: textfield; +} + +:root:root:root:root input[type="number"]::-webkit-inner-spin-button +{ + appearance: none; +} + +:root:root:root:root input[type="number"]:focus +{ + outline: unset; +} + +:root:root:root:root .reset-button +{ + background: var(--background-fill-secondary); + border: unset; + font-size: unset; + padding: unset; +} + +:root:root:root:root [type="checkbox"], +:root:root:root:root [type="radio"] +{ + border-radius: 50%; + height: 1.125rem; + width: 1.125rem; +} + +:root:root:root:root input[type="range"] +{ + background: transparent; +} + +:root:root:root:root input[type="range"]::-moz-range-thumb, +:root:root:root:root input[type="range"]::-webkit-slider-thumb +{ + background: var(--neutral-300); + box-shadow: unset; + border-radius: 50%; + height: 1.125rem; + width: 1.125rem; +} + +:root:root:root:root .thumbnail-item +{ + border: unset; + box-shadow: unset; +} + +:root:root:root:root .grid-wrap.fixed-height +{ + min-height: unset; +} + +:root:root:root:root .box-face-selector .empty, +:root:root:root:root .box-face-selector .gallery-container +{ + min-height: 7.375rem; +} + +:root:root:root:root .tab-wrapper +{ + padding: 0 0.625rem; +} + +:root:root:root:root .tab-container +{ + gap: 0.5em; +} + +:root:root:root:root .tab-container button +{ + background: unset; + border-bottom: 0.125rem solid; +} + +:root:root:root:root .tab-container button.selected +{ + color: var(--primary-500) +} + +:root:root:root:root .toast-body +{ + background: white; + color: var(--primary-500); + border: unset; + border-radius: unset; +} + +:root:root:root:root .dark .toast-body +{ + background: var(--neutral-900); + color: var(--primary-600); +} + +:root:root:root:root .toast-icon, +:root:root:root:root .toast-title, +:root:root:root:root .toast-text, +:root:root:root:root .toast-close +{ + color: unset; +} + +:root:root:root:root .toast-body .timer +{ + background: currentColor; +} + +:root:root:root:root .slider_input_container > span, +:root:root:root:root .feather-upload, +:root:root:root:root footer +{ + display: none; +} + +:root:root:root:root .image-frame +{ + background-image: conic-gradient(#fff 90deg, #999 90deg 180deg, #fff 180deg 270deg, #999 270deg); + background-size: 1.25rem 1.25rem; + background-repeat: repeat; + width: 100%; +} + +:root:root:root:root .image-frame > img +{ + object-fit: cover; +} + +:root:root:root:root .image-preview.is-landscape +{ + position: sticky; + top: 0; + z-index: 100; +} + +:root:root:root:root .block .error +{ + border: 0.125rem solid; + padding: 0.375rem 0.75rem; + font-size: 0.75rem; + text-transform: uppercase; +} diff --git a/facefusion/uis/choices.py b/facefusion/uis/choices.py new file mode 100644 index 0000000000000000000000000000000000000000..7332abeddf4e3a52051c08507698eb95e3ff621a --- /dev/null +++ b/facefusion/uis/choices.py @@ -0,0 +1,25 @@ +from typing import Dict, List + +from facefusion.types import Color, WebcamMode +from facefusion.uis.types import JobManagerAction, JobRunnerAction, PreviewMode + +job_manager_actions : List[JobManagerAction] = [ 'job-create', 'job-submit', 'job-delete', 'job-add-step', 'job-remix-step', 'job-insert-step', 'job-remove-step' ] +job_runner_actions : List[JobRunnerAction] = [ 'job-run', 'job-run-all', 'job-retry', 'job-retry-all' ] + +common_options : List[str] = [ 'keep-temp' ] + +preview_modes : List[PreviewMode] = [ 'default', 'frame-by-frame', 'face-by-face' ] +preview_resolutions : List[str] = [ '512x512', '768x768', '1024x1024' ] + +webcam_modes : List[WebcamMode] = [ 'inline', 'udp', 'v4l2' ] +webcam_resolutions : List[str] = [ '320x240', '640x480', '800x600', '1024x768', '1280x720', '1280x960', '1920x1080' ] + +background_remover_colors : Dict[str, Color] =\ +{ + 'red' : (255, 0, 0, 255), + 'green' : (0, 255, 0, 255), + 'blue' : (0, 0, 255, 255), + 'black' : (0, 0, 0, 255), + 'white' : (255, 255, 255, 255), + 'alpha' : (0, 0, 0, 0) +} diff --git a/facefusion/uis/components/__init__.py b/facefusion/uis/components/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/uis/components/about.py b/facefusion/uis/components/about.py new file mode 100644 index 0000000000000000000000000000000000000000..3343cb301677b90992c488a27377898b8dcb08cc --- /dev/null +++ b/facefusion/uis/components/about.py @@ -0,0 +1,41 @@ +import random +from typing import Optional + +import gradio + +from facefusion import metadata, translator + +METADATA_BUTTON : Optional[gradio.Button] = None +ACTION_BUTTON : Optional[gradio.Button] = None + + +def render() -> None: + global METADATA_BUTTON + global ACTION_BUTTON + + action = random.choice( + [ + { + 'translator': translator.get('about.fund'), + 'url': 'https://fund.facefusion.io' + }, + { + 'translator': translator.get('about.subscribe'), + 'url': 'https://subscribe.facefusion.io' + }, + { + 'translator': translator.get('about.join'), + 'url': 'https://join.facefusion.io' + } + ]) + + METADATA_BUTTON = gradio.Button( + value = metadata.get('name') + ' ' + metadata.get('version'), + variant = 'primary', + link = metadata.get('url') + ) + ACTION_BUTTON = gradio.Button( + value = action.get('translator'), + link = action.get('url'), + size = 'sm' + ) diff --git a/facefusion/uis/components/age_modifier_options.py b/facefusion/uis/components/age_modifier_options.py new file mode 100644 index 0000000000000000000000000000000000000000..b3a76b1eb89e4fe5805f7056aaa7e6466e34f155 --- /dev/null +++ b/facefusion/uis/components/age_modifier_options.py @@ -0,0 +1,64 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_float_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.age_modifier import choices as age_modifier_choices +from facefusion.processors.modules.age_modifier.types import AgeModifierModel +from facefusion.uis.core import get_ui_component, register_ui_component + +AGE_MODIFIER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +AGE_MODIFIER_DIRECTION_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global AGE_MODIFIER_MODEL_DROPDOWN + global AGE_MODIFIER_DIRECTION_SLIDER + + has_age_modifier = 'age_modifier' in state_manager.get_item('processors') + AGE_MODIFIER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.age_modifier'), + choices = age_modifier_choices.age_modifier_models, + value = state_manager.get_item('age_modifier_model'), + visible = has_age_modifier + ) + AGE_MODIFIER_DIRECTION_SLIDER = gradio.Slider( + label = translator.get('uis.direction_slider', 'facefusion.processors.modules.age_modifier'), + value = state_manager.get_item('age_modifier_direction'), + step = calculate_float_step(age_modifier_choices.age_modifier_direction_range), + minimum = age_modifier_choices.age_modifier_direction_range[0], + maximum = age_modifier_choices.age_modifier_direction_range[-1], + visible = has_age_modifier + ) + register_ui_component('age_modifier_model_dropdown', AGE_MODIFIER_MODEL_DROPDOWN) + register_ui_component('age_modifier_direction_slider', AGE_MODIFIER_DIRECTION_SLIDER) + + +def listen() -> None: + AGE_MODIFIER_MODEL_DROPDOWN.change(update_age_modifier_model, inputs = AGE_MODIFIER_MODEL_DROPDOWN, outputs = AGE_MODIFIER_MODEL_DROPDOWN) + AGE_MODIFIER_DIRECTION_SLIDER.release(update_age_modifier_direction, inputs = AGE_MODIFIER_DIRECTION_SLIDER) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ AGE_MODIFIER_MODEL_DROPDOWN, AGE_MODIFIER_DIRECTION_SLIDER ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider]: + has_age_modifier = 'age_modifier' in processors + return gradio.Dropdown(visible = has_age_modifier), gradio.Slider(visible = has_age_modifier) + + +def update_age_modifier_model(age_modifier_model : AgeModifierModel) -> gradio.Dropdown: + age_modifier_module = load_processor_module('age_modifier') + age_modifier_module.clear_inference_pool() + state_manager.set_item('age_modifier_model', age_modifier_model) + + if age_modifier_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('age_modifier_model')) + return gradio.Dropdown() + + +def update_age_modifier_direction(age_modifier_direction : float) -> None: + state_manager.set_item('age_modifier_direction', int(age_modifier_direction)) diff --git a/facefusion/uis/components/background_remover_options.py b/facefusion/uis/components/background_remover_options.py new file mode 100644 index 0000000000000000000000000000000000000000..0f1cc0f34c4613c62e6ca27f82be527161293211 --- /dev/null +++ b/facefusion/uis/components/background_remover_options.py @@ -0,0 +1,107 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_int_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.background_remover import choices as background_remover_choices +from facefusion.processors.modules.background_remover.types import BackgroundRemoverModel +from facefusion.sanitizer import sanitize_int_range +from facefusion.uis.core import get_ui_component, register_ui_component + +BACKGROUND_REMOVER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +BACKGROUND_REMOVER_COLOR_WRAPPER : Optional[gradio.Group] = None +BACKGROUND_REMOVER_COLOR_RED_NUMBER : Optional[gradio.Number] = None +BACKGROUND_REMOVER_COLOR_GREEN_NUMBER : Optional[gradio.Number] = None +BACKGROUND_REMOVER_COLOR_BLUE_NUMBER : Optional[gradio.Number] = None +BACKGROUND_REMOVER_COLOR_ALPHA_NUMBER : Optional[gradio.Number] = None + + +def render() -> None: + global BACKGROUND_REMOVER_MODEL_DROPDOWN + global BACKGROUND_REMOVER_COLOR_WRAPPER + global BACKGROUND_REMOVER_COLOR_RED_NUMBER + global BACKGROUND_REMOVER_COLOR_GREEN_NUMBER + global BACKGROUND_REMOVER_COLOR_BLUE_NUMBER + global BACKGROUND_REMOVER_COLOR_ALPHA_NUMBER + + has_background_remover = 'background_remover' in state_manager.get_item('processors') + background_remover_color = state_manager.get_item('background_remover_color') + BACKGROUND_REMOVER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.background_remover'), + choices = background_remover_choices.background_remover_models, + value = state_manager.get_item('background_remover_model'), + visible = has_background_remover + ) + with gradio.Group(visible = has_background_remover) as BACKGROUND_REMOVER_COLOR_WRAPPER: + with gradio.Row(): + BACKGROUND_REMOVER_COLOR_RED_NUMBER = gradio.Number( + label = translator.get('uis.color_red_number', 'facefusion.processors.modules.background_remover'), + value = background_remover_color[0], + minimum = background_remover_choices.background_remover_color_range[0], + maximum = background_remover_choices.background_remover_color_range[-1], + step = calculate_int_step(background_remover_choices.background_remover_color_range) + ) + BACKGROUND_REMOVER_COLOR_GREEN_NUMBER = gradio.Number( + label = translator.get('uis.color_green_number', 'facefusion.processors.modules.background_remover'), + value = background_remover_color[1], + minimum = background_remover_choices.background_remover_color_range[0], + maximum = background_remover_choices.background_remover_color_range[-1], + step = calculate_int_step(background_remover_choices.background_remover_color_range) + ) + with gradio.Row(): + BACKGROUND_REMOVER_COLOR_BLUE_NUMBER = gradio.Number( + label = translator.get('uis.color_blue_number', 'facefusion.processors.modules.background_remover'), + value = background_remover_color[2], + minimum = background_remover_choices.background_remover_color_range[0], + maximum = background_remover_choices.background_remover_color_range[-1], + step = calculate_int_step(background_remover_choices.background_remover_color_range) + ) + BACKGROUND_REMOVER_COLOR_ALPHA_NUMBER = gradio.Number( + label = translator.get('uis.color_alpha_number', 'facefusion.processors.modules.background_remover'), + value = background_remover_color[3], + minimum = background_remover_choices.background_remover_color_range[0], + maximum = background_remover_choices.background_remover_color_range[-1], + step = calculate_int_step(background_remover_choices.background_remover_color_range) + ) + register_ui_component('background_remover_model_dropdown', BACKGROUND_REMOVER_MODEL_DROPDOWN) + register_ui_component('background_remover_color_red_number', BACKGROUND_REMOVER_COLOR_RED_NUMBER) + register_ui_component('background_remover_color_green_number', BACKGROUND_REMOVER_COLOR_GREEN_NUMBER) + register_ui_component('background_remover_color_blue_number', BACKGROUND_REMOVER_COLOR_BLUE_NUMBER) + register_ui_component('background_remover_color_alpha_number', BACKGROUND_REMOVER_COLOR_ALPHA_NUMBER) + + +def listen() -> None: + BACKGROUND_REMOVER_MODEL_DROPDOWN.change(update_background_remover_model, inputs = BACKGROUND_REMOVER_MODEL_DROPDOWN, outputs = BACKGROUND_REMOVER_MODEL_DROPDOWN) + background_remover_color_inputs = [ BACKGROUND_REMOVER_COLOR_RED_NUMBER, BACKGROUND_REMOVER_COLOR_GREEN_NUMBER, BACKGROUND_REMOVER_COLOR_BLUE_NUMBER, BACKGROUND_REMOVER_COLOR_ALPHA_NUMBER ] + + for background_remover_color_input in background_remover_color_inputs: + background_remover_color_input.change(update_background_remover_color, inputs = background_remover_color_inputs) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [BACKGROUND_REMOVER_MODEL_DROPDOWN, BACKGROUND_REMOVER_COLOR_WRAPPER]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Group]: + has_background_remover = 'background_remover' in processors + return gradio.Dropdown(visible = has_background_remover), gradio.Group(visible = has_background_remover) + + +def update_background_remover_model(background_remover_model : BackgroundRemoverModel) -> gradio.Dropdown: + background_remover_module = load_processor_module('background_remover') + background_remover_module.clear_inference_pool() + state_manager.set_item('background_remover_model', background_remover_model) + + if background_remover_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('background_remover_model')) + return gradio.Dropdown() + + +def update_background_remover_color(red : int, green : int, blue : int, alpha : int) -> None: + red = sanitize_int_range(red, background_remover_choices.background_remover_color_range) + green = sanitize_int_range(green, background_remover_choices.background_remover_color_range) + blue = sanitize_int_range(blue, background_remover_choices.background_remover_color_range) + alpha = sanitize_int_range(alpha, background_remover_choices.background_remover_color_range) + state_manager.set_item('background_remover_color', (red, green, blue, alpha)) diff --git a/facefusion/uis/components/benchmark.py b/facefusion/uis/components/benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..ce7b66105c4068e04a172d571b2a2d7e34fb7f29 --- /dev/null +++ b/facefusion/uis/components/benchmark.py @@ -0,0 +1,51 @@ +from typing import Any, Iterator, List, Optional + +import gradio + +from facefusion import benchmarker, state_manager, translator + +BENCHMARK_BENCHMARKS_DATAFRAME : Optional[gradio.Dataframe] = None +BENCHMARK_START_BUTTON : Optional[gradio.Button] = None + + +def render() -> None: + global BENCHMARK_BENCHMARKS_DATAFRAME + global BENCHMARK_START_BUTTON + + BENCHMARK_BENCHMARKS_DATAFRAME = gradio.Dataframe( + headers = + [ + 'target_path', + 'cycle_count', + 'average_run', + 'fastest_run', + 'slowest_run', + 'relative_fps' + ], + datatype = + [ + 'str', + 'number', + 'number', + 'number', + 'number', + 'number' + ], + show_label = False + ) + BENCHMARK_START_BUTTON = gradio.Button( + value = translator.get('uis.start_button'), + variant = 'primary', + size = 'sm' + ) + + +def listen() -> None: + BENCHMARK_START_BUTTON.click(start, outputs = BENCHMARK_BENCHMARKS_DATAFRAME) + + +def start() -> Iterator[List[Any]]: + state_manager.sync_state() + + for benchmark in benchmarker.run(): + yield [ list(benchmark_set.values()) for benchmark_set in benchmark ] diff --git a/facefusion/uis/components/benchmark_options.py b/facefusion/uis/components/benchmark_options.py new file mode 100644 index 0000000000000000000000000000000000000000..2135a43b11200261a5ad5528195df85e4ca1809b --- /dev/null +++ b/facefusion/uis/components/benchmark_options.py @@ -0,0 +1,54 @@ +from typing import List, Optional + +import gradio + +import facefusion.choices +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_int_step +from facefusion.types import BenchmarkMode, BenchmarkResolution + +BENCHMARK_MODE_DROPDOWN : Optional[gradio.Dropdown] = None +BENCHMARK_RESOLUTIONS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None +BENCHMARK_CYCLE_COUNT_SLIDER : Optional[gradio.Button] = None + + +def render() -> None: + global BENCHMARK_MODE_DROPDOWN + global BENCHMARK_RESOLUTIONS_CHECKBOX_GROUP + global BENCHMARK_CYCLE_COUNT_SLIDER + + BENCHMARK_MODE_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.benchmark_mode_dropdown'), + choices = facefusion.choices.benchmark_modes, + value = state_manager.get_item('benchmark_mode') + ) + BENCHMARK_RESOLUTIONS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.benchmark_resolutions_checkbox_group'), + choices = facefusion.choices.benchmark_resolutions, + value = state_manager.get_item('benchmark_resolutions') + ) + BENCHMARK_CYCLE_COUNT_SLIDER = gradio.Slider( + label = translator.get('uis.benchmark_cycle_count_slider'), + value = state_manager.get_item('benchmark_cycle_count'), + step = calculate_int_step(facefusion.choices.benchmark_cycle_count_range), + minimum = facefusion.choices.benchmark_cycle_count_range[0], + maximum = facefusion.choices.benchmark_cycle_count_range[-1] + ) + + +def listen() -> None: + BENCHMARK_MODE_DROPDOWN.change(update_benchmark_mode, inputs = BENCHMARK_MODE_DROPDOWN) + BENCHMARK_RESOLUTIONS_CHECKBOX_GROUP.change(update_benchmark_resolutions, inputs = BENCHMARK_RESOLUTIONS_CHECKBOX_GROUP) + BENCHMARK_CYCLE_COUNT_SLIDER.release(update_benchmark_cycle_count, inputs = BENCHMARK_CYCLE_COUNT_SLIDER) + + +def update_benchmark_mode(benchmark_mode : BenchmarkMode) -> None: + state_manager.set_item('benchmark_mode', benchmark_mode) + + +def update_benchmark_resolutions(benchmark_resolutions : List[BenchmarkResolution]) -> None: + state_manager.set_item('benchmark_resolutions', benchmark_resolutions) + + +def update_benchmark_cycle_count(benchmark_cycle_count : int) -> None: + state_manager.set_item('benchmark_cycle_count', benchmark_cycle_count) diff --git a/facefusion/uis/components/common_options.py b/facefusion/uis/components/common_options.py new file mode 100644 index 0000000000000000000000000000000000000000..2deabb154a3f2cd4da3e0e418e13ddef50f07081 --- /dev/null +++ b/facefusion/uis/components/common_options.py @@ -0,0 +1,32 @@ +from typing import List, Optional + +import gradio + +from facefusion import state_manager, translator +from facefusion.uis import choices as uis_choices + +COMMON_OPTIONS_CHECKBOX_GROUP : Optional[gradio.Checkboxgroup] = None + + +def render() -> None: + global COMMON_OPTIONS_CHECKBOX_GROUP + + common_options = [] + + if state_manager.get_item('keep_temp'): + common_options.append('keep-temp') + + COMMON_OPTIONS_CHECKBOX_GROUP = gradio.Checkboxgroup( + label = translator.get('uis.common_options_checkbox_group'), + choices = uis_choices.common_options, + value = common_options + ) + + +def listen() -> None: + COMMON_OPTIONS_CHECKBOX_GROUP.change(update, inputs = COMMON_OPTIONS_CHECKBOX_GROUP) + + +def update(common_options : List[str]) -> None: + keep_temp = 'keep-temp' in common_options + state_manager.set_item('keep_temp', keep_temp) diff --git a/facefusion/uis/components/deep_swapper_options.py b/facefusion/uis/components/deep_swapper_options.py new file mode 100644 index 0000000000000000000000000000000000000000..d2d7144838fa58dcadb22b46601efed7f3220b24 --- /dev/null +++ b/facefusion/uis/components/deep_swapper_options.py @@ -0,0 +1,64 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_int_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.deep_swapper import choices as deep_swapper_choices +from facefusion.processors.modules.deep_swapper.types import DeepSwapperModel +from facefusion.uis.core import get_ui_component, register_ui_component + +DEEP_SWAPPER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +DEEP_SWAPPER_MORPH_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global DEEP_SWAPPER_MODEL_DROPDOWN + global DEEP_SWAPPER_MORPH_SLIDER + + has_deep_swapper = 'deep_swapper' in state_manager.get_item('processors') + DEEP_SWAPPER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.deep_swapper'), + choices = deep_swapper_choices.deep_swapper_models, + value = state_manager.get_item('deep_swapper_model'), + visible = has_deep_swapper + ) + DEEP_SWAPPER_MORPH_SLIDER = gradio.Slider( + label = translator.get('uis.morph_slider', 'facefusion.processors.modules.deep_swapper'), + value = state_manager.get_item('deep_swapper_morph'), + step = calculate_int_step(deep_swapper_choices.deep_swapper_morph_range), + minimum = deep_swapper_choices.deep_swapper_morph_range[0], + maximum = deep_swapper_choices.deep_swapper_morph_range[-1], + visible = has_deep_swapper and load_processor_module('deep_swapper').get_inference_pool() and load_processor_module('deep_swapper').has_morph_input() + ) + register_ui_component('deep_swapper_model_dropdown', DEEP_SWAPPER_MODEL_DROPDOWN) + register_ui_component('deep_swapper_morph_slider', DEEP_SWAPPER_MORPH_SLIDER) + + +def listen() -> None: + DEEP_SWAPPER_MODEL_DROPDOWN.change(update_deep_swapper_model, inputs = DEEP_SWAPPER_MODEL_DROPDOWN, outputs = [ DEEP_SWAPPER_MODEL_DROPDOWN, DEEP_SWAPPER_MORPH_SLIDER ]) + DEEP_SWAPPER_MORPH_SLIDER.release(update_deep_swapper_morph, inputs = DEEP_SWAPPER_MORPH_SLIDER) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ DEEP_SWAPPER_MODEL_DROPDOWN, DEEP_SWAPPER_MORPH_SLIDER ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider]: + has_deep_swapper = 'deep_swapper' in processors + return gradio.Dropdown(visible = has_deep_swapper), gradio.Slider(visible = has_deep_swapper and load_processor_module('deep_swapper').get_inference_pool() and load_processor_module('deep_swapper').has_morph_input()) + + +def update_deep_swapper_model(deep_swapper_model : DeepSwapperModel) -> Tuple[gradio.Dropdown, gradio.Slider]: + deep_swapper_module = load_processor_module('deep_swapper') + deep_swapper_module.clear_inference_pool() + state_manager.set_item('deep_swapper_model', deep_swapper_model) + + if deep_swapper_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('deep_swapper_model')), gradio.Slider(visible = deep_swapper_module.has_morph_input()) + return gradio.Dropdown(), gradio.Slider() + + +def update_deep_swapper_morph(deep_swapper_morph : int) -> None: + state_manager.set_item('deep_swapper_morph', deep_swapper_morph) diff --git a/facefusion/uis/components/download.py b/facefusion/uis/components/download.py new file mode 100644 index 0000000000000000000000000000000000000000..b51625b4a32b4e8d538447072aadf04d4ff5bcf4 --- /dev/null +++ b/facefusion/uis/components/download.py @@ -0,0 +1,48 @@ +from typing import List, Optional + +import gradio + +import facefusion.choices +from facefusion import content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, state_manager, translator, voice_extractor +from facefusion.filesystem import get_file_name, resolve_file_paths +from facefusion.processors.core import get_processors_modules +from facefusion.types import DownloadProvider + +DOWNLOAD_PROVIDERS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None + + +def render() -> None: + global DOWNLOAD_PROVIDERS_CHECKBOX_GROUP + + DOWNLOAD_PROVIDERS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.download_providers_checkbox_group'), + choices = facefusion.choices.download_providers, + value = state_manager.get_item('download_providers') + ) + + +def listen() -> None: + DOWNLOAD_PROVIDERS_CHECKBOX_GROUP.change(update_download_providers, inputs = DOWNLOAD_PROVIDERS_CHECKBOX_GROUP, outputs = DOWNLOAD_PROVIDERS_CHECKBOX_GROUP) + + +def update_download_providers(download_providers : List[DownloadProvider]) -> gradio.CheckboxGroup: + common_modules =\ + [ + content_analyser, + face_classifier, + face_detector, + face_landmarker, + face_recognizer, + face_masker, + voice_extractor + ] + available_processors = [ get_file_name(file_path) for file_path in resolve_file_paths('facefusion/processors/modules') ] + processor_modules = get_processors_modules(available_processors) + + for module in common_modules + processor_modules: + if hasattr(module, 'create_static_model_set'): + module.create_static_model_set.cache_clear() + + download_providers = download_providers or facefusion.choices.download_providers + state_manager.set_item('download_providers', download_providers) + return gradio.CheckboxGroup(value = state_manager.get_item('download_providers')) diff --git a/facefusion/uis/components/execution.py b/facefusion/uis/components/execution.py new file mode 100644 index 0000000000000000000000000000000000000000..cf4c3c1f1e0fc09e342115d09f1f3c3d6b92011f --- /dev/null +++ b/facefusion/uis/components/execution.py @@ -0,0 +1,48 @@ +from typing import List, Optional + +import gradio + +from facefusion import content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, state_manager, translator, voice_extractor +from facefusion.execution import get_available_execution_providers +from facefusion.filesystem import get_file_name, resolve_file_paths +from facefusion.processors.core import get_processors_modules +from facefusion.types import ExecutionProvider + +EXECUTION_PROVIDERS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None + + +def render() -> None: + global EXECUTION_PROVIDERS_CHECKBOX_GROUP + + EXECUTION_PROVIDERS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.execution_providers_checkbox_group'), + choices = get_available_execution_providers(), + value = state_manager.get_item('execution_providers') + ) + + +def listen() -> None: + EXECUTION_PROVIDERS_CHECKBOX_GROUP.change(update_execution_providers, inputs = EXECUTION_PROVIDERS_CHECKBOX_GROUP, outputs = EXECUTION_PROVIDERS_CHECKBOX_GROUP) + + +def update_execution_providers(execution_providers : List[ExecutionProvider]) -> gradio.CheckboxGroup: + common_modules =\ + [ + content_analyser, + face_classifier, + face_detector, + face_landmarker, + face_masker, + face_recognizer, + voice_extractor + ] + available_processors = [ get_file_name(file_path) for file_path in resolve_file_paths('facefusion/processors/modules') ] + processor_modules = get_processors_modules(available_processors) + + for module in common_modules + processor_modules: + if hasattr(module, 'clear_inference_pool'): + module.clear_inference_pool() + + execution_providers = execution_providers or get_available_execution_providers() + state_manager.set_item('execution_providers', execution_providers) + return gradio.CheckboxGroup(value = state_manager.get_item('execution_providers')) diff --git a/facefusion/uis/components/execution_thread_count.py b/facefusion/uis/components/execution_thread_count.py new file mode 100644 index 0000000000000000000000000000000000000000..2a8b145d35096baf320a1418185a34a67deaad05 --- /dev/null +++ b/facefusion/uis/components/execution_thread_count.py @@ -0,0 +1,29 @@ +from typing import Optional + +import gradio + +import facefusion.choices +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_int_step + +EXECUTION_THREAD_COUNT_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global EXECUTION_THREAD_COUNT_SLIDER + + EXECUTION_THREAD_COUNT_SLIDER = gradio.Slider( + label = translator.get('uis.execution_thread_count_slider'), + value = state_manager.get_item('execution_thread_count'), + step = calculate_int_step(facefusion.choices.execution_thread_count_range), + minimum = facefusion.choices.execution_thread_count_range[0], + maximum = facefusion.choices.execution_thread_count_range[-1] + ) + + +def listen() -> None: + EXECUTION_THREAD_COUNT_SLIDER.release(update_execution_thread_count, inputs = EXECUTION_THREAD_COUNT_SLIDER) + + +def update_execution_thread_count(execution_thread_count : float) -> None: + state_manager.set_item('execution_thread_count', int(execution_thread_count)) diff --git a/facefusion/uis/components/expression_restorer_options.py b/facefusion/uis/components/expression_restorer_options.py new file mode 100644 index 0000000000000000000000000000000000000000..d985898a0a000a010d39f48ea8113c68f11fd889 --- /dev/null +++ b/facefusion/uis/components/expression_restorer_options.py @@ -0,0 +1,80 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_float_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.expression_restorer import choices as expression_restorer_choices +from facefusion.processors.modules.expression_restorer.types import ExpressionRestorerArea, ExpressionRestorerModel +from facefusion.uis.core import get_ui_component, register_ui_component + +EXPRESSION_RESTORER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +EXPRESSION_RESTORER_FACTOR_SLIDER : Optional[gradio.Slider] = None +EXPRESSION_RESTORER_AREAS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None + + +def render() -> None: + global EXPRESSION_RESTORER_MODEL_DROPDOWN + global EXPRESSION_RESTORER_FACTOR_SLIDER + global EXPRESSION_RESTORER_AREAS_CHECKBOX_GROUP + + has_expression_restorer = 'expression_restorer' in state_manager.get_item('processors') + EXPRESSION_RESTORER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.expression_restorer'), + choices = expression_restorer_choices.expression_restorer_models, + value = state_manager.get_item('expression_restorer_model'), + visible = has_expression_restorer + ) + EXPRESSION_RESTORER_FACTOR_SLIDER = gradio.Slider( + label = translator.get('uis.factor_slider', 'facefusion.processors.modules.expression_restorer'), + value = state_manager.get_item('expression_restorer_factor'), + step = calculate_float_step(expression_restorer_choices.expression_restorer_factor_range), + minimum = expression_restorer_choices.expression_restorer_factor_range[0], + maximum = expression_restorer_choices.expression_restorer_factor_range[-1], + visible = has_expression_restorer + ) + EXPRESSION_RESTORER_AREAS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.areas_checkbox_group', 'facefusion.processors.modules.expression_restorer'), + choices = expression_restorer_choices.expression_restorer_areas, + value = state_manager.get_item('expression_restorer_areas'), + visible = has_expression_restorer + ) + register_ui_component('expression_restorer_model_dropdown', EXPRESSION_RESTORER_MODEL_DROPDOWN) + register_ui_component('expression_restorer_factor_slider', EXPRESSION_RESTORER_FACTOR_SLIDER) + register_ui_component('expression_restorer_areas_checkbox_group', EXPRESSION_RESTORER_AREAS_CHECKBOX_GROUP) + + +def listen() -> None: + EXPRESSION_RESTORER_MODEL_DROPDOWN.change(update_expression_restorer_model, inputs = EXPRESSION_RESTORER_MODEL_DROPDOWN, outputs = EXPRESSION_RESTORER_MODEL_DROPDOWN) + EXPRESSION_RESTORER_FACTOR_SLIDER.release(update_expression_restorer_factor, inputs = EXPRESSION_RESTORER_FACTOR_SLIDER) + EXPRESSION_RESTORER_AREAS_CHECKBOX_GROUP.change(update_expression_restorer_areas, inputs = EXPRESSION_RESTORER_AREAS_CHECKBOX_GROUP, outputs = EXPRESSION_RESTORER_AREAS_CHECKBOX_GROUP) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ EXPRESSION_RESTORER_MODEL_DROPDOWN, EXPRESSION_RESTORER_FACTOR_SLIDER, EXPRESSION_RESTORER_AREAS_CHECKBOX_GROUP ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider, gradio.CheckboxGroup]: + has_expression_restorer = 'expression_restorer' in processors + return gradio.Dropdown(visible = has_expression_restorer), gradio.Slider(visible = has_expression_restorer), gradio.CheckboxGroup(visible = has_expression_restorer) + + +def update_expression_restorer_model(expression_restorer_model : ExpressionRestorerModel) -> gradio.Dropdown: + expression_restorer_module = load_processor_module('expression_restorer') + expression_restorer_module.clear_inference_pool() + state_manager.set_item('expression_restorer_model', expression_restorer_model) + + if expression_restorer_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('expression_restorer_model')) + return gradio.Dropdown() + + +def update_expression_restorer_factor(expression_restorer_factor : float) -> None: + state_manager.set_item('expression_restorer_factor', int(expression_restorer_factor)) + + +def update_expression_restorer_areas(expression_restorer_areas : List[ExpressionRestorerArea]) -> gradio.CheckboxGroup: + expression_restorer_areas = expression_restorer_areas or expression_restorer_choices.expression_restorer_areas + state_manager.set_item('expression_restorer_areas', expression_restorer_areas) + return gradio.CheckboxGroup(value = state_manager.get_item('expression_restorer_areas')) diff --git a/facefusion/uis/components/face_debugger_options.py b/facefusion/uis/components/face_debugger_options.py new file mode 100644 index 0000000000000000000000000000000000000000..8f9066effb55578c9f2775ed7cde4117d24eb96b --- /dev/null +++ b/facefusion/uis/components/face_debugger_options.py @@ -0,0 +1,40 @@ +from typing import List, Optional + +import gradio + +from facefusion import state_manager, translator +from facefusion.processors.modules.face_debugger import choices as face_debugger_choices +from facefusion.processors.modules.face_debugger.types import FaceDebuggerItem +from facefusion.uis.core import get_ui_component, register_ui_component + +FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None + + +def render() -> None: + global FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP + + has_face_debugger = 'face_debugger' in state_manager.get_item('processors') + FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.items_checkbox_group', 'facefusion.processors.modules.face_debugger'), + choices = face_debugger_choices.face_debugger_items, + value = state_manager.get_item('face_debugger_items'), + visible = has_face_debugger + ) + register_ui_component('face_debugger_items_checkbox_group', FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP) + + +def listen() -> None: + FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP.change(update_face_debugger_items, inputs = FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP) + + +def remote_update(processors : List[str]) -> gradio.CheckboxGroup: + has_face_debugger = 'face_debugger' in processors + return gradio.CheckboxGroup(visible = has_face_debugger) + + +def update_face_debugger_items(face_debugger_items : List[FaceDebuggerItem]) -> None: + state_manager.set_item('face_debugger_items', face_debugger_items) diff --git a/facefusion/uis/components/face_detector.py b/facefusion/uis/components/face_detector.py new file mode 100644 index 0000000000000000000000000000000000000000..a16dbeec7f8aa9b1aedcf7d23ccdb34469fefff5 --- /dev/null +++ b/facefusion/uis/components/face_detector.py @@ -0,0 +1,103 @@ +from typing import Optional, Sequence, Tuple + +import gradio + +import facefusion.choices +from facefusion import face_detector, state_manager, translator +from facefusion.common_helper import calculate_float_step, get_last +from facefusion.sanitizer import sanitize_int_range +from facefusion.types import Angle, FaceDetectorModel, Score +from facefusion.uis.core import register_ui_component +from facefusion.uis.types import ComponentOptions + +FACE_DETECTOR_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_DETECTOR_SIZE_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_DETECTOR_MARGIN_SLIDER : Optional[gradio.Slider] = None +FACE_DETECTOR_ANGLES_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None +FACE_DETECTOR_SCORE_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FACE_DETECTOR_MODEL_DROPDOWN + global FACE_DETECTOR_SIZE_DROPDOWN + global FACE_DETECTOR_MARGIN_SLIDER + global FACE_DETECTOR_ANGLES_CHECKBOX_GROUP + global FACE_DETECTOR_SCORE_SLIDER + + face_detector_size_dropdown_options : ComponentOptions =\ + { + 'label': translator.get('uis.face_detector_size_dropdown'), + 'value': state_manager.get_item('face_detector_size') + } + if state_manager.get_item('face_detector_size') in facefusion.choices.face_detector_set[state_manager.get_item('face_detector_model')]: + face_detector_size_dropdown_options['choices'] = facefusion.choices.face_detector_set[state_manager.get_item('face_detector_model')] + with gradio.Row(): + FACE_DETECTOR_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.face_detector_model_dropdown'), + choices = facefusion.choices.face_detector_models, + value = state_manager.get_item('face_detector_model') + ) + FACE_DETECTOR_SIZE_DROPDOWN = gradio.Dropdown(**face_detector_size_dropdown_options) + FACE_DETECTOR_MARGIN_SLIDER = gradio.Slider( + label = translator.get('uis.face_detector_margin_slider'), + value = state_manager.get_item('face_detector_margin')[0], + step = calculate_float_step(facefusion.choices.face_detector_margin_range), + minimum = facefusion.choices.face_detector_margin_range[0], + maximum = facefusion.choices.face_detector_margin_range[-1] + ) + FACE_DETECTOR_ANGLES_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.face_detector_angles_checkbox_group'), + choices = facefusion.choices.face_detector_angles, + value = state_manager.get_item('face_detector_angles') + ) + FACE_DETECTOR_SCORE_SLIDER = gradio.Slider( + label = translator.get('uis.face_detector_score_slider'), + value = state_manager.get_item('face_detector_score'), + step = calculate_float_step(facefusion.choices.face_detector_score_range), + minimum = facefusion.choices.face_detector_score_range[0], + maximum = facefusion.choices.face_detector_score_range[-1] + ) + register_ui_component('face_detector_model_dropdown', FACE_DETECTOR_MODEL_DROPDOWN) + register_ui_component('face_detector_size_dropdown', FACE_DETECTOR_SIZE_DROPDOWN) + register_ui_component('face_detector_margin_slider', FACE_DETECTOR_MARGIN_SLIDER) + register_ui_component('face_detector_angles_checkbox_group', FACE_DETECTOR_ANGLES_CHECKBOX_GROUP) + register_ui_component('face_detector_score_slider', FACE_DETECTOR_SCORE_SLIDER) + + +def listen() -> None: + FACE_DETECTOR_MODEL_DROPDOWN.change(update_face_detector_model, inputs = FACE_DETECTOR_MODEL_DROPDOWN, outputs = [ FACE_DETECTOR_MODEL_DROPDOWN, FACE_DETECTOR_SIZE_DROPDOWN ]) + FACE_DETECTOR_SIZE_DROPDOWN.change(update_face_detector_size, inputs = FACE_DETECTOR_SIZE_DROPDOWN) + FACE_DETECTOR_MARGIN_SLIDER.release(update_face_detector_margin, inputs=FACE_DETECTOR_MARGIN_SLIDER) + FACE_DETECTOR_ANGLES_CHECKBOX_GROUP.change(update_face_detector_angles, inputs = FACE_DETECTOR_ANGLES_CHECKBOX_GROUP, outputs = FACE_DETECTOR_ANGLES_CHECKBOX_GROUP) + FACE_DETECTOR_SCORE_SLIDER.release(update_face_detector_score, inputs = FACE_DETECTOR_SCORE_SLIDER) + + +def update_face_detector_model(face_detector_model : FaceDetectorModel) -> Tuple[gradio.Dropdown, gradio.Dropdown]: + face_detector.clear_inference_pool() + state_manager.set_item('face_detector_model', face_detector_model) + + if face_detector.pre_check(): + face_detector_size_choices = facefusion.choices.face_detector_set.get(state_manager.get_item('face_detector_model')) + state_manager.set_item('face_detector_size', get_last(face_detector_size_choices)) + return gradio.Dropdown(value = state_manager.get_item('face_detector_model')), gradio.Dropdown(value = state_manager.get_item('face_detector_size'), choices = face_detector_size_choices) + return gradio.Dropdown(), gradio.Dropdown() + + +def update_face_detector_size(face_detector_size : str) -> None: + state_manager.set_item('face_detector_size', face_detector_size) + + +def update_face_detector_margin(face_detector_margin : int) -> None: + face_detector_margin = sanitize_int_range(face_detector_margin, facefusion.choices.face_detector_margin_range) + state_manager.set_item('face_detector_margin', (face_detector_margin, face_detector_margin, face_detector_margin, face_detector_margin)) + + +def update_face_detector_angles(face_detector_angles : Sequence[Angle]) -> gradio.CheckboxGroup: + face_detector_angles = face_detector_angles or facefusion.choices.face_detector_angles + + state_manager.set_item('face_detector_angles', face_detector_angles) + return gradio.CheckboxGroup(value = state_manager.get_item('face_detector_angles')) + + +def update_face_detector_score(face_detector_score : Score) -> None: + state_manager.set_item('face_detector_score', face_detector_score) diff --git a/facefusion/uis/components/face_editor_options.py b/facefusion/uis/components/face_editor_options.py new file mode 100644 index 0000000000000000000000000000000000000000..4e3e437d89df0e14a2f63ce93179937240c87f18 --- /dev/null +++ b/facefusion/uis/components/face_editor_options.py @@ -0,0 +1,272 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_float_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.face_editor import choices as face_editor_choices +from facefusion.processors.modules.face_editor.types import FaceEditorModel +from facefusion.uis.core import get_ui_component, register_ui_component + +FACE_EDITOR_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_EDITOR_EYEBROW_DIRECTION_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_EYE_GAZE_HORIZONTAL_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_EYE_GAZE_VERTICAL_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_EYE_OPEN_RATIO_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_LIP_OPEN_RATIO_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_MOUTH_GRIM_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_MOUTH_POUT_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_MOUTH_PURSE_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_MOUTH_SMILE_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_MOUTH_POSITION_HORIZONTAL_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_MOUTH_POSITION_VERTICAL_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_HEAD_PITCH_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_HEAD_YAW_SLIDER : Optional[gradio.Slider] = None +FACE_EDITOR_HEAD_ROLL_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FACE_EDITOR_MODEL_DROPDOWN + global FACE_EDITOR_EYEBROW_DIRECTION_SLIDER + global FACE_EDITOR_EYE_GAZE_HORIZONTAL_SLIDER + global FACE_EDITOR_EYE_GAZE_VERTICAL_SLIDER + global FACE_EDITOR_EYE_OPEN_RATIO_SLIDER + global FACE_EDITOR_LIP_OPEN_RATIO_SLIDER + global FACE_EDITOR_MOUTH_GRIM_SLIDER + global FACE_EDITOR_MOUTH_POUT_SLIDER + global FACE_EDITOR_MOUTH_PURSE_SLIDER + global FACE_EDITOR_MOUTH_SMILE_SLIDER + global FACE_EDITOR_MOUTH_POSITION_HORIZONTAL_SLIDER + global FACE_EDITOR_MOUTH_POSITION_VERTICAL_SLIDER + global FACE_EDITOR_HEAD_PITCH_SLIDER + global FACE_EDITOR_HEAD_YAW_SLIDER + global FACE_EDITOR_HEAD_ROLL_SLIDER + + has_face_editor = 'face_editor' in state_manager.get_item('processors') + FACE_EDITOR_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.face_editor'), + choices = face_editor_choices.face_editor_models, + value = state_manager.get_item('face_editor_model'), + visible = has_face_editor + ) + FACE_EDITOR_EYEBROW_DIRECTION_SLIDER = gradio.Slider( + label = translator.get('uis.eyebrow_direction_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_eyebrow_direction'), + step = calculate_float_step(face_editor_choices.face_editor_eyebrow_direction_range), + minimum = face_editor_choices.face_editor_eyebrow_direction_range[0], + maximum = face_editor_choices.face_editor_eyebrow_direction_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_EYE_GAZE_HORIZONTAL_SLIDER = gradio.Slider( + label = translator.get('uis.eye_gaze_horizontal_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_eye_gaze_horizontal'), + step = calculate_float_step(face_editor_choices.face_editor_eye_gaze_horizontal_range), + minimum = face_editor_choices.face_editor_eye_gaze_horizontal_range[0], + maximum = face_editor_choices.face_editor_eye_gaze_horizontal_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_EYE_GAZE_VERTICAL_SLIDER = gradio.Slider( + label = translator.get('uis.eye_gaze_vertical_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_eye_gaze_vertical'), + step = calculate_float_step(face_editor_choices.face_editor_eye_gaze_vertical_range), + minimum = face_editor_choices.face_editor_eye_gaze_vertical_range[0], + maximum = face_editor_choices.face_editor_eye_gaze_vertical_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_EYE_OPEN_RATIO_SLIDER = gradio.Slider( + label = translator.get('uis.eye_open_ratio_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_eye_open_ratio'), + step = calculate_float_step(face_editor_choices.face_editor_eye_open_ratio_range), + minimum = face_editor_choices.face_editor_eye_open_ratio_range[0], + maximum = face_editor_choices.face_editor_eye_open_ratio_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_LIP_OPEN_RATIO_SLIDER = gradio.Slider( + label = translator.get('uis.lip_open_ratio_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_lip_open_ratio'), + step = calculate_float_step(face_editor_choices.face_editor_lip_open_ratio_range), + minimum = face_editor_choices.face_editor_lip_open_ratio_range[0], + maximum = face_editor_choices.face_editor_lip_open_ratio_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_MOUTH_GRIM_SLIDER = gradio.Slider( + label = translator.get('uis.mouth_grim_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_mouth_grim'), + step = calculate_float_step(face_editor_choices.face_editor_mouth_grim_range), + minimum = face_editor_choices.face_editor_mouth_grim_range[0], + maximum = face_editor_choices.face_editor_mouth_grim_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_MOUTH_POUT_SLIDER = gradio.Slider( + label = translator.get('uis.mouth_pout_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_mouth_pout'), + step = calculate_float_step(face_editor_choices.face_editor_mouth_pout_range), + minimum = face_editor_choices.face_editor_mouth_pout_range[0], + maximum = face_editor_choices.face_editor_mouth_pout_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_MOUTH_PURSE_SLIDER = gradio.Slider( + label = translator.get('uis.mouth_purse_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_mouth_purse'), + step = calculate_float_step(face_editor_choices.face_editor_mouth_purse_range), + minimum = face_editor_choices.face_editor_mouth_purse_range[0], + maximum = face_editor_choices.face_editor_mouth_purse_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_MOUTH_SMILE_SLIDER = gradio.Slider( + label = translator.get('uis.mouth_smile_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_mouth_smile'), + step = calculate_float_step(face_editor_choices.face_editor_mouth_smile_range), + minimum = face_editor_choices.face_editor_mouth_smile_range[0], + maximum = face_editor_choices.face_editor_mouth_smile_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_MOUTH_POSITION_HORIZONTAL_SLIDER = gradio.Slider( + label = translator.get('uis.mouth_position_horizontal_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_mouth_position_horizontal'), + step = calculate_float_step(face_editor_choices.face_editor_mouth_position_horizontal_range), + minimum = face_editor_choices.face_editor_mouth_position_horizontal_range[0], + maximum = face_editor_choices.face_editor_mouth_position_horizontal_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_MOUTH_POSITION_VERTICAL_SLIDER = gradio.Slider( + label = translator.get('uis.mouth_position_vertical_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_mouth_position_vertical'), + step = calculate_float_step(face_editor_choices.face_editor_mouth_position_vertical_range), + minimum = face_editor_choices.face_editor_mouth_position_vertical_range[0], + maximum = face_editor_choices.face_editor_mouth_position_vertical_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_HEAD_PITCH_SLIDER = gradio.Slider( + label = translator.get('uis.head_pitch_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_head_pitch'), + step = calculate_float_step(face_editor_choices.face_editor_head_pitch_range), + minimum = face_editor_choices.face_editor_head_pitch_range[0], + maximum = face_editor_choices.face_editor_head_pitch_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_HEAD_YAW_SLIDER = gradio.Slider( + label = translator.get('uis.head_yaw_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_head_yaw'), + step = calculate_float_step(face_editor_choices.face_editor_head_yaw_range), + minimum = face_editor_choices.face_editor_head_yaw_range[0], + maximum = face_editor_choices.face_editor_head_yaw_range[-1], + visible = has_face_editor + ) + FACE_EDITOR_HEAD_ROLL_SLIDER = gradio.Slider( + label = translator.get('uis.head_roll_slider', 'facefusion.processors.modules.face_editor'), + value = state_manager.get_item('face_editor_head_roll'), + step = calculate_float_step(face_editor_choices.face_editor_head_roll_range), + minimum = face_editor_choices.face_editor_head_roll_range[0], + maximum = face_editor_choices.face_editor_head_roll_range[-1], + visible = has_face_editor + ) + register_ui_component('face_editor_model_dropdown', FACE_EDITOR_MODEL_DROPDOWN) + register_ui_component('face_editor_eyebrow_direction_slider', FACE_EDITOR_EYEBROW_DIRECTION_SLIDER) + register_ui_component('face_editor_eye_gaze_horizontal_slider', FACE_EDITOR_EYE_GAZE_HORIZONTAL_SLIDER) + register_ui_component('face_editor_eye_gaze_vertical_slider', FACE_EDITOR_EYE_GAZE_VERTICAL_SLIDER) + register_ui_component('face_editor_eye_open_ratio_slider', FACE_EDITOR_EYE_OPEN_RATIO_SLIDER) + register_ui_component('face_editor_lip_open_ratio_slider', FACE_EDITOR_LIP_OPEN_RATIO_SLIDER) + register_ui_component('face_editor_mouth_grim_slider', FACE_EDITOR_MOUTH_GRIM_SLIDER) + register_ui_component('face_editor_mouth_pout_slider', FACE_EDITOR_MOUTH_POUT_SLIDER) + register_ui_component('face_editor_mouth_purse_slider', FACE_EDITOR_MOUTH_PURSE_SLIDER) + register_ui_component('face_editor_mouth_smile_slider', FACE_EDITOR_MOUTH_SMILE_SLIDER) + register_ui_component('face_editor_mouth_position_horizontal_slider', FACE_EDITOR_MOUTH_POSITION_HORIZONTAL_SLIDER) + register_ui_component('face_editor_mouth_position_vertical_slider', FACE_EDITOR_MOUTH_POSITION_VERTICAL_SLIDER) + register_ui_component('face_editor_head_pitch_slider', FACE_EDITOR_HEAD_PITCH_SLIDER) + register_ui_component('face_editor_head_yaw_slider', FACE_EDITOR_HEAD_YAW_SLIDER) + register_ui_component('face_editor_head_roll_slider', FACE_EDITOR_HEAD_ROLL_SLIDER) + + +def listen() -> None: + FACE_EDITOR_MODEL_DROPDOWN.change(update_face_editor_model, inputs = FACE_EDITOR_MODEL_DROPDOWN, outputs = FACE_EDITOR_MODEL_DROPDOWN) + FACE_EDITOR_EYEBROW_DIRECTION_SLIDER.release(update_face_editor_eyebrow_direction, inputs = FACE_EDITOR_EYEBROW_DIRECTION_SLIDER) + FACE_EDITOR_EYE_GAZE_HORIZONTAL_SLIDER.release(update_face_editor_eye_gaze_horizontal, inputs = FACE_EDITOR_EYE_GAZE_HORIZONTAL_SLIDER) + FACE_EDITOR_EYE_GAZE_VERTICAL_SLIDER.release(update_face_editor_eye_gaze_vertical, inputs = FACE_EDITOR_EYE_GAZE_VERTICAL_SLIDER) + FACE_EDITOR_EYE_OPEN_RATIO_SLIDER.release(update_face_editor_eye_open_ratio, inputs = FACE_EDITOR_EYE_OPEN_RATIO_SLIDER) + FACE_EDITOR_LIP_OPEN_RATIO_SLIDER.release(update_face_editor_lip_open_ratio, inputs = FACE_EDITOR_LIP_OPEN_RATIO_SLIDER) + FACE_EDITOR_MOUTH_GRIM_SLIDER.release(update_face_editor_mouth_grim, inputs = FACE_EDITOR_MOUTH_GRIM_SLIDER) + FACE_EDITOR_MOUTH_POUT_SLIDER.release(update_face_editor_mouth_pout, inputs = FACE_EDITOR_MOUTH_POUT_SLIDER) + FACE_EDITOR_MOUTH_PURSE_SLIDER.release(update_face_editor_mouth_purse, inputs = FACE_EDITOR_MOUTH_PURSE_SLIDER) + FACE_EDITOR_MOUTH_SMILE_SLIDER.release(update_face_editor_mouth_smile, inputs = FACE_EDITOR_MOUTH_SMILE_SLIDER) + FACE_EDITOR_MOUTH_POSITION_HORIZONTAL_SLIDER.release(update_face_editor_mouth_position_horizontal, inputs = FACE_EDITOR_MOUTH_POSITION_HORIZONTAL_SLIDER) + FACE_EDITOR_MOUTH_POSITION_VERTICAL_SLIDER.release(update_face_editor_mouth_position_vertical, inputs = FACE_EDITOR_MOUTH_POSITION_VERTICAL_SLIDER) + FACE_EDITOR_HEAD_PITCH_SLIDER.release(update_face_editor_head_pitch, inputs = FACE_EDITOR_HEAD_PITCH_SLIDER) + FACE_EDITOR_HEAD_YAW_SLIDER.release(update_face_editor_head_yaw, inputs = FACE_EDITOR_HEAD_YAW_SLIDER) + FACE_EDITOR_HEAD_ROLL_SLIDER.release(update_face_editor_head_roll, inputs = FACE_EDITOR_HEAD_ROLL_SLIDER) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ FACE_EDITOR_MODEL_DROPDOWN, FACE_EDITOR_EYEBROW_DIRECTION_SLIDER, FACE_EDITOR_EYE_GAZE_HORIZONTAL_SLIDER, FACE_EDITOR_EYE_GAZE_VERTICAL_SLIDER, FACE_EDITOR_EYE_OPEN_RATIO_SLIDER, FACE_EDITOR_LIP_OPEN_RATIO_SLIDER, FACE_EDITOR_MOUTH_GRIM_SLIDER, FACE_EDITOR_MOUTH_POUT_SLIDER, FACE_EDITOR_MOUTH_PURSE_SLIDER, FACE_EDITOR_MOUTH_SMILE_SLIDER, FACE_EDITOR_MOUTH_POSITION_HORIZONTAL_SLIDER, FACE_EDITOR_MOUTH_POSITION_VERTICAL_SLIDER, FACE_EDITOR_HEAD_PITCH_SLIDER, FACE_EDITOR_HEAD_YAW_SLIDER, FACE_EDITOR_HEAD_ROLL_SLIDER ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider]: + has_face_editor = 'face_editor' in processors + return gradio.Dropdown(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor), gradio.Slider(visible = has_face_editor) + + +def update_face_editor_model(face_editor_model : FaceEditorModel) -> gradio.Dropdown: + face_editor_module = load_processor_module('face_editor') + face_editor_module.clear_inference_pool() + state_manager.set_item('face_editor_model', face_editor_model) + + if face_editor_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('face_editor_model')) + return gradio.Dropdown() + + +def update_face_editor_eyebrow_direction(face_editor_eyebrow_direction : float) -> None: + state_manager.set_item('face_editor_eyebrow_direction', face_editor_eyebrow_direction) + + +def update_face_editor_eye_gaze_horizontal(face_editor_eye_gaze_horizontal : float) -> None: + state_manager.set_item('face_editor_eye_gaze_horizontal', face_editor_eye_gaze_horizontal) + + +def update_face_editor_eye_gaze_vertical(face_editor_eye_gaze_vertical : float) -> None: + state_manager.set_item('face_editor_eye_gaze_vertical', face_editor_eye_gaze_vertical) + + +def update_face_editor_eye_open_ratio(face_editor_eye_open_ratio : float) -> None: + state_manager.set_item('face_editor_eye_open_ratio', face_editor_eye_open_ratio) + + +def update_face_editor_lip_open_ratio(face_editor_lip_open_ratio : float) -> None: + state_manager.set_item('face_editor_lip_open_ratio', face_editor_lip_open_ratio) + + +def update_face_editor_mouth_grim(face_editor_mouth_grim : float) -> None: + state_manager.set_item('face_editor_mouth_grim', face_editor_mouth_grim) + + +def update_face_editor_mouth_pout(face_editor_mouth_pout : float) -> None: + state_manager.set_item('face_editor_mouth_pout', face_editor_mouth_pout) + + +def update_face_editor_mouth_purse(face_editor_mouth_purse : float) -> None: + state_manager.set_item('face_editor_mouth_purse', face_editor_mouth_purse) + + +def update_face_editor_mouth_smile(face_editor_mouth_smile : float) -> None: + state_manager.set_item('face_editor_mouth_smile', face_editor_mouth_smile) + + +def update_face_editor_mouth_position_horizontal(face_editor_mouth_position_horizontal : float) -> None: + state_manager.set_item('face_editor_mouth_position_horizontal', face_editor_mouth_position_horizontal) + + +def update_face_editor_mouth_position_vertical(face_editor_mouth_position_vertical : float) -> None: + state_manager.set_item('face_editor_mouth_position_vertical', face_editor_mouth_position_vertical) + + +def update_face_editor_head_pitch(face_editor_head_pitch : float) -> None: + state_manager.set_item('face_editor_head_pitch', face_editor_head_pitch) + + +def update_face_editor_head_yaw(face_editor_head_yaw : float) -> None: + state_manager.set_item('face_editor_head_yaw', face_editor_head_yaw) + + +def update_face_editor_head_roll(face_editor_head_roll : float) -> None: + state_manager.set_item('face_editor_head_roll', face_editor_head_roll) diff --git a/facefusion/uis/components/face_enhancer_options.py b/facefusion/uis/components/face_enhancer_options.py new file mode 100644 index 0000000000000000000000000000000000000000..5ade967bf5c4305bb33718ac7a3b0d6f03227a5d --- /dev/null +++ b/facefusion/uis/components/face_enhancer_options.py @@ -0,0 +1,81 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_float_step, calculate_int_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.face_enhancer import choices as face_enhancer_choices +from facefusion.processors.modules.face_enhancer.types import FaceEnhancerModel, FaceEnhancerWeight +from facefusion.uis.core import get_ui_component, register_ui_component + +FACE_ENHANCER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_ENHANCER_BLEND_SLIDER : Optional[gradio.Slider] = None +FACE_ENHANCER_WEIGHT_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FACE_ENHANCER_MODEL_DROPDOWN + global FACE_ENHANCER_BLEND_SLIDER + global FACE_ENHANCER_WEIGHT_SLIDER + + has_face_enhancer = 'face_enhancer' in state_manager.get_item('processors') + FACE_ENHANCER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.face_enhancer'), + choices = face_enhancer_choices.face_enhancer_models, + value = state_manager.get_item('face_enhancer_model'), + visible = has_face_enhancer + ) + FACE_ENHANCER_BLEND_SLIDER = gradio.Slider( + label = translator.get('uis.blend_slider', 'facefusion.processors.modules.face_enhancer'), + value = state_manager.get_item('face_enhancer_blend'), + step = calculate_int_step(face_enhancer_choices.face_enhancer_blend_range), + minimum = face_enhancer_choices.face_enhancer_blend_range[0], + maximum = face_enhancer_choices.face_enhancer_blend_range[-1], + visible = has_face_enhancer + ) + FACE_ENHANCER_WEIGHT_SLIDER = gradio.Slider( + label = translator.get('uis.weight_slider', 'facefusion.processors.modules.face_enhancer'), + value = state_manager.get_item('face_enhancer_weight'), + step = calculate_float_step(face_enhancer_choices.face_enhancer_weight_range), + minimum = face_enhancer_choices.face_enhancer_weight_range[0], + maximum = face_enhancer_choices.face_enhancer_weight_range[-1], + visible = has_face_enhancer and load_processor_module('face_enhancer').get_inference_pool() and load_processor_module('face_enhancer').has_weight_input() + ) + register_ui_component('face_enhancer_model_dropdown', FACE_ENHANCER_MODEL_DROPDOWN) + register_ui_component('face_enhancer_blend_slider', FACE_ENHANCER_BLEND_SLIDER) + register_ui_component('face_enhancer_weight_slider', FACE_ENHANCER_WEIGHT_SLIDER) + + +def listen() -> None: + FACE_ENHANCER_MODEL_DROPDOWN.change(update_face_enhancer_model, inputs = FACE_ENHANCER_MODEL_DROPDOWN, outputs = [ FACE_ENHANCER_MODEL_DROPDOWN, FACE_ENHANCER_WEIGHT_SLIDER ]) + FACE_ENHANCER_BLEND_SLIDER.release(update_face_enhancer_blend, inputs = FACE_ENHANCER_BLEND_SLIDER) + FACE_ENHANCER_WEIGHT_SLIDER.release(update_face_enhancer_weight, inputs = FACE_ENHANCER_WEIGHT_SLIDER) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ FACE_ENHANCER_MODEL_DROPDOWN, FACE_ENHANCER_BLEND_SLIDER, FACE_ENHANCER_WEIGHT_SLIDER ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider, gradio.Slider]: + has_face_enhancer = 'face_enhancer' in processors + return gradio.Dropdown(visible = has_face_enhancer), gradio.Slider(visible = has_face_enhancer), gradio.Slider(visible = has_face_enhancer and load_processor_module('face_enhancer').get_inference_pool() and load_processor_module('face_enhancer').has_weight_input()) + + +def update_face_enhancer_model(face_enhancer_model : FaceEnhancerModel) -> Tuple[gradio.Dropdown, gradio.Slider]: + face_enhancer_module = load_processor_module('face_enhancer') + face_enhancer_module.clear_inference_pool() + state_manager.set_item('face_enhancer_model', face_enhancer_model) + + if face_enhancer_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('face_enhancer_model')), gradio.Slider(visible = face_enhancer_module.has_weight_input()) + return gradio.Dropdown(), gradio.Slider() + + +def update_face_enhancer_blend(face_enhancer_blend : float) -> None: + state_manager.set_item('face_enhancer_blend', int(face_enhancer_blend)) + + +def update_face_enhancer_weight(face_enhancer_weight : FaceEnhancerWeight) -> None: + state_manager.set_item('face_enhancer_weight', face_enhancer_weight) + diff --git a/facefusion/uis/components/face_landmarker.py b/facefusion/uis/components/face_landmarker.py new file mode 100644 index 0000000000000000000000000000000000000000..fd4d6cac1dbdf7d2551186214eee0a7be1d259db --- /dev/null +++ b/facefusion/uis/components/face_landmarker.py @@ -0,0 +1,50 @@ +from typing import Optional + +import gradio + +import facefusion.choices +from facefusion import face_landmarker, state_manager, translator +from facefusion.common_helper import calculate_float_step +from facefusion.types import FaceLandmarkerModel, Score +from facefusion.uis.core import register_ui_component + +FACE_LANDMARKER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_LANDMARKER_SCORE_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FACE_LANDMARKER_MODEL_DROPDOWN + global FACE_LANDMARKER_SCORE_SLIDER + + FACE_LANDMARKER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.face_landmarker_model_dropdown'), + choices = facefusion.choices.face_landmarker_models, + value = state_manager.get_item('face_landmarker_model') + ) + FACE_LANDMARKER_SCORE_SLIDER = gradio.Slider( + label = translator.get('uis.face_landmarker_score_slider'), + value = state_manager.get_item('face_landmarker_score'), + step = calculate_float_step(facefusion.choices.face_landmarker_score_range), + minimum = facefusion.choices.face_landmarker_score_range[0], + maximum = facefusion.choices.face_landmarker_score_range[-1] + ) + register_ui_component('face_landmarker_model_dropdown', FACE_LANDMARKER_MODEL_DROPDOWN) + register_ui_component('face_landmarker_score_slider', FACE_LANDMARKER_SCORE_SLIDER) + + +def listen() -> None: + FACE_LANDMARKER_MODEL_DROPDOWN.change(update_face_landmarker_model, inputs = FACE_LANDMARKER_MODEL_DROPDOWN, outputs = FACE_LANDMARKER_MODEL_DROPDOWN) + FACE_LANDMARKER_SCORE_SLIDER.release(update_face_landmarker_score, inputs = FACE_LANDMARKER_SCORE_SLIDER) + + +def update_face_landmarker_model(face_landmarker_model : FaceLandmarkerModel) -> gradio.Dropdown: + face_landmarker.clear_inference_pool() + state_manager.set_item('face_landmarker_model', face_landmarker_model) + + if face_landmarker.pre_check(): + gradio.Dropdown(value = state_manager.get_item('face_landmarker_model')) + return gradio.Dropdown() + + +def update_face_landmarker_score(face_landmarker_score : Score) -> None: + state_manager.set_item('face_landmarker_score', face_landmarker_score) diff --git a/facefusion/uis/components/face_masker.py b/facefusion/uis/components/face_masker.py new file mode 100644 index 0000000000000000000000000000000000000000..6a7784d531aa12cf772fcde253a0fbe8cc93f3d4 --- /dev/null +++ b/facefusion/uis/components/face_masker.py @@ -0,0 +1,181 @@ +from typing import List, Optional, Tuple + +import gradio + +import facefusion.choices +from facefusion import face_masker, state_manager, translator +from facefusion.common_helper import calculate_float_step, calculate_int_step +from facefusion.sanitizer import sanitize_int_range +from facefusion.types import FaceMaskArea, FaceMaskRegion, FaceMaskType, FaceOccluderModel, FaceParserModel +from facefusion.uis.core import register_ui_component + +FACE_OCCLUDER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_PARSER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_MASK_TYPES_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None +FACE_MASK_AREAS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None +FACE_MASK_REGIONS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None +FACE_MASK_BOX_WRAPPER : Optional[gradio.Group] = None +FACE_MASK_BLUR_SLIDER : Optional[gradio.Slider] = None +FACE_MASK_PADDING_TOP_SLIDER : Optional[gradio.Slider] = None +FACE_MASK_PADDING_RIGHT_SLIDER : Optional[gradio.Slider] = None +FACE_MASK_PADDING_BOTTOM_SLIDER : Optional[gradio.Slider] = None +FACE_MASK_PADDING_LEFT_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FACE_OCCLUDER_MODEL_DROPDOWN + global FACE_PARSER_MODEL_DROPDOWN + global FACE_MASK_TYPES_CHECKBOX_GROUP + global FACE_MASK_AREAS_CHECKBOX_GROUP + global FACE_MASK_REGIONS_CHECKBOX_GROUP + global FACE_MASK_BOX_WRAPPER + global FACE_MASK_BLUR_SLIDER + global FACE_MASK_PADDING_TOP_SLIDER + global FACE_MASK_PADDING_RIGHT_SLIDER + global FACE_MASK_PADDING_BOTTOM_SLIDER + global FACE_MASK_PADDING_LEFT_SLIDER + + has_box_mask = 'box' in state_manager.get_item('face_mask_types') + has_region_mask = 'region' in state_manager.get_item('face_mask_types') + has_area_mask = 'area' in state_manager.get_item('face_mask_types') + with gradio.Row(): + FACE_OCCLUDER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.face_occluder_model_dropdown'), + choices = facefusion.choices.face_occluder_models, + value = state_manager.get_item('face_occluder_model') + ) + FACE_PARSER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.face_parser_model_dropdown'), + choices = facefusion.choices.face_parser_models, + value = state_manager.get_item('face_parser_model') + ) + FACE_MASK_TYPES_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.face_mask_types_checkbox_group'), + choices = facefusion.choices.face_mask_types, + value = state_manager.get_item('face_mask_types') + ) + FACE_MASK_AREAS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.face_mask_areas_checkbox_group'), + choices = facefusion.choices.face_mask_areas, + value = state_manager.get_item('face_mask_areas'), + visible = has_area_mask + ) + FACE_MASK_REGIONS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.face_mask_regions_checkbox_group'), + choices = facefusion.choices.face_mask_regions, + value = state_manager.get_item('face_mask_regions'), + visible = has_region_mask + ) + FACE_MASK_BLUR_SLIDER = gradio.Slider( + label = translator.get('uis.face_mask_blur_slider'), + step = calculate_float_step(facefusion.choices.face_mask_blur_range), + minimum = facefusion.choices.face_mask_blur_range[0], + maximum = facefusion.choices.face_mask_blur_range[-1], + value = state_manager.get_item('face_mask_blur'), + visible = has_box_mask + ) + with gradio.Group(visible = has_box_mask) as FACE_MASK_BOX_WRAPPER: + with gradio.Row(): + FACE_MASK_PADDING_TOP_SLIDER = gradio.Slider( + label = translator.get('uis.face_mask_padding_top_slider'), + step = calculate_int_step(facefusion.choices.face_mask_padding_range), + minimum = facefusion.choices.face_mask_padding_range[0], + maximum = facefusion.choices.face_mask_padding_range[-1], + value = state_manager.get_item('face_mask_padding')[0] + ) + FACE_MASK_PADDING_RIGHT_SLIDER = gradio.Slider( + label = translator.get('uis.face_mask_padding_right_slider'), + step = calculate_int_step(facefusion.choices.face_mask_padding_range), + minimum = facefusion.choices.face_mask_padding_range[0], + maximum = facefusion.choices.face_mask_padding_range[-1], + value = state_manager.get_item('face_mask_padding')[1] + ) + with gradio.Row(): + FACE_MASK_PADDING_BOTTOM_SLIDER = gradio.Slider( + label = translator.get('uis.face_mask_padding_bottom_slider'), + step = calculate_int_step(facefusion.choices.face_mask_padding_range), + minimum = facefusion.choices.face_mask_padding_range[0], + maximum = facefusion.choices.face_mask_padding_range[-1], + value = state_manager.get_item('face_mask_padding')[2] + ) + FACE_MASK_PADDING_LEFT_SLIDER = gradio.Slider( + label = translator.get('uis.face_mask_padding_left_slider'), + step = calculate_int_step(facefusion.choices.face_mask_padding_range), + minimum = facefusion.choices.face_mask_padding_range[0], + maximum = facefusion.choices.face_mask_padding_range[-1], + value = state_manager.get_item('face_mask_padding')[3] + ) + register_ui_component('face_occluder_model_dropdown', FACE_OCCLUDER_MODEL_DROPDOWN) + register_ui_component('face_parser_model_dropdown', FACE_PARSER_MODEL_DROPDOWN) + register_ui_component('face_mask_types_checkbox_group', FACE_MASK_TYPES_CHECKBOX_GROUP) + register_ui_component('face_mask_areas_checkbox_group', FACE_MASK_AREAS_CHECKBOX_GROUP) + register_ui_component('face_mask_regions_checkbox_group', FACE_MASK_REGIONS_CHECKBOX_GROUP) + register_ui_component('face_mask_blur_slider', FACE_MASK_BLUR_SLIDER) + register_ui_component('face_mask_padding_top_slider', FACE_MASK_PADDING_TOP_SLIDER) + register_ui_component('face_mask_padding_right_slider', FACE_MASK_PADDING_RIGHT_SLIDER) + register_ui_component('face_mask_padding_bottom_slider', FACE_MASK_PADDING_BOTTOM_SLIDER) + register_ui_component('face_mask_padding_left_slider', FACE_MASK_PADDING_LEFT_SLIDER) + + +def listen() -> None: + FACE_OCCLUDER_MODEL_DROPDOWN.change(update_face_occluder_model, inputs = FACE_OCCLUDER_MODEL_DROPDOWN) + FACE_PARSER_MODEL_DROPDOWN.change(update_face_parser_model, inputs = FACE_PARSER_MODEL_DROPDOWN) + FACE_MASK_TYPES_CHECKBOX_GROUP.change(update_face_mask_types, inputs = FACE_MASK_TYPES_CHECKBOX_GROUP, outputs = [ FACE_MASK_TYPES_CHECKBOX_GROUP, FACE_MASK_AREAS_CHECKBOX_GROUP, FACE_MASK_REGIONS_CHECKBOX_GROUP, FACE_MASK_BLUR_SLIDER, FACE_MASK_BOX_WRAPPER ]) + FACE_MASK_AREAS_CHECKBOX_GROUP.change(update_face_mask_areas, inputs = FACE_MASK_AREAS_CHECKBOX_GROUP, outputs = FACE_MASK_AREAS_CHECKBOX_GROUP) + FACE_MASK_REGIONS_CHECKBOX_GROUP.change(update_face_mask_regions, inputs = FACE_MASK_REGIONS_CHECKBOX_GROUP, outputs = FACE_MASK_REGIONS_CHECKBOX_GROUP) + FACE_MASK_BLUR_SLIDER.release(update_face_mask_blur, inputs = FACE_MASK_BLUR_SLIDER) + + face_mask_padding_sliders = [ FACE_MASK_PADDING_TOP_SLIDER, FACE_MASK_PADDING_RIGHT_SLIDER, FACE_MASK_PADDING_BOTTOM_SLIDER, FACE_MASK_PADDING_LEFT_SLIDER ] + for face_mask_padding_slider in face_mask_padding_sliders: + face_mask_padding_slider.release(update_face_mask_padding, inputs = face_mask_padding_sliders) + + +def update_face_occluder_model(face_occluder_model : FaceOccluderModel) -> gradio.Dropdown: + face_masker.clear_inference_pool() + state_manager.set_item('face_occluder_model', face_occluder_model) + + if face_masker.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('face_occluder_model')) + return gradio.Dropdown() + + +def update_face_parser_model(face_parser_model : FaceParserModel) -> gradio.Dropdown: + face_masker.clear_inference_pool() + state_manager.set_item('face_parser_model', face_parser_model) + + if face_masker.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('face_parser_model')) + return gradio.Dropdown() + + +def update_face_mask_types(face_mask_types : List[FaceMaskType]) -> Tuple[gradio.CheckboxGroup, gradio.CheckboxGroup, gradio.CheckboxGroup, gradio.Slider, gradio.Group]: + face_mask_types = face_mask_types or facefusion.choices.face_mask_types + state_manager.set_item('face_mask_types', face_mask_types) + has_box_mask = 'box' in face_mask_types + has_area_mask = 'area' in face_mask_types + has_region_mask = 'region' in face_mask_types + return gradio.CheckboxGroup(value = state_manager.get_item('face_mask_types')), gradio.CheckboxGroup(visible = has_area_mask), gradio.CheckboxGroup(visible = has_region_mask), gradio.Slider(visible = has_box_mask), gradio.Group(visible = has_box_mask) + + +def update_face_mask_areas(face_mask_areas : List[FaceMaskArea]) -> gradio.CheckboxGroup: + face_mask_areas = face_mask_areas or facefusion.choices.face_mask_areas + state_manager.set_item('face_mask_areas', face_mask_areas) + return gradio.CheckboxGroup(value = state_manager.get_item('face_mask_areas')) + + +def update_face_mask_regions(face_mask_regions : List[FaceMaskRegion]) -> gradio.CheckboxGroup: + face_mask_regions = face_mask_regions or facefusion.choices.face_mask_regions + state_manager.set_item('face_mask_regions', face_mask_regions) + return gradio.CheckboxGroup(value = state_manager.get_item('face_mask_regions')) + + +def update_face_mask_blur(face_mask_blur : float) -> None: + state_manager.set_item('face_mask_blur', face_mask_blur) + + +def update_face_mask_padding(face_mask_padding_top : float, face_mask_padding_right : float, face_mask_padding_bottom : float, face_mask_padding_left : float) -> None: + face_mask_padding_top = sanitize_int_range(int(face_mask_padding_top), facefusion.choices.face_mask_padding_range) + face_mask_padding_right = sanitize_int_range(int(face_mask_padding_right), facefusion.choices.face_mask_padding_range) + face_mask_padding_bottom = sanitize_int_range(int(face_mask_padding_bottom), facefusion.choices.face_mask_padding_range) + face_mask_padding_left = sanitize_int_range(int(face_mask_padding_left), facefusion.choices.face_mask_padding_range) + state_manager.set_item('face_mask_padding', (face_mask_padding_top, face_mask_padding_right, face_mask_padding_bottom, face_mask_padding_left)) diff --git a/facefusion/uis/components/face_selector.py b/facefusion/uis/components/face_selector.py new file mode 100644 index 0000000000000000000000000000000000000000..f7d6b4c9c9046bd555ab66536806889a288f09a9 --- /dev/null +++ b/facefusion/uis/components/face_selector.py @@ -0,0 +1,231 @@ +from typing import List, Optional, Tuple + +import cv2 +import gradio +from gradio_rangeslider import RangeSlider + +import facefusion.choices +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_float_step, calculate_int_step +from facefusion.face_analyser import get_many_faces +from facefusion.face_selector import sort_and_filter_faces +from facefusion.face_store import clear_static_faces +from facefusion.filesystem import is_image, is_video +from facefusion.types import FaceSelectorMode, FaceSelectorOrder, Gender, Race, VisionFrame +from facefusion.uis.core import get_ui_component, get_ui_components, register_ui_component +from facefusion.uis.types import ComponentOptions +from facefusion.uis.ui_helper import convert_str_none +from facefusion.vision import fit_cover_frame, read_static_image, read_video_frame + +FACE_SELECTOR_MODE_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_SELECTOR_ORDER_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_SELECTOR_GENDER_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_SELECTOR_RACE_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_SELECTOR_AGE_RANGE_SLIDER : Optional[RangeSlider] = None +REFERENCE_FACE_POSITION_GALLERY : Optional[gradio.Gallery] = None +REFERENCE_FACE_DISTANCE_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FACE_SELECTOR_MODE_DROPDOWN + global FACE_SELECTOR_ORDER_DROPDOWN + global FACE_SELECTOR_GENDER_DROPDOWN + global FACE_SELECTOR_RACE_DROPDOWN + global FACE_SELECTOR_AGE_RANGE_SLIDER + global REFERENCE_FACE_POSITION_GALLERY + global REFERENCE_FACE_DISTANCE_SLIDER + + reference_face_gallery_options : ComponentOptions =\ + { + 'label': translator.get('uis.reference_face_gallery'), + 'object_fit': 'cover', + 'columns': 7, + 'allow_preview': False, + 'elem_classes': 'box-face-selector', + 'visible': 'reference' in state_manager.get_item('face_selector_mode') + } + if is_image(state_manager.get_item('target_path')): + target_vision_frame = read_static_image(state_manager.get_item('target_path')) + reference_face_gallery_options['value'] = extract_gallery_frames(target_vision_frame) + if is_video(state_manager.get_item('target_path')): + target_vision_frame = read_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) + reference_face_gallery_options['value'] = extract_gallery_frames(target_vision_frame) + FACE_SELECTOR_MODE_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.face_selector_mode_dropdown'), + choices = facefusion.choices.face_selector_modes, + value = state_manager.get_item('face_selector_mode') + ) + REFERENCE_FACE_POSITION_GALLERY = gradio.Gallery(**reference_face_gallery_options) + with gradio.Group(): + with gradio.Row(): + FACE_SELECTOR_ORDER_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.face_selector_order_dropdown'), + choices = facefusion.choices.face_selector_orders, + value = state_manager.get_item('face_selector_order') + ) + FACE_SELECTOR_GENDER_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.face_selector_gender_dropdown'), + choices = [ 'none' ] + facefusion.choices.face_selector_genders, + value = state_manager.get_item('face_selector_gender') or 'none' + ) + FACE_SELECTOR_RACE_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.face_selector_race_dropdown'), + choices = [ 'none' ] + facefusion.choices.face_selector_races, + value = state_manager.get_item('face_selector_race') or 'none' + ) + with gradio.Row(): + face_selector_age_start = state_manager.get_item('face_selector_age_start') or facefusion.choices.face_selector_age_range[0] + face_selector_age_end = state_manager.get_item('face_selector_age_end') or facefusion.choices.face_selector_age_range[-1] + FACE_SELECTOR_AGE_RANGE_SLIDER = RangeSlider( + label = translator.get('uis.face_selector_age_range_slider'), + minimum = facefusion.choices.face_selector_age_range[0], + maximum = facefusion.choices.face_selector_age_range[-1], + value = (face_selector_age_start, face_selector_age_end), + step = calculate_int_step(facefusion.choices.face_selector_age_range) + ) + REFERENCE_FACE_DISTANCE_SLIDER = gradio.Slider( + label = translator.get('uis.reference_face_distance_slider'), + value = state_manager.get_item('reference_face_distance'), + step = calculate_float_step(facefusion.choices.reference_face_distance_range), + minimum = facefusion.choices.reference_face_distance_range[0], + maximum = facefusion.choices.reference_face_distance_range[-1], + visible = 'reference' in state_manager.get_item('face_selector_mode') + ) + register_ui_component('face_selector_mode_dropdown', FACE_SELECTOR_MODE_DROPDOWN) + register_ui_component('face_selector_order_dropdown', FACE_SELECTOR_ORDER_DROPDOWN) + register_ui_component('face_selector_gender_dropdown', FACE_SELECTOR_GENDER_DROPDOWN) + register_ui_component('face_selector_race_dropdown', FACE_SELECTOR_RACE_DROPDOWN) + register_ui_component('face_selector_age_range_slider', FACE_SELECTOR_AGE_RANGE_SLIDER) + register_ui_component('reference_face_position_gallery', REFERENCE_FACE_POSITION_GALLERY) + register_ui_component('reference_face_distance_slider', REFERENCE_FACE_DISTANCE_SLIDER) + + +def listen() -> None: + FACE_SELECTOR_MODE_DROPDOWN.change(update_face_selector_mode, inputs = FACE_SELECTOR_MODE_DROPDOWN, outputs = [ REFERENCE_FACE_POSITION_GALLERY, REFERENCE_FACE_DISTANCE_SLIDER ]) + FACE_SELECTOR_ORDER_DROPDOWN.change(update_face_selector_order, inputs = FACE_SELECTOR_ORDER_DROPDOWN, outputs = REFERENCE_FACE_POSITION_GALLERY) + FACE_SELECTOR_GENDER_DROPDOWN.change(update_face_selector_gender, inputs = FACE_SELECTOR_GENDER_DROPDOWN, outputs = REFERENCE_FACE_POSITION_GALLERY) + FACE_SELECTOR_RACE_DROPDOWN.change(update_face_selector_race, inputs = FACE_SELECTOR_RACE_DROPDOWN, outputs = REFERENCE_FACE_POSITION_GALLERY) + FACE_SELECTOR_AGE_RANGE_SLIDER.release(update_face_selector_age_range, inputs = FACE_SELECTOR_AGE_RANGE_SLIDER, outputs = REFERENCE_FACE_POSITION_GALLERY) + REFERENCE_FACE_DISTANCE_SLIDER.release(update_reference_face_distance, inputs = REFERENCE_FACE_DISTANCE_SLIDER) + + preview_frame_slider = get_ui_component('preview_frame_slider') + if preview_frame_slider: + REFERENCE_FACE_POSITION_GALLERY.select(update_reference_frame_number, inputs = preview_frame_slider) + REFERENCE_FACE_POSITION_GALLERY.select(update_reference_face_position) + + for ui_component in get_ui_components( + [ + 'target_image', + 'target_video' + ]): + for method in [ 'change', 'clear' ]: + getattr(ui_component, method)(clear_reference_frame_number) + getattr(ui_component, method)(clear_reference_face_position) + getattr(ui_component, method)(update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY) + + for ui_component in get_ui_components( + [ + 'face_detector_model_dropdown', + 'face_detector_size_dropdown', + 'face_detector_angles_checkbox_group' + ]): + ui_component.change(clear_and_update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY) + + face_detector_score_slider = get_ui_component('face_detector_score_slider') + if face_detector_score_slider: + face_detector_score_slider.release(update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY) + + preview_frame_slider = get_ui_component('preview_frame_slider') + if preview_frame_slider: + for method in [ 'change', 'release' ]: + getattr(preview_frame_slider, method)(update_reference_position_gallery, inputs = preview_frame_slider, outputs = REFERENCE_FACE_POSITION_GALLERY, show_progress = 'hidden') + + +def update_face_selector_mode(face_selector_mode : FaceSelectorMode) -> Tuple[gradio.Gallery, gradio.Slider]: + state_manager.set_item('face_selector_mode', face_selector_mode) + if face_selector_mode == 'many': + return gradio.Gallery(visible = False), gradio.Slider(visible = False) + if face_selector_mode == 'one': + return gradio.Gallery(visible = False), gradio.Slider(visible = False) + if face_selector_mode == 'reference': + return gradio.Gallery(visible = True), gradio.Slider(visible = True) + + +def update_face_selector_order(face_analyser_order : FaceSelectorOrder) -> gradio.Gallery: + state_manager.set_item('face_selector_order', convert_str_none(face_analyser_order)) + return update_reference_position_gallery() + + +def update_face_selector_gender(face_selector_gender : Gender) -> gradio.Gallery: + state_manager.set_item('face_selector_gender', convert_str_none(face_selector_gender)) + return update_reference_position_gallery() + + +def update_face_selector_race(face_selector_race : Race) -> gradio.Gallery: + state_manager.set_item('face_selector_race', convert_str_none(face_selector_race)) + return update_reference_position_gallery() + + +def update_face_selector_age_range(face_selector_age_range : Tuple[float, float]) -> gradio.Gallery: + face_selector_age_start, face_selector_age_end = face_selector_age_range + state_manager.set_item('face_selector_age_start', int(face_selector_age_start)) + state_manager.set_item('face_selector_age_end', int(face_selector_age_end)) + return update_reference_position_gallery() + + +def update_reference_face_position(event : gradio.SelectData) -> None: + state_manager.set_item('reference_face_position', event.index) + + +def clear_reference_face_position() -> None: + state_manager.set_item('reference_face_position', 0) + + +def update_reference_face_distance(reference_face_distance : float) -> None: + state_manager.set_item('reference_face_distance', reference_face_distance) + + +def update_reference_frame_number(reference_frame_number : int = 0) -> None: + state_manager.set_item('reference_frame_number', reference_frame_number) + + +def clear_reference_frame_number() -> None: + state_manager.set_item('reference_frame_number', 0) + + +def clear_and_update_reference_position_gallery() -> gradio.Gallery: + clear_static_faces() + return update_reference_position_gallery() + + +def update_reference_position_gallery(frame_number : int = 0) -> gradio.Gallery: + gallery_vision_frames = [] + if is_image(state_manager.get_item('target_path')): + target_vision_frame = read_static_image(state_manager.get_item('target_path')) + gallery_vision_frames = extract_gallery_frames(target_vision_frame) + if is_video(state_manager.get_item('target_path')): + target_vision_frame = read_video_frame(state_manager.get_item('target_path'), frame_number) + gallery_vision_frames = extract_gallery_frames(target_vision_frame) + if gallery_vision_frames: + return gradio.Gallery(value = gallery_vision_frames) + return gradio.Gallery(value = None) + + +def extract_gallery_frames(target_vision_frame : VisionFrame) -> List[VisionFrame]: + gallery_vision_frames = [] + faces = get_many_faces([ target_vision_frame ]) + faces = sort_and_filter_faces(faces) + + for face in faces: + start_x, start_y, end_x, end_y = map(int, face.bounding_box) + padding_x = int((end_x - start_x) * 0.25) + padding_y = int((end_y - start_y) * 0.25) + start_x = max(0, start_x - padding_x) + start_y = max(0, start_y - padding_y) + end_x = max(0, end_x + padding_x) + end_y = max(0, end_y + padding_y) + crop_vision_frame = target_vision_frame[start_y:end_y, start_x:end_x] + crop_vision_frame = fit_cover_frame(crop_vision_frame, (128, 128)) + crop_vision_frame = cv2.cvtColor(crop_vision_frame, cv2.COLOR_BGR2RGB) + gallery_vision_frames.append(crop_vision_frame) + return gallery_vision_frames diff --git a/facefusion/uis/components/face_swapper_options.py b/facefusion/uis/components/face_swapper_options.py new file mode 100644 index 0000000000000000000000000000000000000000..c0e39e37ef672173158c68ff6ff0e20ae944e142 --- /dev/null +++ b/facefusion/uis/components/face_swapper_options.py @@ -0,0 +1,84 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_float_step, get_first +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.face_swapper import choices as face_swapper_choices +from facefusion.processors.modules.face_swapper.types import FaceSwapperModel, FaceSwapperWeight +from facefusion.uis.core import get_ui_component, register_ui_component + +FACE_SWAPPER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_SWAPPER_PIXEL_BOOST_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_SWAPPER_WEIGHT_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FACE_SWAPPER_MODEL_DROPDOWN + global FACE_SWAPPER_PIXEL_BOOST_DROPDOWN + global FACE_SWAPPER_WEIGHT_SLIDER + + has_face_swapper = 'face_swapper' in state_manager.get_item('processors') + FACE_SWAPPER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.face_swapper'), + choices = face_swapper_choices.face_swapper_models, + value = state_manager.get_item('face_swapper_model'), + visible = has_face_swapper + ) + FACE_SWAPPER_PIXEL_BOOST_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.pixel_boost_dropdown', 'facefusion.processors.modules.face_swapper'), + choices = face_swapper_choices.face_swapper_set.get(state_manager.get_item('face_swapper_model')), + value = state_manager.get_item('face_swapper_pixel_boost'), + visible = has_face_swapper + ) + FACE_SWAPPER_WEIGHT_SLIDER = gradio.Slider( + label = translator.get('uis.weight_slider', 'facefusion.processors.modules.face_swapper'), + value = state_manager.get_item('face_swapper_weight'), + minimum = face_swapper_choices.face_swapper_weight_range[0], + maximum = face_swapper_choices.face_swapper_weight_range[-1], + step = calculate_float_step(face_swapper_choices.face_swapper_weight_range), + visible = has_face_swapper and has_face_swapper_weight() + ) + register_ui_component('face_swapper_model_dropdown', FACE_SWAPPER_MODEL_DROPDOWN) + register_ui_component('face_swapper_pixel_boost_dropdown', FACE_SWAPPER_PIXEL_BOOST_DROPDOWN) + register_ui_component('face_swapper_weight_slider', FACE_SWAPPER_WEIGHT_SLIDER) + + +def listen() -> None: + FACE_SWAPPER_MODEL_DROPDOWN.change(update_face_swapper_model, inputs = FACE_SWAPPER_MODEL_DROPDOWN, outputs = [ FACE_SWAPPER_MODEL_DROPDOWN, FACE_SWAPPER_PIXEL_BOOST_DROPDOWN, FACE_SWAPPER_WEIGHT_SLIDER ]) + FACE_SWAPPER_PIXEL_BOOST_DROPDOWN.change(update_face_swapper_pixel_boost, inputs = FACE_SWAPPER_PIXEL_BOOST_DROPDOWN) + FACE_SWAPPER_WEIGHT_SLIDER.change(update_face_swapper_weight, inputs = FACE_SWAPPER_WEIGHT_SLIDER) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ FACE_SWAPPER_MODEL_DROPDOWN, FACE_SWAPPER_PIXEL_BOOST_DROPDOWN, FACE_SWAPPER_WEIGHT_SLIDER ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Dropdown, gradio.Slider]: + has_face_swapper = 'face_swapper' in processors + return gradio.Dropdown(visible = has_face_swapper), gradio.Dropdown(visible = has_face_swapper), gradio.Slider(visible = has_face_swapper) + + +def update_face_swapper_model(face_swapper_model : FaceSwapperModel) -> Tuple[gradio.Dropdown, gradio.Dropdown, gradio.Slider]: + face_swapper_module = load_processor_module('face_swapper') + face_swapper_module.clear_inference_pool() + state_manager.set_item('face_swapper_model', face_swapper_model) + + if face_swapper_module.pre_check(): + face_swapper_pixel_boost_dropdown_choices = face_swapper_choices.face_swapper_set.get(state_manager.get_item('face_swapper_model')) + state_manager.set_item('face_swapper_pixel_boost', get_first(face_swapper_pixel_boost_dropdown_choices)) + return gradio.Dropdown(value = state_manager.get_item('face_swapper_model')), gradio.Dropdown(value = state_manager.get_item('face_swapper_pixel_boost'), choices = face_swapper_pixel_boost_dropdown_choices), gradio.Slider(visible = has_face_swapper_weight()) + return gradio.Dropdown(), gradio.Dropdown(), gradio.Slider() + + +def update_face_swapper_pixel_boost(face_swapper_pixel_boost : str) -> None: + state_manager.set_item('face_swapper_pixel_boost', face_swapper_pixel_boost) + + +def update_face_swapper_weight(face_swapper_weight : FaceSwapperWeight) -> None: + state_manager.set_item('face_swapper_weight', face_swapper_weight) + + +def has_face_swapper_weight() -> bool: + return state_manager.get_item('face_swapper_model') in [ 'ghost_1_256', 'ghost_2_256', 'ghost_3_256', 'hififace_unofficial_256', 'hyperswap_1a_256', 'hyperswap_1b_256', 'hyperswap_1c_256', 'inswapper_128', 'inswapper_128_fp16', 'simswap_256', 'simswap_unofficial_512' ] diff --git a/facefusion/uis/components/frame_colorizer_options.py b/facefusion/uis/components/frame_colorizer_options.py new file mode 100644 index 0000000000000000000000000000000000000000..e725bdc6e054533adf02297a8544deda16428f03 --- /dev/null +++ b/facefusion/uis/components/frame_colorizer_options.py @@ -0,0 +1,81 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_int_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.frame_colorizer import choices as frame_colorizer_choices +from facefusion.processors.modules.frame_colorizer.types import FrameColorizerModel +from facefusion.uis.core import get_ui_component, register_ui_component + +FRAME_COLORIZER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FRAME_COLORIZER_SIZE_DROPDOWN : Optional[gradio.Dropdown] = None +FRAME_COLORIZER_BLEND_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FRAME_COLORIZER_MODEL_DROPDOWN + global FRAME_COLORIZER_SIZE_DROPDOWN + global FRAME_COLORIZER_BLEND_SLIDER + + has_frame_colorizer = 'frame_colorizer' in state_manager.get_item('processors') + FRAME_COLORIZER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.frame_colorizer'), + choices = frame_colorizer_choices.frame_colorizer_models, + value = state_manager.get_item('frame_colorizer_model'), + visible = has_frame_colorizer + ) + FRAME_COLORIZER_SIZE_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.size_dropdown', 'facefusion.processors.modules.frame_colorizer'), + choices = frame_colorizer_choices.frame_colorizer_sizes, + value = state_manager.get_item('frame_colorizer_size'), + visible = has_frame_colorizer + ) + FRAME_COLORIZER_BLEND_SLIDER = gradio.Slider( + label = translator.get('uis.blend_slider', 'facefusion.processors.modules.frame_colorizer'), + value = state_manager.get_item('frame_colorizer_blend'), + step = calculate_int_step(frame_colorizer_choices.frame_colorizer_blend_range), + minimum = frame_colorizer_choices.frame_colorizer_blend_range[0], + maximum = frame_colorizer_choices.frame_colorizer_blend_range[-1], + visible = has_frame_colorizer + ) + register_ui_component('frame_colorizer_model_dropdown', FRAME_COLORIZER_MODEL_DROPDOWN) + register_ui_component('frame_colorizer_size_dropdown', FRAME_COLORIZER_SIZE_DROPDOWN) + register_ui_component('frame_colorizer_blend_slider', FRAME_COLORIZER_BLEND_SLIDER) + + +def listen() -> None: + FRAME_COLORIZER_MODEL_DROPDOWN.change(update_frame_colorizer_model, inputs = FRAME_COLORIZER_MODEL_DROPDOWN, outputs = FRAME_COLORIZER_MODEL_DROPDOWN) + FRAME_COLORIZER_SIZE_DROPDOWN.change(update_frame_colorizer_size, inputs = FRAME_COLORIZER_SIZE_DROPDOWN) + FRAME_COLORIZER_BLEND_SLIDER.release(update_frame_colorizer_blend, inputs = FRAME_COLORIZER_BLEND_SLIDER) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ FRAME_COLORIZER_MODEL_DROPDOWN, FRAME_COLORIZER_BLEND_SLIDER, FRAME_COLORIZER_SIZE_DROPDOWN ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider, gradio.Dropdown]: + has_frame_colorizer = 'frame_colorizer' in processors + return gradio.Dropdown(visible = has_frame_colorizer), gradio.Slider(visible = has_frame_colorizer), gradio.Dropdown(visible = has_frame_colorizer) + + +def update_frame_colorizer_model(frame_colorizer_model : FrameColorizerModel) -> gradio.Dropdown: + frame_colorizer_module = load_processor_module('frame_colorizer') + frame_colorizer_module.clear_inference_pool() + state_manager.set_item('frame_colorizer_model', frame_colorizer_model) + + if frame_colorizer_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('frame_colorizer_model')) + return gradio.Dropdown() + + +def update_frame_colorizer_size(frame_colorizer_size : str) -> None: + state_manager.set_item('frame_colorizer_size', frame_colorizer_size) + + +def update_frame_colorizer_blend(frame_colorizer_blend : float) -> None: + state_manager.set_item('frame_colorizer_blend', int(frame_colorizer_blend)) + + + diff --git a/facefusion/uis/components/frame_enhancer_options.py b/facefusion/uis/components/frame_enhancer_options.py new file mode 100644 index 0000000000000000000000000000000000000000..fdb28c6c4e3ed1ee9756e9b981c2fd6228e8d2cc --- /dev/null +++ b/facefusion/uis/components/frame_enhancer_options.py @@ -0,0 +1,64 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_int_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.frame_enhancer import choices as frame_enhancer_choices +from facefusion.processors.modules.frame_enhancer.types import FrameEnhancerModel +from facefusion.uis.core import get_ui_component, register_ui_component + +FRAME_ENHANCER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FRAME_ENHANCER_BLEND_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global FRAME_ENHANCER_MODEL_DROPDOWN + global FRAME_ENHANCER_BLEND_SLIDER + + has_frame_enhancer = 'frame_enhancer' in state_manager.get_item('processors') + FRAME_ENHANCER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.frame_enhancer'), + choices = frame_enhancer_choices.frame_enhancer_models, + value = state_manager.get_item('frame_enhancer_model'), + visible = has_frame_enhancer + ) + FRAME_ENHANCER_BLEND_SLIDER = gradio.Slider( + label = translator.get('uis.blend_slider', 'facefusion.processors.modules.frame_enhancer'), + value = state_manager.get_item('frame_enhancer_blend'), + step = calculate_int_step(frame_enhancer_choices.frame_enhancer_blend_range), + minimum = frame_enhancer_choices.frame_enhancer_blend_range[0], + maximum = frame_enhancer_choices.frame_enhancer_blend_range[-1], + visible = has_frame_enhancer + ) + register_ui_component('frame_enhancer_model_dropdown', FRAME_ENHANCER_MODEL_DROPDOWN) + register_ui_component('frame_enhancer_blend_slider', FRAME_ENHANCER_BLEND_SLIDER) + + +def listen() -> None: + FRAME_ENHANCER_MODEL_DROPDOWN.change(update_frame_enhancer_model, inputs = FRAME_ENHANCER_MODEL_DROPDOWN, outputs = FRAME_ENHANCER_MODEL_DROPDOWN) + FRAME_ENHANCER_BLEND_SLIDER.release(update_frame_enhancer_blend, inputs = FRAME_ENHANCER_BLEND_SLIDER) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ FRAME_ENHANCER_MODEL_DROPDOWN, FRAME_ENHANCER_BLEND_SLIDER ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider]: + has_frame_enhancer = 'frame_enhancer' in processors + return gradio.Dropdown(visible = has_frame_enhancer), gradio.Slider(visible = has_frame_enhancer) + + +def update_frame_enhancer_model(frame_enhancer_model : FrameEnhancerModel) -> gradio.Dropdown: + frame_enhancer_module = load_processor_module('frame_enhancer') + frame_enhancer_module.clear_inference_pool() + state_manager.set_item('frame_enhancer_model', frame_enhancer_model) + + if frame_enhancer_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('frame_enhancer_model')) + return gradio.Dropdown() + + +def update_frame_enhancer_blend(frame_enhancer_blend : float) -> None: + state_manager.set_item('frame_enhancer_blend', int(frame_enhancer_blend)) diff --git a/facefusion/uis/components/instant_runner.py b/facefusion/uis/components/instant_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..e4c267a9cc640615ce322513745ae7bdadcf6752 --- /dev/null +++ b/facefusion/uis/components/instant_runner.py @@ -0,0 +1,110 @@ +from time import sleep +from typing import Optional, Tuple + +import gradio + +from facefusion import process_manager, state_manager, translator +from facefusion.args import collect_step_args +from facefusion.core import process_step +from facefusion.filesystem import is_directory, is_image, is_video +from facefusion.jobs import job_helper, job_manager, job_runner, job_store +from facefusion.temp_helper import clear_temp_directory +from facefusion.types import Args, UiWorkflow +from facefusion.uis.core import get_ui_component +from facefusion.uis.ui_helper import suggest_output_path + +INSTANT_RUNNER_WRAPPER : Optional[gradio.Row] = None +INSTANT_RUNNER_START_BUTTON : Optional[gradio.Button] = None +INSTANT_RUNNER_STOP_BUTTON : Optional[gradio.Button] = None +INSTANT_RUNNER_CLEAR_BUTTON : Optional[gradio.Button] = None + + +def render() -> None: + global INSTANT_RUNNER_WRAPPER + global INSTANT_RUNNER_START_BUTTON + global INSTANT_RUNNER_STOP_BUTTON + global INSTANT_RUNNER_CLEAR_BUTTON + + if job_manager.init_jobs(state_manager.get_item('jobs_path')): + is_instant_runner = state_manager.get_item('ui_workflow') == 'instant_runner' + + with gradio.Row(visible = is_instant_runner) as INSTANT_RUNNER_WRAPPER: + INSTANT_RUNNER_START_BUTTON = gradio.Button( + value = translator.get('uis.start_button'), + variant = 'primary', + size = 'sm' + ) + INSTANT_RUNNER_STOP_BUTTON = gradio.Button( + value = translator.get('uis.stop_button'), + variant = 'primary', + size = 'sm', + visible = False + ) + INSTANT_RUNNER_CLEAR_BUTTON = gradio.Button( + value = translator.get('uis.clear_button'), + size = 'sm' + ) + + +def listen() -> None: + output_image = get_ui_component('output_image') + output_video = get_ui_component('output_video') + ui_workflow_dropdown = get_ui_component('ui_workflow_dropdown') + + if output_image and output_video: + INSTANT_RUNNER_START_BUTTON.click(start, outputs = [ INSTANT_RUNNER_START_BUTTON, INSTANT_RUNNER_STOP_BUTTON ]) + INSTANT_RUNNER_START_BUTTON.click(run, outputs = [ INSTANT_RUNNER_START_BUTTON, INSTANT_RUNNER_STOP_BUTTON, output_image, output_video ]) + INSTANT_RUNNER_STOP_BUTTON.click(stop, outputs = [ INSTANT_RUNNER_START_BUTTON, INSTANT_RUNNER_STOP_BUTTON, output_image, output_video ]) + INSTANT_RUNNER_CLEAR_BUTTON.click(clear, outputs = [ output_image, output_video ]) + if ui_workflow_dropdown: + ui_workflow_dropdown.change(remote_update, inputs = ui_workflow_dropdown, outputs = INSTANT_RUNNER_WRAPPER) + + +def remote_update(ui_workflow : UiWorkflow) -> gradio.Row: + is_instant_runner = ui_workflow == 'instant_runner' + + return gradio.Row(visible = is_instant_runner) + + +def start() -> Tuple[gradio.Button, gradio.Button]: + while not process_manager.is_processing(): + sleep(0.5) + return gradio.Button(visible = False), gradio.Button(visible = True) + + +def run() -> Tuple[gradio.Button, gradio.Button, gradio.Image, gradio.Video]: + step_args = collect_step_args() + output_path = step_args.get('output_path') + + if is_directory(step_args.get('output_path')): + step_args['output_path'] = suggest_output_path(step_args.get('output_path'), state_manager.get_item('target_path')) + if job_manager.init_jobs(state_manager.get_item('jobs_path')): + create_and_run_job(step_args) + state_manager.set_item('output_path', output_path) + if is_image(step_args.get('output_path')): + return gradio.Button(visible = True), gradio.Button(visible = False), gradio.Image(value = step_args.get('output_path'), visible = True), gradio.Video(value = None, visible = False) + if is_video(step_args.get('output_path')): + return gradio.Button(visible = True), gradio.Button(visible = False), gradio.Image(value = None, visible = False), gradio.Video(value = step_args.get('output_path'), visible = True) + return gradio.Button(visible = True), gradio.Button(visible = False), gradio.Image(value = None), gradio.Video(value = None) + + +def create_and_run_job(step_args : Args) -> bool: + job_id = job_helper.suggest_job_id('ui') + + for key in job_store.get_job_keys(): + state_manager.sync_item(key) #type:ignore[arg-type] + + return job_manager.create_job(job_id) and job_manager.add_step(job_id, step_args) and job_manager.submit_job(job_id) and job_runner.run_job(job_id, process_step) + + +def stop() -> Tuple[gradio.Button, gradio.Button, gradio.Image, gradio.Video]: + process_manager.stop() + return gradio.Button(visible = True), gradio.Button(visible = False), gradio.Image(value = None), gradio.Video(value = None) + + +def clear() -> Tuple[gradio.Image, gradio.Video]: + while process_manager.is_processing(): + sleep(0.5) + if state_manager.get_item('target_path'): + clear_temp_directory(state_manager.get_item('target_path')) + return gradio.Image(value = None), gradio.Video(value = None) diff --git a/facefusion/uis/components/job_list.py b/facefusion/uis/components/job_list.py new file mode 100644 index 0000000000000000000000000000000000000000..6be3e6c84c8bf62f025902576c6a1cd8940f078e --- /dev/null +++ b/facefusion/uis/components/job_list.py @@ -0,0 +1,50 @@ +from typing import List, Optional + +import gradio + +import facefusion.choices +from facefusion import state_manager, translator +from facefusion.common_helper import get_first +from facefusion.jobs import job_list, job_manager +from facefusion.types import JobStatus +from facefusion.uis.core import get_ui_component + +JOB_LIST_JOBS_DATAFRAME : Optional[gradio.Dataframe] = None +JOB_LIST_REFRESH_BUTTON : Optional[gradio.Button] = None + + +def render() -> None: + global JOB_LIST_JOBS_DATAFRAME + global JOB_LIST_REFRESH_BUTTON + + if job_manager.init_jobs(state_manager.get_item('jobs_path')): + job_status = get_first(facefusion.choices.job_statuses) + job_headers, job_contents = job_list.compose_job_list(job_status) + + JOB_LIST_JOBS_DATAFRAME = gradio.Dataframe( + headers = job_headers, + value = job_contents, + datatype = [ 'str', 'number', 'date', 'date', 'str' ], + show_label = False + ) + JOB_LIST_REFRESH_BUTTON = gradio.Button( + value = translator.get('uis.refresh_button'), + variant = 'primary', + size = 'sm' + ) + + +def listen() -> None: + job_list_job_status_checkbox_group = get_ui_component('job_list_job_status_checkbox_group') + if job_list_job_status_checkbox_group: + job_list_job_status_checkbox_group.change(update_job_dataframe, inputs = job_list_job_status_checkbox_group, outputs = JOB_LIST_JOBS_DATAFRAME) + JOB_LIST_REFRESH_BUTTON.click(update_job_dataframe, inputs = job_list_job_status_checkbox_group, outputs = JOB_LIST_JOBS_DATAFRAME) + + +def update_job_dataframe(job_statuses : List[JobStatus]) -> gradio.Dataframe: + all_job_contents = [] + + for job_status in job_statuses: + _, job_contents = job_list.compose_job_list(job_status) + all_job_contents.extend(job_contents) + return gradio.Dataframe(value = all_job_contents) diff --git a/facefusion/uis/components/job_list_options.py b/facefusion/uis/components/job_list_options.py new file mode 100644 index 0000000000000000000000000000000000000000..9cbd43306a284ba000373fd3d281f2c11b23f12f --- /dev/null +++ b/facefusion/uis/components/job_list_options.py @@ -0,0 +1,35 @@ +from typing import List, Optional + +import gradio + +import facefusion.choices +from facefusion import state_manager, translator +from facefusion.common_helper import get_first +from facefusion.jobs import job_manager +from facefusion.types import JobStatus +from facefusion.uis.core import register_ui_component + +JOB_LIST_JOB_STATUS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None + + +def render() -> None: + global JOB_LIST_JOB_STATUS_CHECKBOX_GROUP + + if job_manager.init_jobs(state_manager.get_item('jobs_path')): + job_status = get_first(facefusion.choices.job_statuses) + + JOB_LIST_JOB_STATUS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.job_list_status_checkbox_group'), + choices = facefusion.choices.job_statuses, + value = job_status + ) + register_ui_component('job_list_job_status_checkbox_group', JOB_LIST_JOB_STATUS_CHECKBOX_GROUP) + + +def listen() -> None: + JOB_LIST_JOB_STATUS_CHECKBOX_GROUP.change(update_job_status_checkbox_group, inputs = JOB_LIST_JOB_STATUS_CHECKBOX_GROUP, outputs = JOB_LIST_JOB_STATUS_CHECKBOX_GROUP) + + +def update_job_status_checkbox_group(job_statuses : List[JobStatus]) -> gradio.CheckboxGroup: + job_statuses = job_statuses or facefusion.choices.job_statuses + return gradio.CheckboxGroup(value = job_statuses) diff --git a/facefusion/uis/components/job_manager.py b/facefusion/uis/components/job_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..5574b2af8b2a192233c13cb5dda12b49c7545556 --- /dev/null +++ b/facefusion/uis/components/job_manager.py @@ -0,0 +1,194 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import logger, state_manager, translator +from facefusion.args import collect_step_args +from facefusion.common_helper import get_first, get_last +from facefusion.filesystem import is_directory +from facefusion.jobs import job_manager +from facefusion.types import UiWorkflow +from facefusion.uis import choices as uis_choices +from facefusion.uis.core import get_ui_component +from facefusion.uis.types import JobManagerAction +from facefusion.uis.ui_helper import convert_int_none, convert_str_none, suggest_output_path + +JOB_MANAGER_WRAPPER : Optional[gradio.Column] = None +JOB_MANAGER_JOB_ACTION_DROPDOWN : Optional[gradio.Dropdown] = None +JOB_MANAGER_JOB_ID_TEXTBOX : Optional[gradio.Textbox] = None +JOB_MANAGER_JOB_ID_DROPDOWN : Optional[gradio.Dropdown] = None +JOB_MANAGER_STEP_INDEX_DROPDOWN : Optional[gradio.Dropdown] = None +JOB_MANAGER_APPLY_BUTTON : Optional[gradio.Button] = None + + +def render() -> None: + global JOB_MANAGER_WRAPPER + global JOB_MANAGER_JOB_ACTION_DROPDOWN + global JOB_MANAGER_JOB_ID_TEXTBOX + global JOB_MANAGER_JOB_ID_DROPDOWN + global JOB_MANAGER_STEP_INDEX_DROPDOWN + global JOB_MANAGER_APPLY_BUTTON + + if job_manager.init_jobs(state_manager.get_item('jobs_path')): + is_job_manager = state_manager.get_item('ui_workflow') == 'job_manager' + drafted_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ] + + with gradio.Column(visible = is_job_manager) as JOB_MANAGER_WRAPPER: + JOB_MANAGER_JOB_ACTION_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.job_manager_job_action_dropdown'), + choices = uis_choices.job_manager_actions, + value = get_first(uis_choices.job_manager_actions) + ) + JOB_MANAGER_JOB_ID_TEXTBOX = gradio.Textbox( + label = translator.get('uis.job_manager_job_id_dropdown'), + max_lines = 1, + interactive = True + ) + JOB_MANAGER_JOB_ID_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.job_manager_job_id_dropdown'), + choices = drafted_job_ids, + value = get_last(drafted_job_ids), + interactive = True, + visible = False + ) + JOB_MANAGER_STEP_INDEX_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.job_manager_step_index_dropdown'), + choices = [ 'none' ], + value = 'none', + interactive = True, + visible = False + ) + JOB_MANAGER_APPLY_BUTTON = gradio.Button( + value = translator.get('uis.apply_button'), + variant = 'primary', + size = 'sm' + ) + + +def listen() -> None: + JOB_MANAGER_JOB_ACTION_DROPDOWN.change(update, inputs = [ JOB_MANAGER_JOB_ACTION_DROPDOWN, JOB_MANAGER_JOB_ID_DROPDOWN ], outputs = [ JOB_MANAGER_JOB_ID_TEXTBOX, JOB_MANAGER_JOB_ID_DROPDOWN, JOB_MANAGER_STEP_INDEX_DROPDOWN ]) + JOB_MANAGER_JOB_ID_DROPDOWN.change(update_step_index, inputs = JOB_MANAGER_JOB_ID_DROPDOWN, outputs = JOB_MANAGER_STEP_INDEX_DROPDOWN) + JOB_MANAGER_APPLY_BUTTON.click(apply, inputs = [ JOB_MANAGER_JOB_ACTION_DROPDOWN, JOB_MANAGER_JOB_ID_TEXTBOX, JOB_MANAGER_JOB_ID_DROPDOWN, JOB_MANAGER_STEP_INDEX_DROPDOWN ], outputs = [ JOB_MANAGER_JOB_ACTION_DROPDOWN, JOB_MANAGER_JOB_ID_TEXTBOX, JOB_MANAGER_JOB_ID_DROPDOWN, JOB_MANAGER_STEP_INDEX_DROPDOWN ]) + + ui_workflow_dropdown = get_ui_component('ui_workflow_dropdown') + if ui_workflow_dropdown: + ui_workflow_dropdown.change(remote_update, inputs = ui_workflow_dropdown, outputs = [ JOB_MANAGER_WRAPPER, JOB_MANAGER_JOB_ACTION_DROPDOWN, JOB_MANAGER_JOB_ID_TEXTBOX, JOB_MANAGER_JOB_ID_DROPDOWN, JOB_MANAGER_STEP_INDEX_DROPDOWN ]) + + +def remote_update(ui_workflow : UiWorkflow) -> Tuple[gradio.Row, gradio.Dropdown, gradio.Textbox, gradio.Dropdown, gradio.Dropdown]: + is_job_manager = ui_workflow == 'job_manager' + return gradio.Row(visible = is_job_manager), gradio.Dropdown(value = get_first(uis_choices.job_manager_actions)), gradio.Textbox(value = None, visible = True), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False) + + +def apply(job_action : JobManagerAction, created_job_id : str, selected_job_id : str, selected_step_index : int) -> Tuple[gradio.Dropdown, gradio.Textbox, gradio.Dropdown, gradio.Dropdown]: + created_job_id = convert_str_none(created_job_id) + selected_job_id = convert_str_none(selected_job_id) + selected_step_index = convert_int_none(selected_step_index) + step_args = collect_step_args() + output_path = step_args.get('output_path') + + if is_directory(step_args.get('output_path')): + step_args['output_path'] = suggest_output_path(step_args.get('output_path'), state_manager.get_item('target_path')) + + if job_action == 'job-create': + if created_job_id and job_manager.create_job(created_job_id): + updated_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ] + + logger.info(translator.get('job_created').format(job_id = created_job_id), __name__) + return gradio.Dropdown(value = 'job-add-step'), gradio.Textbox(visible = False), gradio.Dropdown(value = created_job_id, choices = updated_job_ids, visible = True), gradio.Dropdown() + else: + logger.error(translator.get('job_not_created').format(job_id = created_job_id), __name__) + + if job_action == 'job-submit': + if selected_job_id and job_manager.submit_job(selected_job_id): + updated_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ] + + logger.info(translator.get('job_submitted').format(job_id = selected_job_id), __name__) + return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(value = get_last(updated_job_ids), choices = updated_job_ids, visible = True), gradio.Dropdown() + else: + logger.error(translator.get('job_not_submitted').format(job_id = selected_job_id), __name__) + + if job_action == 'job-delete': + if selected_job_id and job_manager.delete_job(selected_job_id): + updated_job_ids = job_manager.find_job_ids('drafted') + job_manager.find_job_ids('queued') + job_manager.find_job_ids('failed') + job_manager.find_job_ids('completed') or [ 'none' ] + + logger.info(translator.get('job_deleted').format(job_id = selected_job_id), __name__) + return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(value = get_last(updated_job_ids), choices = updated_job_ids, visible = True), gradio.Dropdown() + else: + logger.error(translator.get('job_not_deleted').format(job_id = selected_job_id), __name__) + + if job_action == 'job-add-step': + if selected_job_id and job_manager.add_step(selected_job_id, step_args): + state_manager.set_item('output_path', output_path) + logger.info(translator.get('job_step_added').format(job_id = selected_job_id), __name__) + return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(visible = True), gradio.Dropdown(visible = False) + else: + state_manager.set_item('output_path', output_path) + logger.error(translator.get('job_step_not_added').format(job_id = selected_job_id), __name__) + + if job_action == 'job-remix-step': + if selected_job_id and job_manager.has_step(selected_job_id, selected_step_index) and job_manager.remix_step(selected_job_id, selected_step_index, step_args): + updated_step_choices = get_step_choices(selected_job_id) or [ 'none' ] #type:ignore[list-item] + + state_manager.set_item('output_path', output_path) + logger.info(translator.get('job_remix_step_added').format(job_id = selected_job_id, step_index = selected_step_index), __name__) + return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(visible = True), gradio.Dropdown(value = get_last(updated_step_choices), choices = updated_step_choices, visible = True) + else: + state_manager.set_item('output_path', output_path) + logger.error(translator.get('job_remix_step_not_added').format(job_id = selected_job_id, step_index = selected_step_index), __name__) + + if job_action == 'job-insert-step': + if selected_job_id and job_manager.has_step(selected_job_id, selected_step_index) and job_manager.insert_step(selected_job_id, selected_step_index, step_args): + updated_step_choices = get_step_choices(selected_job_id) or [ 'none' ] #type:ignore[list-item] + + state_manager.set_item('output_path', output_path) + logger.info(translator.get('job_step_inserted').format(job_id = selected_job_id, step_index = selected_step_index), __name__) + return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(visible = True), gradio.Dropdown(value = get_last(updated_step_choices), choices = updated_step_choices, visible = True) + else: + state_manager.set_item('output_path', output_path) + logger.error(translator.get('job_step_not_inserted').format(job_id = selected_job_id, step_index = selected_step_index), __name__) + + if job_action == 'job-remove-step': + if selected_job_id and job_manager.has_step(selected_job_id, selected_step_index) and job_manager.remove_step(selected_job_id, selected_step_index): + updated_step_choices = get_step_choices(selected_job_id) or [ 'none' ] #type:ignore[list-item] + + logger.info(translator.get('job_step_removed').format(job_id = selected_job_id, step_index = selected_step_index), __name__) + return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(visible = True), gradio.Dropdown(value = get_last(updated_step_choices), choices = updated_step_choices, visible = True) + else: + logger.error(translator.get('job_step_not_removed').format(job_id = selected_job_id, step_index = selected_step_index), __name__) + return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(), gradio.Dropdown() + + +def get_step_choices(job_id : str) -> List[int]: + steps = job_manager.get_steps(job_id) + return [ index for index, _ in enumerate(steps) ] + + +def update(job_action : JobManagerAction, selected_job_id : str) -> Tuple[gradio.Textbox, gradio.Dropdown, gradio.Dropdown]: + if job_action == 'job-create': + return gradio.Textbox(value = None, visible = True), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False) + + if job_action == 'job-delete': + updated_job_ids = job_manager.find_job_ids('drafted') + job_manager.find_job_ids('queued') + job_manager.find_job_ids('failed') + job_manager.find_job_ids('completed') or [ 'none' ] + updated_job_id = selected_job_id if selected_job_id in updated_job_ids else get_last(updated_job_ids) + + return gradio.Textbox(visible = False), gradio.Dropdown(value = updated_job_id, choices = updated_job_ids, visible = True), gradio.Dropdown(visible = False) + + if job_action in [ 'job-submit', 'job-add-step' ]: + updated_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ] + updated_job_id = selected_job_id if selected_job_id in updated_job_ids else get_last(updated_job_ids) + + return gradio.Textbox(visible = False), gradio.Dropdown(value = updated_job_id, choices = updated_job_ids, visible = True), gradio.Dropdown(visible = False) + + if job_action in [ 'job-remix-step', 'job-insert-step', 'job-remove-step' ]: + updated_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ] + updated_job_id = selected_job_id if selected_job_id in updated_job_ids else get_last(updated_job_ids) + updated_step_choices = get_step_choices(updated_job_id) or [ 'none' ] #type:ignore[list-item] + + return gradio.Textbox(visible = False), gradio.Dropdown(value = updated_job_id, choices = updated_job_ids, visible = True), gradio.Dropdown(value = get_last(updated_step_choices), choices = updated_step_choices, visible = True) + return gradio.Textbox(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False) + + +def update_step_index(job_id : str) -> gradio.Dropdown: + step_choices = get_step_choices(job_id) or [ 'none' ] #type:ignore[list-item] + return gradio.Dropdown(value = get_last(step_choices), choices = step_choices) diff --git a/facefusion/uis/components/job_runner.py b/facefusion/uis/components/job_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..d47672e50c209df0f5f1bb677219e81e3e40119d --- /dev/null +++ b/facefusion/uis/components/job_runner.py @@ -0,0 +1,142 @@ +from time import sleep +from typing import Optional, Tuple + +import gradio + +from facefusion import logger, process_manager, state_manager, translator +from facefusion.common_helper import get_first, get_last +from facefusion.core import process_step +from facefusion.jobs import job_manager, job_runner, job_store +from facefusion.types import UiWorkflow +from facefusion.uis import choices as uis_choices +from facefusion.uis.core import get_ui_component +from facefusion.uis.types import JobRunnerAction +from facefusion.uis.ui_helper import convert_str_none + +JOB_RUNNER_WRAPPER : Optional[gradio.Column] = None +JOB_RUNNER_JOB_ACTION_DROPDOWN : Optional[gradio.Dropdown] = None +JOB_RUNNER_JOB_ID_DROPDOWN : Optional[gradio.Dropdown] = None +JOB_RUNNER_START_BUTTON : Optional[gradio.Button] = None +JOB_RUNNER_STOP_BUTTON : Optional[gradio.Button] = None + + +def render() -> None: + global JOB_RUNNER_WRAPPER + global JOB_RUNNER_JOB_ACTION_DROPDOWN + global JOB_RUNNER_JOB_ID_DROPDOWN + global JOB_RUNNER_START_BUTTON + global JOB_RUNNER_STOP_BUTTON + + if job_manager.init_jobs(state_manager.get_item('jobs_path')): + is_job_runner = state_manager.get_item('ui_workflow') == 'job_runner' + queued_job_ids = job_manager.find_job_ids('queued') or [ 'none' ] + + with gradio.Column(visible = is_job_runner) as JOB_RUNNER_WRAPPER: + JOB_RUNNER_JOB_ACTION_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.job_runner_job_action_dropdown'), + choices = uis_choices.job_runner_actions, + value = get_first(uis_choices.job_runner_actions) + ) + JOB_RUNNER_JOB_ID_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.job_runner_job_id_dropdown'), + choices = queued_job_ids, + value = get_last(queued_job_ids) + ) + with gradio.Row(): + JOB_RUNNER_START_BUTTON = gradio.Button( + value = translator.get('uis.start_button'), + variant = 'primary', + size = 'sm' + ) + JOB_RUNNER_STOP_BUTTON = gradio.Button( + value = translator.get('uis.stop_button'), + variant = 'primary', + size = 'sm', + visible = False + ) + + +def listen() -> None: + JOB_RUNNER_JOB_ACTION_DROPDOWN.change(update_job_action, inputs = JOB_RUNNER_JOB_ACTION_DROPDOWN, outputs = JOB_RUNNER_JOB_ID_DROPDOWN) + JOB_RUNNER_START_BUTTON.click(start, outputs = [ JOB_RUNNER_START_BUTTON, JOB_RUNNER_STOP_BUTTON ]) + JOB_RUNNER_START_BUTTON.click(run, inputs = [ JOB_RUNNER_JOB_ACTION_DROPDOWN, JOB_RUNNER_JOB_ID_DROPDOWN ], outputs = [ JOB_RUNNER_START_BUTTON, JOB_RUNNER_STOP_BUTTON, JOB_RUNNER_JOB_ID_DROPDOWN ]) + JOB_RUNNER_STOP_BUTTON.click(stop, outputs = [ JOB_RUNNER_START_BUTTON, JOB_RUNNER_STOP_BUTTON ]) + + ui_workflow_dropdown = get_ui_component('ui_workflow_dropdown') + if ui_workflow_dropdown: + ui_workflow_dropdown.change(remote_update, inputs = ui_workflow_dropdown, outputs = [ JOB_RUNNER_WRAPPER, JOB_RUNNER_JOB_ACTION_DROPDOWN, JOB_RUNNER_JOB_ID_DROPDOWN ]) + + +def remote_update(ui_workflow : UiWorkflow) -> Tuple[gradio.Row, gradio.Dropdown, gradio.Dropdown]: + is_job_runner = ui_workflow == 'job_runner' + queued_job_ids = job_manager.find_job_ids('queued') or [ 'none' ] + + return gradio.Row(visible = is_job_runner), gradio.Dropdown(value = get_first(uis_choices.job_runner_actions), choices = uis_choices.job_runner_actions), gradio.Dropdown(value = get_last(queued_job_ids), choices = queued_job_ids) + + +def start() -> Tuple[gradio.Button, gradio.Button]: + while not process_manager.is_processing(): + sleep(0.5) + return gradio.Button(visible = False), gradio.Button(visible = True) + + +def run(job_action : JobRunnerAction, job_id : str) -> Tuple[gradio.Button, gradio.Button, gradio.Dropdown]: + job_id = convert_str_none(job_id) + + for key in job_store.get_job_keys(): + state_manager.sync_item(key) #type:ignore[arg-type] + + if job_action == 'job-run': + logger.info(translator.get('running_job').format(job_id = job_id), __name__) + if job_id and job_runner.run_job(job_id, process_step): + logger.info(translator.get('processing_job_succeeded').format(job_id = job_id), __name__) + else: + logger.info(translator.get('processing_job_failed').format(job_id = job_id), __name__) + updated_job_ids = job_manager.find_job_ids('queued') or [ 'none' ] + + return gradio.Button(visible = True), gradio.Button(visible = False), gradio.Dropdown(value = get_last(updated_job_ids), choices = updated_job_ids) + + if job_action == 'job-run-all': + logger.info(translator.get('running_jobs'), __name__) + halt_on_error = False + if job_runner.run_jobs(process_step, halt_on_error): + logger.info(translator.get('processing_jobs_succeeded'), __name__) + else: + logger.info(translator.get('processing_jobs_failed'), __name__) + + if job_action == 'job-retry': + logger.info(translator.get('retrying_job').format(job_id = job_id), __name__) + if job_id and job_runner.retry_job(job_id, process_step): + logger.info(translator.get('processing_job_succeeded').format(job_id = job_id), __name__) + else: + logger.info(translator.get('processing_job_failed').format(job_id = job_id), __name__) + updated_job_ids = job_manager.find_job_ids('failed') or [ 'none' ] + + return gradio.Button(visible = True), gradio.Button(visible = False), gradio.Dropdown(value = get_last(updated_job_ids), choices = updated_job_ids) + + if job_action == 'job-retry-all': + logger.info(translator.get('retrying_jobs'), __name__) + halt_on_error = False + if job_runner.retry_jobs(process_step, halt_on_error): + logger.info(translator.get('processing_jobs_succeeded'), __name__) + else: + logger.info(translator.get('processing_jobs_failed'), __name__) + return gradio.Button(visible = True), gradio.Button(visible = False), gradio.Dropdown() + + +def stop() -> Tuple[gradio.Button, gradio.Button]: + process_manager.stop() + return gradio.Button(visible = True), gradio.Button(visible = False) + + +def update_job_action(job_action : JobRunnerAction) -> gradio.Dropdown: + if job_action == 'job-run': + updated_job_ids = job_manager.find_job_ids('queued') or [ 'none' ] + + return gradio.Dropdown(value = get_last(updated_job_ids), choices = updated_job_ids, visible = True) + + if job_action == 'job-retry': + updated_job_ids = job_manager.find_job_ids('failed') or [ 'none' ] + + return gradio.Dropdown(value = get_last(updated_job_ids), choices = updated_job_ids, visible = True) + return gradio.Dropdown(visible = False) diff --git a/facefusion/uis/components/lip_syncer_options.py b/facefusion/uis/components/lip_syncer_options.py new file mode 100644 index 0000000000000000000000000000000000000000..4e28e4f14720158798f81796587ed79efc405e69 --- /dev/null +++ b/facefusion/uis/components/lip_syncer_options.py @@ -0,0 +1,64 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_float_step +from facefusion.processors.core import load_processor_module +from facefusion.processors.modules.lip_syncer import choices as lip_syncer_choices +from facefusion.processors.modules.lip_syncer.types import LipSyncerModel, LipSyncerWeight +from facefusion.uis.core import get_ui_component, register_ui_component + +LIP_SYNCER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +LIP_SYNCER_WEIGHT_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global LIP_SYNCER_MODEL_DROPDOWN + global LIP_SYNCER_WEIGHT_SLIDER + + has_lip_syncer = 'lip_syncer' in state_manager.get_item('processors') + LIP_SYNCER_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.model_dropdown', 'facefusion.processors.modules.lip_syncer'), + choices = lip_syncer_choices.lip_syncer_models, + value = state_manager.get_item('lip_syncer_model'), + visible = has_lip_syncer + ) + LIP_SYNCER_WEIGHT_SLIDER = gradio.Slider( + label = translator.get('uis.weight_slider', 'facefusion.processors.modules.lip_syncer'), + value = state_manager.get_item('lip_syncer_weight'), + step = calculate_float_step(lip_syncer_choices.lip_syncer_weight_range), + minimum = lip_syncer_choices.lip_syncer_weight_range[0], + maximum = lip_syncer_choices.lip_syncer_weight_range[-1], + visible = has_lip_syncer + ) + register_ui_component('lip_syncer_model_dropdown', LIP_SYNCER_MODEL_DROPDOWN) + register_ui_component('lip_syncer_weight_slider', LIP_SYNCER_WEIGHT_SLIDER) + + +def listen() -> None: + LIP_SYNCER_MODEL_DROPDOWN.change(update_lip_syncer_model, inputs = LIP_SYNCER_MODEL_DROPDOWN, outputs = LIP_SYNCER_MODEL_DROPDOWN) + LIP_SYNCER_WEIGHT_SLIDER.release(update_lip_syncer_weight, inputs = LIP_SYNCER_WEIGHT_SLIDER) + + processors_checkbox_group = get_ui_component('processors_checkbox_group') + if processors_checkbox_group: + processors_checkbox_group.change(remote_update, inputs = processors_checkbox_group, outputs = [ LIP_SYNCER_MODEL_DROPDOWN, LIP_SYNCER_WEIGHT_SLIDER ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider]: + has_lip_syncer = 'lip_syncer' in processors + return gradio.Dropdown(visible = has_lip_syncer), gradio.Slider(visible = has_lip_syncer) + + +def update_lip_syncer_model(lip_syncer_model : LipSyncerModel) -> gradio.Dropdown: + lip_syncer_module = load_processor_module('lip_syncer') + lip_syncer_module.clear_inference_pool() + state_manager.set_item('lip_syncer_model', lip_syncer_model) + + if lip_syncer_module.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('lip_syncer_model')) + return gradio.Dropdown() + + +def update_lip_syncer_weight(lip_syncer_weight : LipSyncerWeight) -> None: + state_manager.set_item('lip_syncer_weight', lip_syncer_weight) diff --git a/facefusion/uis/components/memory.py b/facefusion/uis/components/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..b5fafccecf6c4e607ebedd3efc9ffac3b75fca29 --- /dev/null +++ b/facefusion/uis/components/memory.py @@ -0,0 +1,42 @@ +from typing import Optional + +import gradio + +import facefusion.choices +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_int_step +from facefusion.types import VideoMemoryStrategy + +VIDEO_MEMORY_STRATEGY_DROPDOWN : Optional[gradio.Dropdown] = None +SYSTEM_MEMORY_LIMIT_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global VIDEO_MEMORY_STRATEGY_DROPDOWN + global SYSTEM_MEMORY_LIMIT_SLIDER + + VIDEO_MEMORY_STRATEGY_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.video_memory_strategy_dropdown'), + choices = facefusion.choices.video_memory_strategies, + value = state_manager.get_item('video_memory_strategy') + ) + SYSTEM_MEMORY_LIMIT_SLIDER = gradio.Slider( + label = translator.get('uis.system_memory_limit_slider'), + step = calculate_int_step(facefusion.choices.system_memory_limit_range), + minimum = facefusion.choices.system_memory_limit_range[0], + maximum = facefusion.choices.system_memory_limit_range[-1], + value = state_manager.get_item('system_memory_limit') + ) + + +def listen() -> None: + VIDEO_MEMORY_STRATEGY_DROPDOWN.change(update_video_memory_strategy, inputs = VIDEO_MEMORY_STRATEGY_DROPDOWN) + SYSTEM_MEMORY_LIMIT_SLIDER.release(update_system_memory_limit, inputs = SYSTEM_MEMORY_LIMIT_SLIDER) + + +def update_video_memory_strategy(video_memory_strategy : VideoMemoryStrategy) -> None: + state_manager.set_item('video_memory_strategy', video_memory_strategy) + + +def update_system_memory_limit(system_memory_limit : float) -> None: + state_manager.set_item('system_memory_limit', int(system_memory_limit)) diff --git a/facefusion/uis/components/output.py b/facefusion/uis/components/output.py new file mode 100644 index 0000000000000000000000000000000000000000..3a6f96c1676a47c530f7d66d856dc4587d92fbb1 --- /dev/null +++ b/facefusion/uis/components/output.py @@ -0,0 +1,48 @@ +import tempfile +from pathlib import Path +from typing import Optional + +import gradio + +from facefusion import state_manager, translator +from facefusion.uis.core import register_ui_component + +OUTPUT_PATH_TEXTBOX : Optional[gradio.Textbox] = None +OUTPUT_IMAGE : Optional[gradio.Image] = None +OUTPUT_VIDEO : Optional[gradio.Video] = None + + +def render() -> None: + global OUTPUT_PATH_TEXTBOX + global OUTPUT_IMAGE + global OUTPUT_VIDEO + + if not state_manager.get_item('output_path'): + documents_directory = Path.home().joinpath('Documents') + + if documents_directory.exists(): + state_manager.set_item('output_path', documents_directory) + else: + state_manager.set_item('output_path', tempfile.gettempdir()) + OUTPUT_PATH_TEXTBOX = gradio.Textbox( + label = translator.get('uis.output_path_textbox'), + value = state_manager.get_item('output_path'), + max_lines = 1 + ) + OUTPUT_IMAGE = gradio.Image( + label = translator.get('uis.output_image_or_video'), + visible = False + ) + OUTPUT_VIDEO = gradio.Video( + label = translator.get('uis.output_image_or_video') + ) + + +def listen() -> None: + OUTPUT_PATH_TEXTBOX.change(update_output_path, inputs = OUTPUT_PATH_TEXTBOX) + register_ui_component('output_image', OUTPUT_IMAGE) + register_ui_component('output_video', OUTPUT_VIDEO) + + +def update_output_path(output_path : str) -> None: + state_manager.set_item('output_path', output_path) diff --git a/facefusion/uis/components/output_options.py b/facefusion/uis/components/output_options.py new file mode 100644 index 0000000000000000000000000000000000000000..01a01c9619c74ab993d8479f5f3a0c6c2440d8c7 --- /dev/null +++ b/facefusion/uis/components/output_options.py @@ -0,0 +1,184 @@ +from typing import Optional, Tuple + +import gradio + +import facefusion.choices +from facefusion import state_manager, translator +from facefusion.common_helper import calculate_float_step, calculate_int_step +from facefusion.ffmpeg import get_available_encoder_set +from facefusion.filesystem import is_image, is_video +from facefusion.types import AudioEncoder, Fps, Scale, VideoEncoder, VideoPreset +from facefusion.uis.core import get_ui_components, register_ui_component +from facefusion.vision import detect_video_fps + +OUTPUT_IMAGE_QUALITY_SLIDER : Optional[gradio.Slider] = None +OUTPUT_IMAGE_SCALE_SLIDER : Optional[gradio.Slider] = None +OUTPUT_AUDIO_ENCODER_DROPDOWN : Optional[gradio.Dropdown] = None +OUTPUT_AUDIO_QUALITY_SLIDER : Optional[gradio.Slider] = None +OUTPUT_AUDIO_VOLUME_SLIDER : Optional[gradio.Slider] = None +OUTPUT_VIDEO_ENCODER_DROPDOWN : Optional[gradio.Dropdown] = None +OUTPUT_VIDEO_PRESET_DROPDOWN : Optional[gradio.Dropdown] = None +OUTPUT_VIDEO_SCALE_SLIDER : Optional[gradio.Slider] = None +OUTPUT_VIDEO_QUALITY_SLIDER : Optional[gradio.Slider] = None +OUTPUT_VIDEO_FPS_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global OUTPUT_IMAGE_QUALITY_SLIDER + global OUTPUT_IMAGE_SCALE_SLIDER + global OUTPUT_AUDIO_ENCODER_DROPDOWN + global OUTPUT_AUDIO_QUALITY_SLIDER + global OUTPUT_AUDIO_VOLUME_SLIDER + global OUTPUT_VIDEO_ENCODER_DROPDOWN + global OUTPUT_VIDEO_PRESET_DROPDOWN + global OUTPUT_VIDEO_SCALE_SLIDER + global OUTPUT_VIDEO_QUALITY_SLIDER + global OUTPUT_VIDEO_FPS_SLIDER + + available_encoder_set = get_available_encoder_set() + + OUTPUT_IMAGE_QUALITY_SLIDER = gradio.Slider( + label = translator.get('uis.output_image_quality_slider'), + value = state_manager.get_item('output_image_quality'), + step = calculate_int_step(facefusion.choices.output_image_quality_range), + minimum = facefusion.choices.output_image_quality_range[0], + maximum = facefusion.choices.output_image_quality_range[-1], + visible = is_image(state_manager.get_item('target_path')) + ) + OUTPUT_IMAGE_SCALE_SLIDER = gradio.Slider( + label = translator.get('uis.output_image_scale_slider'), + step = calculate_float_step(facefusion.choices.output_image_scale_range), + value = state_manager.get_item('output_image_scale'), + minimum = facefusion.choices.output_image_scale_range[0], + maximum = facefusion.choices.output_image_scale_range[-1], + visible = is_image(state_manager.get_item('target_path')) + ) + OUTPUT_AUDIO_ENCODER_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.output_audio_encoder_dropdown'), + choices = available_encoder_set.get('audio'), + value = state_manager.get_item('output_audio_encoder'), + visible = is_video(state_manager.get_item('target_path')) + ) + OUTPUT_AUDIO_QUALITY_SLIDER = gradio.Slider( + label = translator.get('uis.output_audio_quality_slider'), + value = state_manager.get_item('output_audio_quality'), + step = calculate_int_step(facefusion.choices.output_audio_quality_range), + minimum = facefusion.choices.output_audio_quality_range[0], + maximum = facefusion.choices.output_audio_quality_range[-1], + visible = is_video(state_manager.get_item('target_path')) + ) + OUTPUT_AUDIO_VOLUME_SLIDER = gradio.Slider( + label = translator.get('uis.output_audio_volume_slider'), + value = state_manager.get_item('output_audio_volume'), + step = calculate_int_step(facefusion.choices.output_audio_volume_range), + minimum = facefusion.choices.output_audio_volume_range[0], + maximum = facefusion.choices.output_audio_volume_range[-1], + visible = is_video(state_manager.get_item('target_path')) + ) + OUTPUT_VIDEO_ENCODER_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.output_video_encoder_dropdown'), + choices = available_encoder_set.get('video'), + value = state_manager.get_item('output_video_encoder'), + visible = is_video(state_manager.get_item('target_path')) + ) + OUTPUT_VIDEO_PRESET_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.output_video_preset_dropdown'), + choices = facefusion.choices.output_video_presets, + value = state_manager.get_item('output_video_preset'), + visible = is_video(state_manager.get_item('target_path')) + ) + OUTPUT_VIDEO_QUALITY_SLIDER = gradio.Slider( + label = translator.get('uis.output_video_quality_slider'), + value = state_manager.get_item('output_video_quality'), + step = calculate_int_step(facefusion.choices.output_video_quality_range), + minimum = facefusion.choices.output_video_quality_range[0], + maximum = facefusion.choices.output_video_quality_range[-1], + visible = is_video(state_manager.get_item('target_path')) + ) + OUTPUT_VIDEO_SCALE_SLIDER = gradio.Slider( + label = translator.get('uis.output_video_scale_slider'), + step = calculate_float_step(facefusion.choices.output_video_scale_range), + value = state_manager.get_item('output_video_scale'), + minimum = facefusion.choices.output_video_scale_range[0], + maximum = facefusion.choices.output_video_scale_range[-1], + visible = is_video(state_manager.get_item('target_path')) + ) + OUTPUT_VIDEO_FPS_SLIDER = gradio.Slider( + label = translator.get('uis.output_video_fps_slider'), + value = state_manager.get_item('output_video_fps'), + step = 0.01, + minimum = 1, + maximum = 60, + visible = is_video(state_manager.get_item('target_path')) + ) + register_ui_component('output_video_fps_slider', OUTPUT_VIDEO_FPS_SLIDER) + + +def listen() -> None: + OUTPUT_IMAGE_QUALITY_SLIDER.release(update_output_image_quality, inputs = OUTPUT_IMAGE_QUALITY_SLIDER) + OUTPUT_IMAGE_SCALE_SLIDER.release(update_output_image_scale, inputs = OUTPUT_IMAGE_SCALE_SLIDER) + OUTPUT_AUDIO_ENCODER_DROPDOWN.change(update_output_audio_encoder, inputs = OUTPUT_AUDIO_ENCODER_DROPDOWN) + OUTPUT_AUDIO_QUALITY_SLIDER.release(update_output_audio_quality, inputs = OUTPUT_AUDIO_QUALITY_SLIDER) + OUTPUT_AUDIO_VOLUME_SLIDER.release(update_output_audio_volume, inputs = OUTPUT_AUDIO_VOLUME_SLIDER) + OUTPUT_VIDEO_ENCODER_DROPDOWN.change(update_output_video_encoder, inputs = OUTPUT_VIDEO_ENCODER_DROPDOWN) + OUTPUT_VIDEO_PRESET_DROPDOWN.change(update_output_video_preset, inputs = OUTPUT_VIDEO_PRESET_DROPDOWN) + OUTPUT_VIDEO_QUALITY_SLIDER.release(update_output_video_quality, inputs = OUTPUT_VIDEO_QUALITY_SLIDER) + OUTPUT_VIDEO_SCALE_SLIDER.release(update_output_video_scale, inputs = OUTPUT_VIDEO_SCALE_SLIDER) + OUTPUT_VIDEO_FPS_SLIDER.release(update_output_video_fps, inputs = OUTPUT_VIDEO_FPS_SLIDER) + + for ui_component in get_ui_components( + [ + 'target_image', + 'target_video' + ]): + for method in [ 'change', 'clear' ]: + getattr(ui_component, method)(remote_update, outputs = [OUTPUT_IMAGE_QUALITY_SLIDER, OUTPUT_IMAGE_SCALE_SLIDER, OUTPUT_AUDIO_ENCODER_DROPDOWN, OUTPUT_AUDIO_QUALITY_SLIDER, OUTPUT_AUDIO_VOLUME_SLIDER, OUTPUT_VIDEO_ENCODER_DROPDOWN, OUTPUT_VIDEO_PRESET_DROPDOWN, OUTPUT_VIDEO_QUALITY_SLIDER, OUTPUT_VIDEO_SCALE_SLIDER, OUTPUT_VIDEO_FPS_SLIDER]) + + +def remote_update() -> Tuple[gradio.Slider, gradio.Slider, gradio.Dropdown, gradio.Slider, gradio.Slider, gradio.Dropdown, gradio.Dropdown, gradio.Slider, gradio.Slider, gradio.Slider]: + if is_image(state_manager.get_item('target_path')): + return gradio.Slider(visible = True), gradio.Slider(visible = True), gradio.Dropdown(visible = False), gradio.Slider(visible = False), gradio.Slider(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False), gradio.Slider(visible = False), gradio.Slider(visible = False) + if is_video(state_manager.get_item('target_path')): + state_manager.set_item('output_video_fps', detect_video_fps(state_manager.get_item('target_path'))) + return gradio.Slider(visible = False), gradio.Slider(visible = False), gradio.Dropdown(visible = True), gradio.Slider(visible = True), gradio.Slider(visible = True), gradio.Dropdown(visible = True), gradio.Dropdown(visible = True), gradio.Slider(visible = True), gradio.Slider(visible = True), gradio.Slider(value = state_manager.get_item('output_video_fps'), visible = True) + return gradio.Slider(visible = False), gradio.Slider(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False), gradio.Slider(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False), gradio.Slider(visible = False), gradio.Slider(visible = False) + + +def update_output_image_quality(output_image_quality : float) -> None: + state_manager.set_item('output_image_quality', int(output_image_quality)) + + +def update_output_image_scale(output_image_scale : Scale) -> None: + state_manager.set_item('output_image_scale', output_image_scale) + + +def update_output_audio_encoder(output_audio_encoder : AudioEncoder) -> None: + state_manager.set_item('output_audio_encoder', output_audio_encoder) + + +def update_output_audio_quality(output_audio_quality : float) -> None: + state_manager.set_item('output_audio_quality', int(output_audio_quality)) + + +def update_output_audio_volume(output_audio_volume: float) -> None: + state_manager.set_item('output_audio_volume', int(output_audio_volume)) + + +def update_output_video_encoder(output_video_encoder : VideoEncoder) -> None: + state_manager.set_item('output_video_encoder', output_video_encoder) + + +def update_output_video_preset(output_video_preset : VideoPreset) -> None: + state_manager.set_item('output_video_preset', output_video_preset) + + +def update_output_video_quality(output_video_quality : float) -> None: + state_manager.set_item('output_video_quality', int(output_video_quality)) + + +def update_output_video_scale(output_video_scale : Scale) -> None: + state_manager.set_item('output_video_scale', output_video_scale) + + +def update_output_video_fps(output_video_fps : Fps) -> None: + state_manager.set_item('output_video_fps', output_video_fps) diff --git a/facefusion/uis/components/preview.py b/facefusion/uis/components/preview.py new file mode 100644 index 0000000000000000000000000000000000000000..8d86c3d2eb95e27e9a3c4fc5e875c9ae42b3a1e4 --- /dev/null +++ b/facefusion/uis/components/preview.py @@ -0,0 +1,302 @@ +from time import sleep +from typing import List, Optional, Tuple + +import cv2 +import gradio +import numpy + +from facefusion import logger, process_manager, state_manager, translator +from facefusion.audio import create_empty_audio_frame, get_voice_frame +from facefusion.common_helper import get_first +from facefusion.content_analyser import analyse_frame +from facefusion.face_analyser import get_one_face +from facefusion.face_selector import select_faces +from facefusion.face_store import clear_static_faces +from facefusion.filesystem import filter_audio_paths, is_image, is_video +from facefusion.processors.core import get_processors_modules +from facefusion.types import AudioFrame, Face, Mask, VisionFrame +from facefusion.uis import choices as uis_choices +from facefusion.uis.core import get_ui_component, get_ui_components, register_ui_component +from facefusion.uis.types import ComponentOptions, PreviewMode +from facefusion.vision import detect_frame_orientation, extract_vision_mask, fit_cover_frame, merge_vision_mask, obscure_frame, read_static_image, read_static_images, read_video_frame, restrict_frame, unpack_resolution + +PREVIEW_IMAGE : Optional[gradio.Image] = None + + +def render() -> None: + global PREVIEW_IMAGE + + preview_image_options : ComponentOptions =\ + { + 'label': translator.get('uis.preview_image') + } + + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) + source_audio_frame = create_empty_audio_frame() + source_voice_frame = create_empty_audio_frame() + + if source_audio_path and state_manager.get_item('output_video_fps') and state_manager.get_item('reference_frame_number'): + temp_voice_frame = get_voice_frame(source_audio_path, state_manager.get_item('output_video_fps'), state_manager.get_item('reference_frame_number')) + if numpy.any(temp_voice_frame): + source_voice_frame = temp_voice_frame + + if is_image(state_manager.get_item('target_path')): + target_vision_frame = read_static_image(state_manager.get_item('target_path')) + reference_vision_frame = read_static_image(state_manager.get_item('target_path')) + preview_vision_frame = process_preview_frame(reference_vision_frame, source_vision_frames, source_audio_frame, source_voice_frame, target_vision_frame, uis_choices.preview_modes[0], uis_choices.preview_resolutions[-1]) + preview_image_options['value'] = cv2.cvtColor(preview_vision_frame, cv2.COLOR_BGR2RGB) + preview_image_options['elem_classes'] = [ 'image-preview', 'is-' + detect_frame_orientation(preview_vision_frame) ] + + if is_video(state_manager.get_item('target_path')): + temp_vision_frame = read_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) + reference_vision_frame = read_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) + preview_vision_frame = process_preview_frame(reference_vision_frame, source_vision_frames, source_audio_frame, source_voice_frame, temp_vision_frame, uis_choices.preview_modes[0], uis_choices.preview_resolutions[-1]) + preview_image_options['value'] = cv2.cvtColor(preview_vision_frame, cv2.COLOR_BGR2RGB) + preview_image_options['elem_classes'] = [ 'image-preview', 'is-' + detect_frame_orientation(preview_vision_frame) ] + preview_image_options['visible'] = True + PREVIEW_IMAGE = gradio.Image(**preview_image_options) + register_ui_component('preview_image', PREVIEW_IMAGE) + + +def listen() -> None: + preview_frame_slider = get_ui_component('preview_frame_slider') + preview_mode_dropdown = get_ui_component('preview_mode_dropdown') + preview_resolution_dropdown = get_ui_component('preview_resolution_dropdown') + + if preview_mode_dropdown: + preview_mode_dropdown.change(update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE) + + if preview_resolution_dropdown: + preview_resolution_dropdown.change(update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE) + + if preview_frame_slider: + preview_frame_slider.release(update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE, show_progress = 'hidden') + preview_frame_slider.change(update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE, show_progress = 'hidden', trigger_mode = 'once') + + reference_face_position_gallery = get_ui_component('reference_face_position_gallery') + if reference_face_position_gallery: + reference_face_position_gallery.select(clear_and_update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE) + + for ui_component in get_ui_components( + [ + 'source_audio', + 'source_image', + 'target_image', + 'target_video' + ]): + for method in [ 'change', 'clear' ]: + getattr(ui_component, method)(update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE) + + for ui_component in get_ui_components( + [ + 'background_remover_color_red_number', + 'background_remover_color_green_number', + 'background_remover_color_blue_number', + 'background_remover_color_alpha_number', + 'face_debugger_items_checkbox_group', + 'frame_colorizer_size_dropdown', + 'face_mask_types_checkbox_group', + 'face_mask_areas_checkbox_group', + 'face_mask_regions_checkbox_group', + 'expression_restorer_areas_checkbox_group' + ]): + ui_component.change(update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE) + + for ui_component in get_ui_components( + [ + 'age_modifier_direction_slider', + 'deep_swapper_morph_slider', + 'expression_restorer_factor_slider', + 'face_editor_eyebrow_direction_slider', + 'face_editor_eye_gaze_horizontal_slider', + 'face_editor_eye_gaze_vertical_slider', + 'face_editor_eye_open_ratio_slider', + 'face_editor_lip_open_ratio_slider', + 'face_editor_mouth_grim_slider', + 'face_editor_mouth_pout_slider', + 'face_editor_mouth_purse_slider', + 'face_editor_mouth_smile_slider', + 'face_editor_mouth_position_horizontal_slider', + 'face_editor_mouth_position_vertical_slider', + 'face_editor_head_pitch_slider', + 'face_editor_head_yaw_slider', + 'face_editor_head_roll_slider', + 'face_enhancer_blend_slider', + 'face_enhancer_weight_slider', + 'face_swapper_weight_slider', + 'frame_colorizer_blend_slider', + 'frame_enhancer_blend_slider', + 'lip_syncer_weight_slider', + 'reference_face_distance_slider', + 'face_selector_age_range_slider', + 'face_mask_blur_slider', + 'face_mask_padding_top_slider', + 'face_mask_padding_bottom_slider', + 'face_mask_padding_left_slider', + 'face_mask_padding_right_slider', + 'output_video_fps_slider' + ]): + ui_component.release(update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE) + + for ui_component in get_ui_components( + [ + 'age_modifier_model_dropdown', + 'background_remover_model_dropdown', + 'deep_swapper_model_dropdown', + 'expression_restorer_model_dropdown', + 'processors_checkbox_group', + 'face_editor_model_dropdown', + 'face_enhancer_model_dropdown', + 'face_swapper_model_dropdown', + 'face_swapper_pixel_boost_dropdown', + 'frame_colorizer_model_dropdown', + 'frame_enhancer_model_dropdown', + 'lip_syncer_model_dropdown', + 'face_selector_mode_dropdown', + 'face_selector_order_dropdown', + 'face_selector_gender_dropdown', + 'face_selector_race_dropdown', + 'face_detector_model_dropdown', + 'face_detector_size_dropdown', + 'face_detector_angles_checkbox_group', + 'face_landmarker_model_dropdown', + 'face_occluder_model_dropdown', + 'face_parser_model_dropdown', + 'voice_extractor_model_dropdown' + ]): + ui_component.change(clear_and_update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE) + + for ui_component in get_ui_components( + [ + 'face_detector_margin_slider', + 'face_detector_score_slider', + 'face_landmarker_score_slider' + ]): + ui_component.release(clear_and_update_preview_image, inputs = [ preview_mode_dropdown, preview_resolution_dropdown, preview_frame_slider ], outputs = PREVIEW_IMAGE) + + +def update_preview_image(preview_mode : PreviewMode, preview_resolution : str, frame_number : int = 0) -> gradio.Image: + while process_manager.is_checking(): + sleep(0.5) + + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) + source_audio_frame = create_empty_audio_frame() + source_voice_frame = create_empty_audio_frame() + + if source_audio_path and state_manager.get_item('output_video_fps') and state_manager.get_item('reference_frame_number'): + reference_audio_frame_number = state_manager.get_item('reference_frame_number') + if state_manager.get_item('trim_frame_start'): + reference_audio_frame_number -= state_manager.get_item('trim_frame_start') + temp_voice_frame = get_voice_frame(source_audio_path, state_manager.get_item('output_video_fps'), reference_audio_frame_number) + if numpy.any(temp_voice_frame): + source_voice_frame = temp_voice_frame + + if is_image(state_manager.get_item('target_path')): + reference_vision_frame = read_static_image(state_manager.get_item('target_path')) + target_vision_frame = read_static_image(state_manager.get_item('target_path'), 'rgba') + target_vision_mask = extract_vision_mask(target_vision_frame) + target_vision_frame = merge_vision_mask(target_vision_frame, target_vision_mask) + preview_vision_frame = process_preview_frame(reference_vision_frame, source_vision_frames, source_audio_frame, source_voice_frame, target_vision_frame, preview_mode, preview_resolution) + preview_vision_frame = cv2.cvtColor(preview_vision_frame, cv2.COLOR_BGRA2RGBA) + return gradio.Image(value = preview_vision_frame, elem_classes = [ 'image-preview', 'is-' + detect_frame_orientation(preview_vision_frame) ]) + + if is_video(state_manager.get_item('target_path')): + reference_vision_frame = read_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) + temp_vision_frame = read_video_frame(state_manager.get_item('target_path'), frame_number) + temp_vision_mask = extract_vision_mask(temp_vision_frame) + temp_vision_frame = merge_vision_mask(temp_vision_frame, temp_vision_mask) + preview_vision_frame = process_preview_frame(reference_vision_frame, source_vision_frames, source_audio_frame, source_voice_frame, temp_vision_frame, preview_mode, preview_resolution) + preview_vision_frame = cv2.cvtColor(preview_vision_frame, cv2.COLOR_BGRA2RGBA) + return gradio.Image(value = preview_vision_frame, elem_classes = [ 'image-preview', 'is-' + detect_frame_orientation(preview_vision_frame) ]) + return gradio.Image(value = None, elem_classes = None) + + +def clear_and_update_preview_image(preview_mode : PreviewMode, preview_resolution : str, frame_number : int = 0) -> gradio.Image: + clear_static_faces() + return update_preview_image(preview_mode, preview_resolution, frame_number) + + +def process_preview_frame(reference_vision_frame : VisionFrame, source_vision_frames : List[VisionFrame], source_audio_frame : AudioFrame, source_voice_frame : AudioFrame, target_vision_frame : VisionFrame, preview_mode : PreviewMode, preview_resolution : str) -> VisionFrame: + target_vision_frame = restrict_frame(target_vision_frame, unpack_resolution(preview_resolution)) + temp_vision_frame = target_vision_frame.copy() + temp_vision_mask = extract_vision_mask(temp_vision_frame) + + if analyse_frame(target_vision_frame[:, :, :3]): + if preview_mode == 'frame-by-frame': + temp_vision_frame = obscure_frame(temp_vision_frame[:, :, :3]) + return numpy.hstack((temp_vision_frame, temp_vision_frame)) + + if preview_mode == 'face-by-face': + target_crop_vision_frame, output_crop_vision_frame = create_face_by_face(reference_vision_frame, target_vision_frame[:, :, :3], temp_vision_frame[:, :, :3]) + target_crop_vision_frame = obscure_frame(target_crop_vision_frame) + output_crop_vision_frame = obscure_frame(output_crop_vision_frame) + return numpy.hstack((target_crop_vision_frame, output_crop_vision_frame)) + + temp_vision_frame = obscure_frame(temp_vision_frame) + return temp_vision_frame + + for processor_module in get_processors_modules(state_manager.get_item('processors')): + logger.disable() + if processor_module.pre_process('preview'): + logger.enable() + temp_vision_frame, temp_vision_mask = processor_module.process_frame( + { + 'reference_vision_frame': reference_vision_frame, + 'source_audio_frame': source_audio_frame, + 'source_voice_frame': source_voice_frame, + 'source_vision_frames': source_vision_frames, + 'target_vision_frame': target_vision_frame[:, :, :3], + 'temp_vision_frame': temp_vision_frame[:, :, :3], + 'temp_vision_mask': temp_vision_mask + }) + logger.enable() + + temp_vision_frame = prepare_output_frame(target_vision_frame, temp_vision_frame, temp_vision_mask) + + if preview_mode == 'frame-by-frame': + return numpy.hstack((target_vision_frame, temp_vision_frame)) + + if preview_mode == 'face-by-face': + target_crop_vision_frame, output_crop_vision_frame = create_face_by_face(reference_vision_frame, target_vision_frame, temp_vision_frame) + return numpy.hstack((target_crop_vision_frame, output_crop_vision_frame)) + + return temp_vision_frame + + +def create_face_by_face(reference_vision_frame : VisionFrame, target_vision_frame : VisionFrame, temp_vision_frame : VisionFrame) -> Tuple[VisionFrame, VisionFrame]: + target_faces = select_faces(reference_vision_frame[:, :, :3], target_vision_frame[:, :, :3]) + target_face = get_one_face(target_faces) + + if target_face: + target_crop_vision_frame = extract_crop_frame(target_vision_frame, target_face) + output_crop_vision_frame = extract_crop_frame(temp_vision_frame, target_face) + + if numpy.any(target_crop_vision_frame) and numpy.any(output_crop_vision_frame): + target_crop_dimension = min(target_crop_vision_frame.shape[:2]) + target_crop_vision_frame = fit_cover_frame(target_crop_vision_frame, (target_crop_dimension, target_crop_dimension)) + output_crop_vision_frame = fit_cover_frame(output_crop_vision_frame, (target_crop_dimension, target_crop_dimension)) + return target_crop_vision_frame, output_crop_vision_frame + + empty_vision_frame = numpy.zeros((512, 512, 4), dtype = numpy.uint8) + return empty_vision_frame, empty_vision_frame + + +def extract_crop_frame(vision_frame : VisionFrame, face : Face) -> Optional[VisionFrame]: + start_x, start_y, end_x, end_y = map(int, face.bounding_box) + padding_x = int((end_x - start_x) * 0.25) + padding_y = int((end_y - start_y) * 0.25) + start_x = max(0, start_x - padding_x) + start_y = max(0, start_y - padding_y) + end_x = max(0, end_x + padding_x) + end_y = max(0, end_y + padding_y) + crop_vision_frame = vision_frame[start_y:end_y, start_x:end_x] + return crop_vision_frame + + +def prepare_output_frame(target_vision_frame : VisionFrame, temp_vision_frame : VisionFrame, temp_vision_mask : Mask) -> VisionFrame: + temp_vision_mask = temp_vision_mask.clip(state_manager.get_item('background_remover_color')[-1], 255) + temp_vision_frame = merge_vision_mask(temp_vision_frame, temp_vision_mask) + temp_vision_frame = cv2.resize(temp_vision_frame, target_vision_frame.shape[1::-1]) + return temp_vision_frame diff --git a/facefusion/uis/components/preview_options.py b/facefusion/uis/components/preview_options.py new file mode 100644 index 0000000000000000000000000000000000000000..f230614680fae6acbd4543cb10df57c6c06d4487 --- /dev/null +++ b/facefusion/uis/components/preview_options.py @@ -0,0 +1,61 @@ +from typing import Optional + +import gradio + +from facefusion import state_manager, translator +from facefusion.filesystem import is_video +from facefusion.uis import choices as uis_choices +from facefusion.uis.core import get_ui_components, register_ui_component +from facefusion.uis.types import ComponentOptions +from facefusion.vision import count_video_frame_total + +PREVIEW_FRAME_SLIDER: Optional[gradio.Slider] = None +PREVIEW_MODE_DROPDOWN: Optional[gradio.Dropdown] = None +PREVIEW_RESOLUTION_DROPDOWN: Optional[gradio.Dropdown] = None + + +def render() -> None: + global PREVIEW_FRAME_SLIDER, PREVIEW_MODE_DROPDOWN, PREVIEW_RESOLUTION_DROPDOWN + + preview_frame_slider_options : ComponentOptions =\ + { + 'label': translator.get('uis.preview_frame_slider'), + 'step': 1, + 'minimum': 0, + 'maximum': 100, + 'visible': False + } + if is_video(state_manager.get_item('target_path')): + preview_frame_slider_options['value'] = state_manager.get_item('reference_frame_number') + preview_frame_slider_options['maximum'] = count_video_frame_total(state_manager.get_item('target_path')) + preview_frame_slider_options['visible'] = True + PREVIEW_FRAME_SLIDER = gradio.Slider(**preview_frame_slider_options) + with gradio.Row(): + PREVIEW_MODE_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.preview_mode_dropdown'), + value = uis_choices.preview_modes[0], + choices = uis_choices.preview_modes, + visible = True + ) + PREVIEW_RESOLUTION_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.preview_resolution_dropdown'), + value = uis_choices.preview_resolutions[-1], + choices = uis_choices.preview_resolutions, + visible = True + ) + register_ui_component('preview_mode_dropdown', PREVIEW_MODE_DROPDOWN) + register_ui_component('preview_resolution_dropdown', PREVIEW_RESOLUTION_DROPDOWN) + register_ui_component('preview_frame_slider', PREVIEW_FRAME_SLIDER) + + +def listen() -> None: + for ui_component in get_ui_components([ 'target_image', 'target_video' ]): + for method in [ 'change', 'clear' ]: + getattr(ui_component, method)(update_preview_frame_slider, outputs = PREVIEW_FRAME_SLIDER) + + +def update_preview_frame_slider() -> gradio.Slider: + if is_video(state_manager.get_item('target_path')): + video_frame_total = count_video_frame_total(state_manager.get_item('target_path')) + return gradio.Slider(maximum = video_frame_total, visible = True) + return gradio.Slider(value = 0, visible = False) diff --git a/facefusion/uis/components/processors.py b/facefusion/uis/components/processors.py new file mode 100644 index 0000000000000000000000000000000000000000..281cb72c544e54a909fb94982d32c31adb09ff30 --- /dev/null +++ b/facefusion/uis/components/processors.py @@ -0,0 +1,49 @@ +from typing import List, Optional + +import gradio + +from facefusion import state_manager, translator +from facefusion.filesystem import get_file_name, resolve_file_paths +from facefusion.processors.core import get_processors_modules +from facefusion.uis.core import register_ui_component + +PROCESSORS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None + + +def render() -> None: + global PROCESSORS_CHECKBOX_GROUP + + PROCESSORS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = translator.get('uis.processors_checkbox_group'), + choices = sort_processors(state_manager.get_item('processors')), + value = state_manager.get_item('processors') + ) + register_ui_component('processors_checkbox_group', PROCESSORS_CHECKBOX_GROUP) + + +def listen() -> None: + PROCESSORS_CHECKBOX_GROUP.change(update_processors, inputs = PROCESSORS_CHECKBOX_GROUP, outputs = PROCESSORS_CHECKBOX_GROUP) + + +def update_processors(processors : List[str]) -> gradio.CheckboxGroup: + for processor_module in get_processors_modules(state_manager.get_item('processors')): + if hasattr(processor_module, 'clear_inference_pool'): + processor_module.clear_inference_pool() + + for processor_module in get_processors_modules(processors): + if not processor_module.pre_check(): + return gradio.CheckboxGroup() + + state_manager.set_item('processors', processors) + return gradio.CheckboxGroup(value = state_manager.get_item('processors'), choices = sort_processors(state_manager.get_item('processors'))) + + +def sort_processors(processors : List[str]) -> List[str]: + available_processors = [ get_file_name(file_path) for file_path in resolve_file_paths('facefusion/processors/modules') ] + current_processors = [] + + for processor in processors + available_processors: + if processor in available_processors and processor not in current_processors: + current_processors.append(processor) + + return current_processors diff --git a/facefusion/uis/components/source.py b/facefusion/uis/components/source.py new file mode 100644 index 0000000000000000000000000000000000000000..c57c5b9867efa72f9fafec8bf1061d89bf34d086 --- /dev/null +++ b/facefusion/uis/components/source.py @@ -0,0 +1,61 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.common_helper import get_first +from facefusion.filesystem import filter_audio_paths, filter_image_paths, has_audio, has_image +from facefusion.uis.core import register_ui_component +from facefusion.uis.types import File + +SOURCE_FILE : Optional[gradio.File] = None +SOURCE_AUDIO : Optional[gradio.Audio] = None +SOURCE_IMAGE : Optional[gradio.Image] = None + + +def render() -> None: + global SOURCE_FILE + global SOURCE_AUDIO + global SOURCE_IMAGE + + has_source_audio = has_audio(state_manager.get_item('source_paths')) + has_source_image = has_image(state_manager.get_item('source_paths')) + SOURCE_FILE = gradio.File( + label = translator.get('uis.source_file'), + file_count = 'multiple', + value = state_manager.get_item('source_paths') if has_source_audio or has_source_image else None + ) + source_file_names = [ source_file_value.get('path') for source_file_value in SOURCE_FILE.value ] if SOURCE_FILE.value else None + source_audio_path = get_first(filter_audio_paths(source_file_names)) + source_image_path = get_first(filter_image_paths(source_file_names)) + SOURCE_AUDIO = gradio.Audio( + value = source_audio_path if has_source_audio else None, + visible = has_source_audio, + show_label = False + ) + SOURCE_IMAGE = gradio.Image( + value = source_image_path if has_source_image else None, + visible = has_source_image, + show_label = False + ) + register_ui_component('source_audio', SOURCE_AUDIO) + register_ui_component('source_image', SOURCE_IMAGE) + + +def listen() -> None: + SOURCE_FILE.change(update, inputs = SOURCE_FILE, outputs = [ SOURCE_AUDIO, SOURCE_IMAGE ]) + + +def update(files : List[File]) -> Tuple[gradio.Audio, gradio.Image]: + file_names = [ file.name for file in files ] if files else None + has_source_audio = has_audio(file_names) + has_source_image = has_image(file_names) + + if has_source_audio or has_source_image: + source_audio_path = get_first(filter_audio_paths(file_names)) + source_image_path = get_first(filter_image_paths(file_names)) + state_manager.set_item('source_paths', file_names) + return gradio.Audio(value = source_audio_path, visible = has_source_audio), gradio.Image(value = source_image_path, visible = has_source_image) + + state_manager.clear_item('source_paths') + return gradio.Audio(value = None, visible = False), gradio.Image(value = None, visible = False) diff --git a/facefusion/uis/components/target.py b/facefusion/uis/components/target.py new file mode 100644 index 0000000000000000000000000000000000000000..8cc4d77149d06ecf9dec0987fbfb6519230efa7a --- /dev/null +++ b/facefusion/uis/components/target.py @@ -0,0 +1,65 @@ +from typing import Optional, Tuple + +import gradio + +from facefusion import state_manager, translator +from facefusion.face_store import clear_static_faces +from facefusion.filesystem import is_image, is_video +from facefusion.uis.core import register_ui_component +from facefusion.uis.types import ComponentOptions, File + +TARGET_FILE : Optional[gradio.File] = None +TARGET_IMAGE : Optional[gradio.Image] = None +TARGET_VIDEO : Optional[gradio.Video] = None + + +def render() -> None: + global TARGET_FILE + global TARGET_IMAGE + global TARGET_VIDEO + + is_target_image = is_image(state_manager.get_item('target_path')) + is_target_video = is_video(state_manager.get_item('target_path')) + TARGET_FILE = gradio.File( + label = translator.get('uis.target_file'), + value = state_manager.get_item('target_path') if is_target_image or is_target_video else None + ) + target_image_options : ComponentOptions =\ + { + 'show_label': False, + 'visible': False + } + target_video_options : ComponentOptions =\ + { + 'show_label': False, + 'visible': False + } + if is_target_image: + target_image_options['value'] = TARGET_FILE.value.get('path') + target_image_options['visible'] = True + if is_target_video: + target_video_options['value'] = TARGET_FILE.value.get('path') + target_video_options['visible'] = True + TARGET_IMAGE = gradio.Image(**target_image_options) + TARGET_VIDEO = gradio.Video(**target_video_options) + register_ui_component('target_image', TARGET_IMAGE) + register_ui_component('target_video', TARGET_VIDEO) + + +def listen() -> None: + TARGET_FILE.change(update, inputs = TARGET_FILE, outputs = [ TARGET_IMAGE, TARGET_VIDEO ]) + + +def update(file : File) -> Tuple[gradio.Image, gradio.Video]: + clear_static_faces() + + if file and is_image(file.name): + state_manager.set_item('target_path', file.name) + return gradio.Image(value = file.name, visible = True), gradio.Video(value = None, visible = False) + + if file and is_video(file.name): + state_manager.set_item('target_path', file.name) + return gradio.Image(value = None, visible = False), gradio.Video(value = file.name, visible = True) + + state_manager.clear_item('target_path') + return gradio.Image(value = None, visible = False), gradio.Video(value = None, visible = False) diff --git a/facefusion/uis/components/temp_frame.py b/facefusion/uis/components/temp_frame.py new file mode 100644 index 0000000000000000000000000000000000000000..8e608dbaaa0565b88d9ada5de8570fe482eb9b58 --- /dev/null +++ b/facefusion/uis/components/temp_frame.py @@ -0,0 +1,42 @@ +from typing import Optional + +import gradio + +import facefusion.choices +from facefusion import state_manager, translator +from facefusion.filesystem import is_video +from facefusion.types import TempFrameFormat +from facefusion.uis.core import get_ui_component + +TEMP_FRAME_FORMAT_DROPDOWN : Optional[gradio.Dropdown] = None + + +def render() -> None: + global TEMP_FRAME_FORMAT_DROPDOWN + + TEMP_FRAME_FORMAT_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.temp_frame_format_dropdown'), + choices = facefusion.choices.temp_frame_formats, + value = state_manager.get_item('temp_frame_format'), + visible = is_video(state_manager.get_item('target_path')) + ) + + +def listen() -> None: + TEMP_FRAME_FORMAT_DROPDOWN.change(update_temp_frame_format, inputs = TEMP_FRAME_FORMAT_DROPDOWN) + + target_video = get_ui_component('target_video') + if target_video: + for method in [ 'change', 'clear' ]: + getattr(target_video, method)(remote_update, outputs = TEMP_FRAME_FORMAT_DROPDOWN) + + +def remote_update() -> gradio.Dropdown: + if is_video(state_manager.get_item('target_path')): + return gradio.Dropdown(visible = True) + return gradio.Dropdown(visible = False) + + +def update_temp_frame_format(temp_frame_format : TempFrameFormat) -> None: + state_manager.set_item('temp_frame_format', temp_frame_format) + diff --git a/facefusion/uis/components/terminal.py b/facefusion/uis/components/terminal.py new file mode 100644 index 0000000000000000000000000000000000000000..2f47ff3dcdc279014362e1ec260871de5f4fbd65 --- /dev/null +++ b/facefusion/uis/components/terminal.py @@ -0,0 +1,80 @@ +import io +import logging +import math +import os +from typing import Optional + +import gradio +from tqdm import tqdm + +import facefusion.choices +from facefusion import logger, state_manager, translator +from facefusion.types import LogLevel + +LOG_LEVEL_DROPDOWN : Optional[gradio.Dropdown] = None +TERMINAL_TEXTBOX : Optional[gradio.Textbox] = None +LOG_BUFFER = io.StringIO() +LOG_HANDLER = logging.StreamHandler(LOG_BUFFER) +TQDM_UPDATE = tqdm.update + + +def render() -> None: + global LOG_LEVEL_DROPDOWN + global TERMINAL_TEXTBOX + + LOG_LEVEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.log_level_dropdown'), + choices = facefusion.choices.log_levels, + value = state_manager.get_item('log_level') + ) + TERMINAL_TEXTBOX = gradio.Textbox( + label = translator.get('uis.terminal_textbox'), + value = read_logs, + lines = 8, + max_lines = 8, + every = 0.5, + show_copy_button = True + ) + + +def listen() -> None: + LOG_LEVEL_DROPDOWN.change(update_log_level, inputs = LOG_LEVEL_DROPDOWN) + logger.get_package_logger().addHandler(LOG_HANDLER) + tqdm.update = tqdm_update + + +def update_log_level(log_level : LogLevel) -> None: + state_manager.set_item('log_level', log_level) + logger.init(state_manager.get_item('log_level')) + + +def tqdm_update(self : tqdm, n : int = 1) -> None: + TQDM_UPDATE(self, n) + output = create_tqdm_output(self) + + if output: + LOG_BUFFER.seek(0) + log_buffer = LOG_BUFFER.read() + lines = log_buffer.splitlines() + if lines and lines[-1].startswith(self.desc): + position = log_buffer.rfind(lines[-1]) + LOG_BUFFER.seek(position) + else: + LOG_BUFFER.seek(0, os.SEEK_END) + LOG_BUFFER.write(output + os.linesep) + LOG_BUFFER.flush() + + +def create_tqdm_output(self : tqdm) -> Optional[str]: + if not self.disable and self.desc and self.total: + percentage = math.floor(self.n / self.total * 100) + return self.desc + translator.get('colon') + ' ' + str(percentage) + '% (' + str(self.n) + '/' + str(self.total) + ')' + if not self.disable and self.desc and self.unit: + return self.desc + translator.get('colon') + ' ' + str(self.n) + ' ' + self.unit + return None + + +def read_logs() -> str: + LOG_BUFFER.seek(0) + logs = LOG_BUFFER.read().strip() + return logs diff --git a/facefusion/uis/components/trim_frame.py b/facefusion/uis/components/trim_frame.py new file mode 100644 index 0000000000000000000000000000000000000000..fc1797a815a5fd380661379cb1b514bd54d72675 --- /dev/null +++ b/facefusion/uis/components/trim_frame.py @@ -0,0 +1,62 @@ +from typing import Optional, Tuple + +from gradio_rangeslider import RangeSlider + +from facefusion import state_manager, translator +from facefusion.face_store import clear_static_faces +from facefusion.filesystem import is_video +from facefusion.uis.core import get_ui_components +from facefusion.uis.types import ComponentOptions +from facefusion.vision import count_video_frame_total + +TRIM_FRAME_RANGE_SLIDER : Optional[RangeSlider] = None + + +def render() -> None: + global TRIM_FRAME_RANGE_SLIDER + + trim_frame_range_slider_options : ComponentOptions =\ + { + 'label': translator.get('uis.trim_frame_slider'), + 'minimum': 0, + 'step': 1, + 'visible': False + } + if is_video(state_manager.get_item('target_path')): + video_frame_total = count_video_frame_total(state_manager.get_item('target_path')) + trim_frame_start = state_manager.get_item('trim_frame_start') or 0 + trim_frame_end = state_manager.get_item('trim_frame_end') or video_frame_total + trim_frame_range_slider_options['maximum'] = video_frame_total + trim_frame_range_slider_options['value'] = (trim_frame_start, trim_frame_end) + trim_frame_range_slider_options['visible'] = True + TRIM_FRAME_RANGE_SLIDER = RangeSlider(**trim_frame_range_slider_options) + + +def listen() -> None: + TRIM_FRAME_RANGE_SLIDER.release(update_trim_frame, inputs = TRIM_FRAME_RANGE_SLIDER) + for ui_component in get_ui_components( + [ + 'target_image', + 'target_video' + ]): + for method in [ 'change', 'clear' ]: + getattr(ui_component, method)(remote_update, outputs = [ TRIM_FRAME_RANGE_SLIDER ]) + + +def remote_update() -> RangeSlider: + if is_video(state_manager.get_item('target_path')): + video_frame_total = count_video_frame_total(state_manager.get_item('target_path')) + state_manager.clear_item('trim_frame_start') + state_manager.clear_item('trim_frame_end') + return RangeSlider(value = (0, video_frame_total), maximum = video_frame_total, visible = True) + return RangeSlider(visible = False) + + +def update_trim_frame(trim_frame : Tuple[float, float]) -> None: + clear_static_faces() + trim_frame_start, trim_frame_end = trim_frame + video_frame_total = count_video_frame_total(state_manager.get_item('target_path')) + trim_frame_start = int(trim_frame_start) if trim_frame_start > 0 else None + trim_frame_end = int(trim_frame_end) if trim_frame_end < video_frame_total else None + state_manager.set_item('trim_frame_start', trim_frame_start) + state_manager.set_item('trim_frame_end', trim_frame_end) diff --git a/facefusion/uis/components/ui_workflow.py b/facefusion/uis/components/ui_workflow.py new file mode 100644 index 0000000000000000000000000000000000000000..a1e3a699b0762da53d1978a35d6e0ac5b8b00725 --- /dev/null +++ b/facefusion/uis/components/ui_workflow.py @@ -0,0 +1,21 @@ +from typing import Optional + +import gradio + +import facefusion +from facefusion import state_manager, translator +from facefusion.uis.core import register_ui_component + +UI_WORKFLOW_DROPDOWN : Optional[gradio.Dropdown] = None + + +def render() -> None: + global UI_WORKFLOW_DROPDOWN + + UI_WORKFLOW_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.ui_workflow'), + choices = facefusion.choices.ui_workflows, + value = state_manager.get_item('ui_workflow'), + interactive = True + ) + register_ui_component('ui_workflow_dropdown', UI_WORKFLOW_DROPDOWN) diff --git a/facefusion/uis/components/voice_extractor.py b/facefusion/uis/components/voice_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..b071b2c7ff68cd92c4a0de5cf8ea2c62ca6bf609 --- /dev/null +++ b/facefusion/uis/components/voice_extractor.py @@ -0,0 +1,50 @@ +from typing import Optional + +import gradio + +import facefusion.choices +from facefusion import state_manager, translator, voice_extractor +from facefusion.filesystem import is_video +from facefusion.types import VoiceExtractorModel +from facefusion.uis.core import get_ui_components, register_ui_component + +VOICE_EXTRACTOR_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None + + +def render() -> None: + global VOICE_EXTRACTOR_MODEL_DROPDOWN + + VOICE_EXTRACTOR_MODEL_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.voice_extractor_model_dropdown'), + choices = facefusion.choices.voice_extractor_models, + value = state_manager.get_item('voice_extractor_model'), + visible = is_video(state_manager.get_item('target_path')) + ) + register_ui_component('voice_extractor_model_dropdown', VOICE_EXTRACTOR_MODEL_DROPDOWN) + + +def listen() -> None: + VOICE_EXTRACTOR_MODEL_DROPDOWN.change(update_voice_extractor_model, inputs = VOICE_EXTRACTOR_MODEL_DROPDOWN, outputs = VOICE_EXTRACTOR_MODEL_DROPDOWN) + + for ui_component in get_ui_components( + [ + 'target_image', + 'target_video' + ]): + for method in [ 'change', 'clear' ]: + getattr(ui_component, method)(remote_update, outputs = VOICE_EXTRACTOR_MODEL_DROPDOWN) + + +def remote_update() -> gradio.Dropdown: + if is_video(state_manager.get_item('target_path')): + return gradio.Dropdown(visible = True) + return gradio.Dropdown(visible = False) + + +def update_voice_extractor_model(voice_extractor_model : VoiceExtractorModel) -> gradio.Dropdown: + voice_extractor.clear_inference_pool() + state_manager.set_item('voice_extractor_model', voice_extractor_model) + + if voice_extractor.pre_check(): + gradio.Dropdown(value = state_manager.get_item('voice_extractor_model')) + return gradio.Dropdown() diff --git a/facefusion/uis/components/webcam.py b/facefusion/uis/components/webcam.py new file mode 100644 index 0000000000000000000000000000000000000000..5e9294034c72668a6902b707ae9bbb9f132ac3dd --- /dev/null +++ b/facefusion/uis/components/webcam.py @@ -0,0 +1,115 @@ +from typing import Iterator, List, Optional, Tuple + +import cv2 +import gradio + +from facefusion import state_manager, translator +from facefusion.camera_manager import clear_camera_pool, get_local_camera_capture +from facefusion.filesystem import has_image +from facefusion.streamer import multi_process_capture, open_stream +from facefusion.types import Fps, VisionFrame, WebcamMode +from facefusion.uis.core import get_ui_component +from facefusion.uis.types import File +from facefusion.vision import unpack_resolution + +SOURCE_FILE : Optional[gradio.File] = None +WEBCAM_IMAGE : Optional[gradio.Image] = None +WEBCAM_START_BUTTON : Optional[gradio.Button] = None +WEBCAM_STOP_BUTTON : Optional[gradio.Button] = None + + +def render() -> None: + global SOURCE_FILE + global WEBCAM_IMAGE + global WEBCAM_START_BUTTON + global WEBCAM_STOP_BUTTON + + has_source_image = has_image(state_manager.get_item('source_paths')) + SOURCE_FILE = gradio.File( + label = translator.get('uis.source_file'), + file_count = 'multiple', + value = state_manager.get_item('source_paths') if has_source_image else None + ) + WEBCAM_IMAGE = gradio.Image( + label = translator.get('uis.webcam_image'), + format = 'jpeg', + visible = False + ) + WEBCAM_START_BUTTON = gradio.Button( + value = translator.get('uis.start_button'), + variant = 'primary', + size = 'sm' + ) + WEBCAM_STOP_BUTTON = gradio.Button( + value = translator.get('uis.stop_button'), + size = 'sm', + visible = False + ) + + +def listen() -> None: + SOURCE_FILE.change(update_source, inputs = SOURCE_FILE, outputs = SOURCE_FILE) + webcam_device_id_dropdown = get_ui_component('webcam_device_id_dropdown') + webcam_mode_radio = get_ui_component('webcam_mode_radio') + webcam_resolution_dropdown = get_ui_component('webcam_resolution_dropdown') + webcam_fps_slider = get_ui_component('webcam_fps_slider') + + if webcam_device_id_dropdown and webcam_mode_radio and webcam_resolution_dropdown and webcam_fps_slider: + WEBCAM_START_BUTTON.click(pre_start, outputs = [ SOURCE_FILE, WEBCAM_IMAGE, WEBCAM_START_BUTTON, WEBCAM_STOP_BUTTON ]) + start_event = WEBCAM_START_BUTTON.click(start, inputs = [ webcam_device_id_dropdown, webcam_mode_radio, webcam_resolution_dropdown, webcam_fps_slider ], outputs = WEBCAM_IMAGE) + start_event.then(pre_stop) + WEBCAM_STOP_BUTTON.click(stop, cancels = start_event, outputs = WEBCAM_IMAGE) + WEBCAM_STOP_BUTTON.click(pre_stop, outputs = [ SOURCE_FILE, WEBCAM_IMAGE, WEBCAM_START_BUTTON, WEBCAM_STOP_BUTTON ]) + + +def update_source(files : List[File]) -> gradio.File: + file_names = [ file.name for file in files ] if files else None + has_source_image = has_image(file_names) + + if has_source_image: + state_manager.set_item('source_paths', file_names) + return gradio.File(value = file_names) + + state_manager.clear_item('source_paths') + return gradio.File(value = None) + + +def pre_start() -> Tuple[gradio.File, gradio.Image, gradio.Button, gradio.Button]: + return gradio.File(visible = False), gradio.Image(visible = True), gradio.Button(visible = False), gradio.Button(visible = True) + + +def pre_stop() -> Tuple[gradio.File, gradio.Image, gradio.Button, gradio.Button]: + return gradio.File(visible = True), gradio.Image(visible = False), gradio.Button(visible = True), gradio.Button(visible = False) + + +def start(webcam_device_id : int, webcam_mode : WebcamMode, webcam_resolution : str, webcam_fps : Fps) -> Iterator[VisionFrame]: + state_manager.init_item('face_selector_mode', 'one') + state_manager.sync_state() + + camera_capture = get_local_camera_capture(webcam_device_id) + stream = None + + if webcam_mode in [ 'udp', 'v4l2' ]: + stream = open_stream(webcam_mode, webcam_resolution, webcam_fps) # type:ignore[arg-type] + webcam_width, webcam_height = unpack_resolution(webcam_resolution) + + if camera_capture and camera_capture.isOpened(): + camera_capture.set(cv2.CAP_PROP_FRAME_WIDTH, webcam_width) + camera_capture.set(cv2.CAP_PROP_FRAME_HEIGHT, webcam_height) + camera_capture.set(cv2.CAP_PROP_FPS, webcam_fps) + + for capture_frame in multi_process_capture(camera_capture, webcam_fps): + capture_frame = cv2.cvtColor(capture_frame, cv2.COLOR_BGR2RGB) + + if webcam_mode == 'inline': + yield capture_frame + else: + try: + stream.stdin.write(capture_frame.tobytes()) + except Exception: + pass + + +def stop() -> gradio.Image: + clear_camera_pool() + return gradio.Image(value = None) diff --git a/facefusion/uis/components/webcam_options.py b/facefusion/uis/components/webcam_options.py new file mode 100644 index 0000000000000000000000000000000000000000..a5e604cdc83a78595643e503fdf27acdb5f0a32d --- /dev/null +++ b/facefusion/uis/components/webcam_options.py @@ -0,0 +1,49 @@ +from typing import Optional + +import gradio + +from facefusion import translator +from facefusion.camera_manager import detect_local_camera_ids +from facefusion.common_helper import get_first +from facefusion.uis import choices as uis_choices +from facefusion.uis.core import register_ui_component + +WEBCAM_DEVICE_ID_DROPDOWN : Optional[gradio.Dropdown] = None +WEBCAM_MODE_RADIO : Optional[gradio.Radio] = None +WEBCAM_RESOLUTION_DROPDOWN : Optional[gradio.Dropdown] = None +WEBCAM_FPS_SLIDER : Optional[gradio.Slider] = None + + +def render() -> None: + global WEBCAM_DEVICE_ID_DROPDOWN + global WEBCAM_MODE_RADIO + global WEBCAM_RESOLUTION_DROPDOWN + global WEBCAM_FPS_SLIDER + + local_camera_ids = detect_local_camera_ids(0, 10) or [ 'none' ] #type:ignore[list-item] + WEBCAM_DEVICE_ID_DROPDOWN = gradio.Dropdown( + value = get_first(local_camera_ids), + label = translator.get('uis.webcam_device_id_dropdown'), + choices = local_camera_ids + ) + WEBCAM_MODE_RADIO = gradio.Radio( + label = translator.get('uis.webcam_mode_radio'), + choices = uis_choices.webcam_modes, + value = uis_choices.webcam_modes[0] + ) + WEBCAM_RESOLUTION_DROPDOWN = gradio.Dropdown( + label = translator.get('uis.webcam_resolution_dropdown'), + choices = uis_choices.webcam_resolutions, + value = uis_choices.webcam_resolutions[0] + ) + WEBCAM_FPS_SLIDER = gradio.Slider( + label = translator.get('uis.webcam_fps_slider'), + value = 30, + step = 1, + minimum = 1, + maximum = 30 + ) + register_ui_component('webcam_device_id_dropdown', WEBCAM_DEVICE_ID_DROPDOWN) + register_ui_component('webcam_mode_radio', WEBCAM_MODE_RADIO) + register_ui_component('webcam_resolution_dropdown', WEBCAM_RESOLUTION_DROPDOWN) + register_ui_component('webcam_fps_slider', WEBCAM_FPS_SLIDER) diff --git a/facefusion/uis/core.py b/facefusion/uis/core.py new file mode 100644 index 0000000000000000000000000000000000000000..fc57737b9b2fe61c48d101effd637b6282060dde --- /dev/null +++ b/facefusion/uis/core.py @@ -0,0 +1,198 @@ +import importlib +import os +import warnings +from types import ModuleType +from typing import Any, Dict, List, Optional + +import gradio +from gradio.themes import Size + +import facefusion.uis.overrides as uis_overrides +from facefusion import logger, metadata, state_manager, translator +from facefusion.exit_helper import hard_exit +from facefusion.filesystem import resolve_relative_path +from facefusion.uis.types import Component, ComponentName + +UI_COMPONENTS: Dict[ComponentName, Component] = {} +UI_LAYOUT_MODULES : List[ModuleType] = [] +UI_LAYOUT_METHODS =\ +[ + 'pre_check', + 'render', + 'listen', + 'run' +] + + +def load_ui_layout_module(ui_layout : str) -> Any: + try: + ui_layout_module = importlib.import_module('facefusion.uis.layouts.' + ui_layout) + for method_name in UI_LAYOUT_METHODS: + if not hasattr(ui_layout_module, method_name): + raise NotImplementedError + except ModuleNotFoundError as exception: + logger.error(translator.get('ui_layout_not_loaded').format(ui_layout = ui_layout), __name__) + logger.debug(exception.msg, __name__) + hard_exit(1) + except NotImplementedError: + logger.error(translator.get('ui_layout_not_implemented').format(ui_layout = ui_layout), __name__) + hard_exit(1) + return ui_layout_module + + +def get_ui_layouts_modules(ui_layouts : List[str]) -> List[ModuleType]: + if not UI_LAYOUT_MODULES: + for ui_layout in ui_layouts: + ui_layout_module = load_ui_layout_module(ui_layout) + UI_LAYOUT_MODULES.append(ui_layout_module) + return UI_LAYOUT_MODULES + + +def get_ui_component(component_name : ComponentName) -> Optional[Component]: + if component_name in UI_COMPONENTS: + return UI_COMPONENTS[component_name] + return None + + +def get_ui_components(component_names : List[ComponentName]) -> Optional[List[Component]]: + ui_components = [] + + for component_name in component_names: + component = get_ui_component(component_name) + if component: + ui_components.append(component) + return ui_components + + +def register_ui_component(component_name : ComponentName, component: Component) -> None: + UI_COMPONENTS[component_name] = component + + +def init() -> None: + os.environ['GRADIO_ANALYTICS_ENABLED'] = '0' + os.environ['GRADIO_TEMP_DIR'] = os.path.join(state_manager.get_item('temp_path'), 'gradio') + + warnings.filterwarnings('ignore', category = UserWarning, module = 'gradio') + gradio.processing_utils._check_allowed = uis_overrides.mock + gradio.processing_utils.convert_video_to_playable_mp4 = uis_overrides.convert_video_to_playable_mp4 + gradio.components.Number.raise_if_out_of_bounds = uis_overrides.mock + + +def launch() -> None: + ui_layouts_total = len(state_manager.get_item('ui_layouts')) + with gradio.Blocks(theme = get_theme(), css = get_css(), title = metadata.get('name') + ' ' + metadata.get('version'), fill_width = True) as ui: + for ui_layout in state_manager.get_item('ui_layouts'): + ui_layout_module = load_ui_layout_module(ui_layout) + + if ui_layouts_total > 1: + with gradio.Tab(ui_layout): + ui_layout_module.render() + ui_layout_module.listen() + else: + ui_layout_module.render() + ui_layout_module.listen() + + for ui_layout in state_manager.get_item('ui_layouts'): + ui_layout_module = load_ui_layout_module(ui_layout) + ui_layout_module.run(ui) + + +def get_theme() -> gradio.Theme: + return gradio.themes.Base( + primary_hue = gradio.themes.colors.red, + secondary_hue = gradio.themes.Color( + name = 'neutral', + c50 = '#fafafa', + c100 = '#f5f5f5', + c200 = '#e5e5e5', + c300 = '#d4d4d4', + c400 = '#a3a3a3', + c500 = '#737373', + c600 = '#525252', + c700 = '#404040', + c800 = '#262626', + c900 = '#212121', + c950 = '#171717', + ), + radius_size = Size( + xxs = '0.375rem', + xs = '0.375rem', + sm = '0.375rem', + md = '0.375rem', + lg = '0.375rem', + xl = '0.375rem', + xxl = '0.375rem', + ), + font = gradio.themes.GoogleFont('Open Sans') + ).set( + color_accent = 'transparent', + color_accent_soft = 'transparent', + color_accent_soft_dark = 'transparent', + background_fill_primary = '*neutral_100', + background_fill_primary_dark = '*neutral_950', + background_fill_secondary = '*neutral_50', + background_fill_secondary_dark = '*neutral_800', + block_background_fill = 'white', + block_background_fill_dark = '*neutral_900', + block_border_width = '0', + block_label_background_fill = '*neutral_100', + block_label_background_fill_dark = '*neutral_800', + block_label_border_width = 'none', + block_label_margin = '0.5rem', + block_label_radius = '*radius_md', + block_label_text_color = '*neutral_700', + block_label_text_size = '*text_sm', + block_label_text_color_dark = 'white', + block_label_text_weight = '600', + block_title_background_fill = '*neutral_100', + block_title_background_fill_dark = '*neutral_800', + block_title_padding = '*block_label_padding', + block_title_radius = '*block_label_radius', + block_title_text_color = '*neutral_700', + block_title_text_size = '*text_sm', + block_title_text_weight = '600', + block_padding = '0.5rem', + border_color_accent = 'transparent', + border_color_accent_dark = 'transparent', + border_color_accent_subdued = 'transparent', + border_color_accent_subdued_dark = 'transparent', + border_color_primary = 'transparent', + border_color_primary_dark = 'transparent', + button_large_padding = '2rem 0.5rem', + button_large_text_weight = 'normal', + button_primary_background_fill = '*primary_500', + button_primary_background_fill_dark = '*primary_600', + button_primary_text_color = 'white', + button_secondary_background_fill = 'white', + button_secondary_background_fill_dark = '*neutral_800', + button_secondary_background_fill_hover = 'white', + button_secondary_background_fill_hover_dark = '*neutral_800', + button_secondary_text_color = '*neutral_800', + button_small_padding = '0.75rem', + button_small_text_size = '0.875rem', + checkbox_background_color = '*neutral_200', + checkbox_background_color_dark = '*neutral_900', + checkbox_background_color_selected = '*primary_600', + checkbox_background_color_selected_dark = '*primary_700', + checkbox_label_background_fill = '*neutral_50', + checkbox_label_background_fill_dark = '*neutral_800', + checkbox_label_background_fill_hover = '*neutral_50', + checkbox_label_background_fill_hover_dark = '*neutral_800', + checkbox_label_background_fill_selected = '*primary_500', + checkbox_label_background_fill_selected_dark = '*primary_600', + checkbox_label_text_color_selected = 'white', + error_background_fill = 'white', + error_background_fill_dark = '*neutral_900', + error_text_color = '*primary_500', + error_text_color_dark = '*primary_600', + input_background_fill = '*neutral_50', + input_background_fill_dark = '*neutral_800', + shadow_drop = 'none', + slider_color = '*primary_500', + slider_color_dark = '*primary_600' + ) + + +def get_css() -> str: + overrides_css_path = resolve_relative_path('uis/assets/overrides.css') + return open(overrides_css_path).read() diff --git a/facefusion/uis/layouts/benchmark.py b/facefusion/uis/layouts/benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..6689ba7b9494d2c64eb51a1d143894babcb31ef5 --- /dev/null +++ b/facefusion/uis/layouts/benchmark.py @@ -0,0 +1,79 @@ +import gradio + +from facefusion import benchmarker, state_manager +from facefusion.uis.components import about, age_modifier_options, background_remover_options, benchmark, benchmark_options, deep_swapper_options, download, execution, execution_thread_count, expression_restorer_options, face_debugger_options, face_editor_options, face_enhancer_options, face_swapper_options, frame_colorizer_options, frame_enhancer_options, lip_syncer_options, memory, processors + + +def pre_check() -> bool: + return benchmarker.pre_check() + + +def render() -> gradio.Blocks: + with gradio.Blocks() as layout: + with gradio.Row(): + with gradio.Column(scale = 4): + with gradio.Blocks(): + about.render() + with gradio.Blocks(): + benchmark_options.render() + with gradio.Blocks(): + processors.render() + with gradio.Blocks(): + age_modifier_options.render() + with gradio.Blocks(): + background_remover_options.render() + with gradio.Blocks(): + deep_swapper_options.render() + with gradio.Blocks(): + expression_restorer_options.render() + with gradio.Blocks(): + face_debugger_options.render() + with gradio.Blocks(): + face_editor_options.render() + with gradio.Blocks(): + face_enhancer_options.render() + with gradio.Blocks(): + face_swapper_options.render() + with gradio.Blocks(): + frame_colorizer_options.render() + with gradio.Blocks(): + frame_enhancer_options.render() + with gradio.Blocks(): + lip_syncer_options.render() + with gradio.Blocks(): + execution.render() + execution_thread_count.render() + with gradio.Blocks(): + download.render() + with gradio.Blocks(): + state_manager.set_item('video_memory_strategy', 'tolerant') + memory.render() + with gradio.Column(scale = 11): + with gradio.Blocks(): + benchmark.render() + return layout + + +def listen() -> None: + processors.listen() + age_modifier_options.listen() + background_remover_options.listen() + deep_swapper_options.listen() + expression_restorer_options.listen() + download.listen() + face_debugger_options.listen() + face_editor_options.listen() + face_enhancer_options.listen() + face_swapper_options.listen() + frame_colorizer_options.listen() + frame_enhancer_options.listen() + lip_syncer_options.listen() + execution.listen() + execution_thread_count.listen() + memory.listen() + benchmark.listen() + benchmark_options.listen() + + +def run(ui : gradio.Blocks) -> None: + ui.launch(favicon_path = 'facefusion.ico', inbrowser = state_manager.get_item('open_browser')) diff --git a/facefusion/uis/layouts/default.py b/facefusion/uis/layouts/default.py new file mode 100644 index 0000000000000000000000000000000000000000..7d6227a29a7156248d8fcadd98760c7cfc4d720d --- /dev/null +++ b/facefusion/uis/layouts/default.py @@ -0,0 +1,125 @@ +import gradio + +from facefusion import state_manager +from facefusion.uis.components import about, age_modifier_options, background_remover_options, common_options, deep_swapper_options, download, execution, execution_thread_count, expression_restorer_options, face_debugger_options, face_detector, face_editor_options, face_enhancer_options, face_landmarker, face_masker, face_selector, face_swapper_options, frame_colorizer_options, frame_enhancer_options, instant_runner, job_manager, job_runner, lip_syncer_options, memory, output, output_options, preview, preview_options, processors, source, target, temp_frame, terminal, trim_frame, ui_workflow, voice_extractor + + +def pre_check() -> bool: + return True + + +def render() -> gradio.Blocks: + with gradio.Blocks() as layout: + with gradio.Row(): + with gradio.Column(scale = 4): + with gradio.Blocks(): + about.render() + with gradio.Blocks(): + processors.render() + with gradio.Blocks(): + age_modifier_options.render() + with gradio.Blocks(): + background_remover_options.render() + with gradio.Blocks(): + deep_swapper_options.render() + with gradio.Blocks(): + expression_restorer_options.render() + with gradio.Blocks(): + face_debugger_options.render() + with gradio.Blocks(): + face_editor_options.render() + with gradio.Blocks(): + face_enhancer_options.render() + with gradio.Blocks(): + face_swapper_options.render() + with gradio.Blocks(): + frame_colorizer_options.render() + with gradio.Blocks(): + frame_enhancer_options.render() + with gradio.Blocks(): + lip_syncer_options.render() + with gradio.Blocks(): + voice_extractor.render() + with gradio.Blocks(): + execution.render() + execution_thread_count.render() + with gradio.Blocks(): + download.render() + with gradio.Blocks(): + memory.render() + with gradio.Blocks(): + temp_frame.render() + with gradio.Blocks(): + output_options.render() + with gradio.Column(scale = 4): + with gradio.Blocks(): + source.render() + with gradio.Blocks(): + target.render() + with gradio.Blocks(): + output.render() + with gradio.Blocks(): + terminal.render() + with gradio.Blocks(): + ui_workflow.render() + instant_runner.render() + job_runner.render() + job_manager.render() + with gradio.Column(scale = 7): + with gradio.Blocks(): + preview.render() + preview_options.render() + with gradio.Blocks(): + trim_frame.render() + with gradio.Blocks(): + face_selector.render() + with gradio.Blocks(): + face_masker.render() + with gradio.Blocks(): + face_detector.render() + with gradio.Blocks(): + face_landmarker.render() + with gradio.Blocks(): + common_options.render() + return layout + + +def listen() -> None: + processors.listen() + age_modifier_options.listen() + background_remover_options.listen() + deep_swapper_options.listen() + expression_restorer_options.listen() + face_debugger_options.listen() + face_editor_options.listen() + face_enhancer_options.listen() + face_swapper_options.listen() + frame_colorizer_options.listen() + frame_enhancer_options.listen() + lip_syncer_options.listen() + execution.listen() + execution_thread_count.listen() + download.listen() + memory.listen() + temp_frame.listen() + output_options.listen() + source.listen() + target.listen() + output.listen() + instant_runner.listen() + job_runner.listen() + job_manager.listen() + terminal.listen() + preview.listen() + preview_options.listen() + trim_frame.listen() + face_selector.listen() + face_masker.listen() + face_detector.listen() + face_landmarker.listen() + voice_extractor.listen() + common_options.listen() + + +def run(ui : gradio.Blocks) -> None: + ui.launch(favicon_path = 'facefusion.ico', inbrowser = state_manager.get_item('open_browser')) diff --git a/facefusion/uis/layouts/jobs.py b/facefusion/uis/layouts/jobs.py new file mode 100644 index 0000000000000000000000000000000000000000..a1410d449fb6351c627fdd68d5c46a45432227f7 --- /dev/null +++ b/facefusion/uis/layouts/jobs.py @@ -0,0 +1,31 @@ +import gradio + +from facefusion import state_manager +from facefusion.uis.components import about, job_list, job_list_options + + +def pre_check() -> bool: + return True + + +def render() -> gradio.Blocks: + with gradio.Blocks() as layout: + with gradio.Row(): + with gradio.Column(scale = 4): + with gradio.Blocks(): + about.render() + with gradio.Blocks(): + job_list_options.render() + with gradio.Column(scale = 11): + with gradio.Blocks(): + job_list.render() + return layout + + +def listen() -> None: + job_list_options.listen() + job_list.listen() + + +def run(ui : gradio.Blocks) -> None: + ui.launch(favicon_path = 'facefusion.ico', inbrowser = state_manager.get_item('open_browser')) diff --git a/facefusion/uis/layouts/webcam.py b/facefusion/uis/layouts/webcam.py new file mode 100644 index 0000000000000000000000000000000000000000..3c94651e7be4a8c9a8aee2d318ba1b7558f5647a --- /dev/null +++ b/facefusion/uis/layouts/webcam.py @@ -0,0 +1,74 @@ +import gradio + +from facefusion import state_manager +from facefusion.uis.components import about, age_modifier_options, background_remover_options, deep_swapper_options, download, execution, execution_thread_count, expression_restorer_options, face_debugger_options, face_editor_options, face_enhancer_options, face_swapper_options, frame_colorizer_options, frame_enhancer_options, lip_syncer_options, processors, webcam, webcam_options + + +def pre_check() -> bool: + return True + + +def render() -> gradio.Blocks: + with gradio.Blocks() as layout: + with gradio.Row(): + with gradio.Column(scale = 4): + with gradio.Blocks(): + about.render() + with gradio.Blocks(): + webcam_options.render() + with gradio.Blocks(): + processors.render() + with gradio.Blocks(): + age_modifier_options.render() + with gradio.Blocks(): + background_remover_options.render() + with gradio.Blocks(): + deep_swapper_options.render() + with gradio.Blocks(): + expression_restorer_options.render() + with gradio.Blocks(): + face_debugger_options.render() + with gradio.Blocks(): + face_editor_options.render() + with gradio.Blocks(): + face_enhancer_options.render() + with gradio.Blocks(): + face_swapper_options.render() + with gradio.Blocks(): + frame_colorizer_options.render() + with gradio.Blocks(): + frame_enhancer_options.render() + with gradio.Blocks(): + lip_syncer_options.render() + with gradio.Blocks(): + execution.render() + execution_thread_count.render() + with gradio.Blocks(): + download.render() + with gradio.Column(scale = 11): + with gradio.Blocks(): + webcam.render() + return layout + + +def listen() -> None: + processors.listen() + age_modifier_options.listen() + background_remover_options.listen() + deep_swapper_options.listen() + expression_restorer_options.listen() + download.listen() + face_debugger_options.listen() + face_editor_options.listen() + face_enhancer_options.listen() + face_swapper_options.listen() + frame_colorizer_options.listen() + frame_enhancer_options.listen() + lip_syncer_options.listen() + execution.listen() + execution_thread_count.listen() + webcam.listen() + + +def run(ui : gradio.Blocks) -> None: + ui.launch(favicon_path = 'facefusion.ico', inbrowser = state_manager.get_item('open_browser')) diff --git a/facefusion/uis/overrides.py b/facefusion/uis/overrides.py new file mode 100644 index 0000000000000000000000000000000000000000..7ed7db8139e9ecfdb115f8a54d0262508f9521a1 --- /dev/null +++ b/facefusion/uis/overrides.py @@ -0,0 +1,31 @@ +from facefusion import ffmpeg_builder +from facefusion.ffmpeg import run_ffmpeg +from facefusion.filesystem import get_file_size +from facefusion.temp_helper import create_temp_directory, get_temp_file_path +from facefusion.uis.types import MockArgs + + +def convert_video_to_playable_mp4(video_path : str) -> str: + video_file_size = get_file_size(video_path) + max_file_size = 512 * 1024 * 1024 + + create_temp_directory(video_path) + temp_video_path = get_temp_file_path(video_path) + commands = ffmpeg_builder.set_input(video_path) + + if video_file_size > max_file_size: + commands.extend(ffmpeg_builder.set_video_duration(10)) + + commands.extend(ffmpeg_builder.force_output(temp_video_path)) + + process = run_ffmpeg(commands) + process.communicate() + + if process.returncode == 0: + return temp_video_path + + return video_path + + +def mock(*args : MockArgs, **kwargs : MockArgs) -> None: + pass diff --git a/facefusion/uis/types.py b/facefusion/uis/types.py new file mode 100644 index 0000000000000000000000000000000000000000..59e164f048ee551fc335c940997c083dc47b121b --- /dev/null +++ b/facefusion/uis/types.py @@ -0,0 +1,99 @@ +from typing import Any, Dict, IO, Literal, TypeAlias + +File : TypeAlias = IO[Any] +ComponentName = Literal\ +[ + 'age_modifier_direction_slider', + 'age_modifier_model_dropdown', + 'background_remover_model_dropdown', + 'background_remover_color_red_number', + 'background_remover_color_green_number', + 'background_remover_color_blue_number', + 'background_remover_color_alpha_number', + 'deep_swapper_model_dropdown', + 'deep_swapper_morph_slider', + 'expression_restorer_factor_slider', + 'expression_restorer_model_dropdown', + 'expression_restorer_areas_checkbox_group', + 'face_debugger_items_checkbox_group', + 'face_detector_angles_checkbox_group', + 'face_detector_model_dropdown', + 'face_detector_margin_slider', + 'face_detector_score_slider', + 'face_detector_size_dropdown', + 'face_editor_eyebrow_direction_slider', + 'face_editor_eye_gaze_horizontal_slider', + 'face_editor_eye_gaze_vertical_slider', + 'face_editor_eye_open_ratio_slider', + 'face_editor_head_pitch_slider', + 'face_editor_head_roll_slider', + 'face_editor_head_yaw_slider', + 'face_editor_lip_open_ratio_slider', + 'face_editor_model_dropdown', + 'face_editor_mouth_grim_slider', + 'face_editor_mouth_position_horizontal_slider', + 'face_editor_mouth_position_vertical_slider', + 'face_editor_mouth_pout_slider', + 'face_editor_mouth_purse_slider', + 'face_editor_mouth_smile_slider', + 'face_enhancer_blend_slider', + 'face_enhancer_model_dropdown', + 'face_enhancer_weight_slider', + 'face_landmarker_model_dropdown', + 'face_landmarker_score_slider', + 'face_mask_types_checkbox_group', + 'face_mask_areas_checkbox_group', + 'face_mask_regions_checkbox_group', + 'face_mask_blur_slider', + 'face_mask_padding_bottom_slider', + 'face_mask_padding_left_slider', + 'face_mask_padding_right_slider', + 'face_mask_padding_top_slider', + 'face_selector_age_range_slider', + 'face_selector_gender_dropdown', + 'face_selector_mode_dropdown', + 'face_selector_order_dropdown', + 'face_selector_race_dropdown', + 'face_swapper_model_dropdown', + 'face_swapper_pixel_boost_dropdown', + 'face_swapper_weight_slider', + 'face_occluder_model_dropdown', + 'face_parser_model_dropdown', + 'voice_extractor_model_dropdown', + 'frame_colorizer_blend_slider', + 'frame_colorizer_model_dropdown', + 'frame_colorizer_size_dropdown', + 'frame_enhancer_blend_slider', + 'frame_enhancer_model_dropdown', + 'job_list_job_status_checkbox_group', + 'lip_syncer_model_dropdown', + 'lip_syncer_weight_slider', + 'output_image', + 'output_video', + 'output_video_fps_slider', + 'preview_image', + 'preview_frame_slider', + 'preview_mode_dropdown', + 'preview_resolution_dropdown', + 'processors_checkbox_group', + 'reference_face_distance_slider', + 'reference_face_position_gallery', + 'source_audio', + 'source_image', + 'target_image', + 'target_video', + 'ui_workflow_dropdown', + 'webcam_device_id_dropdown', + 'webcam_fps_slider', + 'webcam_mode_radio', + 'webcam_resolution_dropdown' +] +Component : TypeAlias = Any +ComponentOptions : TypeAlias = Dict[str, Any] + +JobManagerAction = Literal['job-create', 'job-submit', 'job-delete', 'job-add-step', 'job-remix-step', 'job-insert-step', 'job-remove-step'] +JobRunnerAction = Literal['job-run', 'job-run-all', 'job-retry', 'job-retry-all'] + +PreviewMode = Literal[ 'default', 'frame-by-frame', 'face-by-face' ] + +MockArgs : TypeAlias = Any diff --git a/facefusion/uis/ui_helper.py b/facefusion/uis/ui_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..75327175c673fd144678ec13754b2b06c939423f --- /dev/null +++ b/facefusion/uis/ui_helper.py @@ -0,0 +1,26 @@ +import hashlib +import os +from typing import Optional + +from facefusion import state_manager +from facefusion.filesystem import get_file_extension, is_image, is_video + + +def convert_int_none(value : int) -> Optional[int]: + if value == 'none': + return None + return value + + +def convert_str_none(value : str) -> Optional[str]: + if value == 'none': + return None + return value + + +def suggest_output_path(output_directory_path : str, target_path : str) -> Optional[str]: + if is_image(target_path) or is_video(target_path): + output_file_name = hashlib.sha1(str(state_manager.get_state()).encode()).hexdigest()[:8] + target_file_extension = get_file_extension(target_path) + return os.path.join(output_directory_path, output_file_name + target_file_extension) + return None diff --git a/facefusion/video_manager.py b/facefusion/video_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..e4dd9e8de93e8c341c39f1eadc10357af423c28f --- /dev/null +++ b/facefusion/video_manager.py @@ -0,0 +1,40 @@ +import cv2 + +from facefusion.types import VideoPoolSet + +VIDEO_POOL_SET : VideoPoolSet =\ +{ + 'capture': {}, + '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) + + if video_capture.isOpened(): + VIDEO_POOL_SET['capture'][video_path] = video_capture + + return VIDEO_POOL_SET.get('capture').get(video_path) + + +def get_video_writer(video_path : str) -> cv2.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 + + return VIDEO_POOL_SET.get('writer').get(video_path) + + +def clear_video_pool() -> None: + for video_capture in VIDEO_POOL_SET.get('capture').values(): + video_capture.release() + + for video_writer in VIDEO_POOL_SET.get('writer').values(): + video_writer.release() + + VIDEO_POOL_SET['capture'].clear() + VIDEO_POOL_SET['writer'].clear() diff --git a/facefusion/vision.py b/facefusion/vision.py new file mode 100644 index 0000000000000000000000000000000000000000..266d7fd2d9c9aba8dc42a8e08d09501251bf3a40 --- /dev/null +++ b/facefusion/vision.py @@ -0,0 +1,365 @@ +import math +from functools import lru_cache +from typing import List, Optional, Tuple + +import cv2 +import numpy +from cv2.typing import Size + +from facefusion.common_helper import is_windows +from facefusion.filesystem import get_file_extension, is_image, is_video +from facefusion.thread_helper import thread_semaphore +from facefusion.types import ColorMode, Duration, Fps, Mask, Orientation, Resolution, Scale, VisionFrame +from facefusion.video_manager import get_video_capture + + +def read_static_images(image_paths : List[str], color_mode : ColorMode = 'rgb') -> List[VisionFrame]: + vision_frames = [] + + if image_paths: + for image_path in image_paths: + vision_frames.append(read_static_image(image_path, color_mode)) + return vision_frames + + +@lru_cache(maxsize = 64) +def read_static_image(image_path : str, color_mode : ColorMode = 'rgb') -> Optional[VisionFrame]: + return read_image(image_path, color_mode) + + +def read_image(image_path : str, color_mode : ColorMode = 'rgb') -> Optional[VisionFrame]: + if is_image(image_path): + flag = cv2.IMREAD_COLOR + + if color_mode == 'rgba': + flag = cv2.IMREAD_UNCHANGED + + if is_windows(): + image_buffer = numpy.fromfile(image_path, dtype = numpy.uint8) + return cv2.imdecode(image_buffer, flag) + return cv2.imread(image_path, flag) + return None + + +def write_image(image_path : str, vision_frame : VisionFrame) -> bool: + if image_path: + if is_windows(): + image_file_extension = get_file_extension(image_path) + _, vision_frame = cv2.imencode(image_file_extension, vision_frame) + vision_frame.tofile(image_path) + return is_image(image_path) + return cv2.imwrite(image_path, vision_frame) + return False + + +def detect_image_resolution(image_path : str) -> Optional[Resolution]: + if is_image(image_path): + image = read_image(image_path) + height, width = image.shape[:2] + + if width > 0 and height > 0: + return width, height + return None + + +def restrict_image_resolution(image_path : str, resolution : Resolution) -> Resolution: + if is_image(image_path): + image_resolution = detect_image_resolution(image_path) + if image_resolution < resolution: + return image_resolution + return resolution + + +@lru_cache(maxsize = 64) +def read_static_video_frame(video_path : str, frame_number : int = 0) -> Optional[VisionFrame]: + return read_video_frame(video_path, frame_number) + + +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) + + if video_capture and video_capture.isOpened(): + frame_total = video_capture.get(cv2.CAP_PROP_FRAME_COUNT) + + with thread_semaphore(): + video_capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1)) + has_vision_frame, vision_frame = video_capture.read() + + if has_vision_frame: + return vision_frame + + return None + + +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 0 + + +def predict_video_frame_total(video_path : str, fps : Fps, trim_frame_start : int, trim_frame_end : int) -> int: + if is_video(video_path): + video_fps = detect_video_fps(video_path) + extract_frame_total = count_trim_frame_total(video_path, trim_frame_start, trim_frame_end) * fps / video_fps + return math.floor(extract_frame_total) + return 0 + + +def detect_video_fps(video_path : str) -> Optional[float]: + 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 None + + +def restrict_video_fps(video_path : str, fps : Fps) -> Fps: + if is_video(video_path): + video_fps = detect_video_fps(video_path) + if video_fps < fps: + return video_fps + return fps + + +def detect_video_duration(video_path : str) -> Duration: + video_frame_total = count_video_frame_total(video_path) + video_fps = detect_video_fps(video_path) + + if video_frame_total and video_fps: + return video_frame_total / video_fps + return 0 + + +def count_trim_frame_total(video_path : str, trim_frame_start : Optional[int], trim_frame_end : Optional[int]) -> int: + trim_frame_start, trim_frame_end = restrict_trim_frame(video_path, trim_frame_start, trim_frame_end) + + return trim_frame_end - trim_frame_start + + +def restrict_trim_frame(video_path : str, trim_frame_start : Optional[int], trim_frame_end : Optional[int]) -> Tuple[int, int]: + video_frame_total = count_video_frame_total(video_path) + + if isinstance(trim_frame_start, int): + trim_frame_start = max(0, min(trim_frame_start, video_frame_total)) + if isinstance(trim_frame_end, int): + trim_frame_end = max(0, min(trim_frame_end, video_frame_total)) + + if isinstance(trim_frame_start, int) and isinstance(trim_frame_end, int): + return trim_frame_start, trim_frame_end + if isinstance(trim_frame_start, int): + return trim_frame_start, video_frame_total + if isinstance(trim_frame_end, int): + return 0, trim_frame_end + + return 0, video_frame_total + + +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 None + + +def restrict_video_resolution(video_path : str, resolution : Resolution) -> Resolution: + if is_video(video_path): + video_resolution = detect_video_resolution(video_path) + if video_resolution < resolution: + return video_resolution + return resolution + + +def scale_resolution(resolution : Resolution, scale : Scale) -> Resolution: + resolution = (int(resolution[0] * scale), int(resolution[1] * scale)) + resolution = normalize_resolution(resolution) + return resolution + + +def normalize_resolution(resolution : Tuple[float, float]) -> Resolution: + width, height = resolution + + if width > 0 and height > 0: + normalize_width = round(width / 2) * 2 + normalize_height = round(height / 2) * 2 + return normalize_width, normalize_height + return 0, 0 + + +def pack_resolution(resolution : Resolution) -> str: + width, height = normalize_resolution(resolution) + return str(width) + 'x' + str(height) + + +def unpack_resolution(resolution : str) -> Resolution: + width, height = map(int, resolution.split('x')) + return width, height + + +def detect_frame_orientation(vision_frame : VisionFrame) -> Orientation: + height, width = vision_frame.shape[:2] + + if width > height: + return 'landscape' + return 'portrait' + + +def restrict_frame(vision_frame : VisionFrame, resolution : Resolution) -> VisionFrame: + height, width = vision_frame.shape[:2] + restrict_width, restrict_height = resolution + + if height > restrict_height or width > restrict_width: + scale = min(restrict_height / height, restrict_width / width) + new_width = int(width * scale) + new_height = int(height * scale) + return cv2.resize(vision_frame, (new_width, new_height)) + return vision_frame + + +def fit_contain_frame(vision_frame : VisionFrame, resolution : Resolution) -> VisionFrame: + contain_width, contain_height = resolution + height, width = vision_frame.shape[:2] + scale = min(contain_height / height, contain_width / width) + new_width = int(width * scale) + new_height = int(height * scale) + start_x = max(0, (contain_width - new_width) // 2) + start_y = max(0, (contain_height - new_height) // 2) + end_x = max(0, contain_width - new_width - start_x) + end_y = max(0, contain_height - new_height - start_y) + temp_vision_frame = cv2.resize(vision_frame, (new_width, new_height)) + temp_vision_frame = numpy.pad(temp_vision_frame, ((start_y, end_y), (start_x, end_x), (0, 0))) + return temp_vision_frame + + +def fit_cover_frame(vision_frame : VisionFrame, resolution : Resolution) -> VisionFrame: + cover_width, cover_height = resolution + height, width = vision_frame.shape[:2] + scale = max(cover_width / width, cover_height / height) + new_width = int(width * scale) + new_height = int(height * scale) + start_x = max(0, (new_width - cover_width) // 2) + start_y = max(0, (new_height - cover_height) // 2) + end_x = min(new_width, start_x + cover_width) + end_y = min(new_height, start_y + cover_height) + temp_vision_frame = cv2.resize(vision_frame, (new_width, new_height)) + temp_vision_frame = temp_vision_frame[start_y:end_y, start_x:end_x] + return temp_vision_frame + + +def obscure_frame(vision_frame : VisionFrame) -> VisionFrame: + return cv2.GaussianBlur(vision_frame, (99, 99), 0) + + +def blend_frame(source_vision_frame : VisionFrame, target_vision_frame : VisionFrame, blend_factor : float) -> VisionFrame: + blend_vision_frame = cv2.addWeighted(source_vision_frame, 1 - blend_factor, target_vision_frame, blend_factor, 0) + return blend_vision_frame + + +def conditional_match_frame_color(source_vision_frame : VisionFrame, target_vision_frame : VisionFrame) -> VisionFrame: + histogram_factor = calculate_histogram_difference(source_vision_frame, target_vision_frame) + target_vision_frame = blend_frame(target_vision_frame, match_frame_color(source_vision_frame, target_vision_frame), histogram_factor) + return target_vision_frame + + +def match_frame_color(source_vision_frame : VisionFrame, target_vision_frame : VisionFrame) -> VisionFrame: + color_difference_sizes = numpy.linspace(16, target_vision_frame.shape[0], 3, endpoint = False) + + for color_difference_size in color_difference_sizes: + source_vision_frame = equalize_frame_color(source_vision_frame, target_vision_frame, normalize_resolution((color_difference_size, color_difference_size))) + target_vision_frame = equalize_frame_color(source_vision_frame, target_vision_frame, target_vision_frame.shape[:2][::-1]) + return target_vision_frame + + +def equalize_frame_color(source_vision_frame : VisionFrame, target_vision_frame : VisionFrame, size : Size) -> VisionFrame: + source_frame_resize = cv2.resize(source_vision_frame, size, interpolation = cv2.INTER_AREA).astype(numpy.float32) + target_frame_resize = cv2.resize(target_vision_frame, size, interpolation = cv2.INTER_AREA).astype(numpy.float32) + color_difference_vision_frame = numpy.subtract(source_frame_resize, target_frame_resize) + color_difference_vision_frame = cv2.resize(color_difference_vision_frame, target_vision_frame.shape[:2][::-1], interpolation = cv2.INTER_CUBIC) + target_vision_frame = numpy.add(target_vision_frame, color_difference_vision_frame).clip(0, 255).astype(numpy.uint8) + return target_vision_frame + + +def calculate_histogram_difference(source_vision_frame : VisionFrame, target_vision_frame : VisionFrame) -> float: + histogram_source = cv2.calcHist([cv2.cvtColor(source_vision_frame, cv2.COLOR_BGR2HSV)], [ 0, 1 ], None, [ 50, 60 ], [ 0, 180, 0, 256 ]) + histogram_target = cv2.calcHist([cv2.cvtColor(target_vision_frame, cv2.COLOR_BGR2HSV)], [ 0, 1 ], None, [ 50, 60 ], [ 0, 180, 0, 256 ]) + histogram_difference = float(numpy.interp(cv2.compareHist(histogram_source, histogram_target, cv2.HISTCMP_CORREL), [ -1, 1 ], [ 0, 1 ])) + return histogram_difference + + +def blend_vision_frames(source_vision_frame : VisionFrame, target_vision_frame : VisionFrame, blend_factor : float) -> VisionFrame: + blend_vision_frame = cv2.addWeighted(source_vision_frame, 1 - blend_factor, target_vision_frame, blend_factor, 0) + return blend_vision_frame + + +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] + pad_size_bottom = pad_size_top + tile_width - (vision_frame.shape[0] + 2 * size[1]) % tile_width + pad_size_right = pad_size_top + tile_width - (vision_frame.shape[1] + 2 * size[1]) % tile_width + pad_vision_frame = numpy.pad(vision_frame, ((pad_size_top, pad_size_bottom), (pad_size_top, pad_size_right), (0, 0))) + pad_height, pad_width = pad_vision_frame.shape[:2] + row_range = range(size[2], pad_height - size[2], tile_width) + col_range = range(size[2], pad_width - size[2], tile_width) + tile_vision_frames = [] + + for row_vision_frame in row_range: + top = row_vision_frame - size[2] + bottom = row_vision_frame + size[2] + tile_width + + for column_vision_frame in col_range: + left = column_vision_frame - size[2] + right = column_vision_frame + size[2] + tile_width + tile_vision_frames.append(pad_vision_frame[top:bottom, left:right, :]) + + return tile_vision_frames, pad_width, pad_height + + +def merge_tile_frames(tile_vision_frames : List[VisionFrame], temp_width : int, temp_height : int, pad_width : int, pad_height : int, size : Size) -> VisionFrame: + merge_vision_frame = numpy.zeros((pad_height, pad_width, 3)).astype(numpy.uint8) + tile_width = tile_vision_frames[0].shape[1] - 2 * size[2] + tiles_per_row = min(pad_width // tile_width, len(tile_vision_frames)) + + for index, tile_vision_frame in enumerate(tile_vision_frames): + tile_vision_frame = tile_vision_frame[size[2]:-size[2], size[2]:-size[2]] + row_index = index // tiles_per_row + col_index = index % tiles_per_row + top = row_index * tile_vision_frame.shape[0] + bottom = top + tile_vision_frame.shape[0] + left = col_index * tile_vision_frame.shape[1] + right = left + tile_vision_frame.shape[1] + merge_vision_frame[top:bottom, left:right, :] = tile_vision_frame + + merge_vision_frame = merge_vision_frame[size[1] : size[1] + temp_height, size[1]: size[1] + temp_width, :] + return merge_vision_frame + + +def extract_vision_mask(vision_frame : VisionFrame) -> Mask: + if vision_frame.ndim == 3 and vision_frame.shape[2] == 4: + return vision_frame[:, :, 3] + return numpy.full(vision_frame.shape[:2], 255, dtype = numpy.uint8) + + +def merge_vision_mask(vision_frame : VisionFrame, vision_mask : Mask) -> VisionFrame: + return numpy.dstack((vision_frame[:, :, :3], vision_mask)) + + +def conditional_merge_vision_mask(vision_frame : VisionFrame, vision_mask : Mask) -> VisionFrame: + if numpy.any(vision_mask < 255): + return merge_vision_mask(vision_frame, vision_mask) + return vision_frame diff --git a/facefusion/voice_extractor.py b/facefusion/voice_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..60f98000bc75fa11b70e011109ada1e56933f104 --- /dev/null +++ b/facefusion/voice_extractor.py @@ -0,0 +1,213 @@ +from functools import lru_cache +from typing import Tuple + +import numpy +import scipy + +from facefusion import inference_manager, state_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 thread_semaphore +from facefusion.types import Audio, AudioChunk, DownloadScope, DownloadSet, InferencePool, ModelSet, Voice, VoiceChunk + + +@lru_cache() +def create_static_model_set(download_scope : DownloadScope) -> ModelSet: + return\ + { + 'kim_vocal_1': + { + '__metadata__': + { + 'vendor': 'KimberleyJensen', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'voice_extractor': + { + 'url': resolve_download_url('models-3.4.0', 'kim_vocal_1.hash'), + 'path': resolve_relative_path('../.assets/models/kim_vocal_1.hash') + } + }, + 'sources': + { + 'voice_extractor': + { + 'url': resolve_download_url('models-3.4.0', 'kim_vocal_1.onnx'), + 'path': resolve_relative_path('../.assets/models/kim_vocal_1.onnx') + } + } + }, + 'kim_vocal_2': + { + '__metadata__': + { + 'vendor': 'KimberleyJensen', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'voice_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'kim_vocal_2.hash'), + 'path': resolve_relative_path('../.assets/models/kim_vocal_2.hash') + } + }, + 'sources': + { + 'voice_extractor': + { + 'url': resolve_download_url('models-3.0.0', 'kim_vocal_2.onnx'), + 'path': resolve_relative_path('../.assets/models/kim_vocal_2.onnx') + } + } + }, + 'uvr_mdxnet': + { + '__metadata__': + { + 'vendor': 'Unknown', + 'license': 'Non-Commercial', + 'year': 2023 + }, + 'hashes': + { + 'voice_extractor': + { + 'url': resolve_download_url('models-3.4.0', 'uvr_mdxnet.hash'), + 'path': resolve_relative_path('../.assets/models/uvr_mdxnet.hash') + } + }, + 'sources': + { + 'voice_extractor': + { + 'url': resolve_download_url('models-3.4.0', 'uvr_mdxnet.onnx'), + 'path': resolve_relative_path('../.assets/models/uvr_mdxnet.onnx') + } + } + } + } + + +def get_inference_pool() -> InferencePool: + model_names = [ state_manager.get_item('voice_extractor_model') ] + _, model_source_set = collect_model_downloads() + + return inference_manager.get_inference_pool(__name__, model_names, model_source_set) + + +def clear_inference_pool() -> None: + model_names = [ state_manager.get_item('voice_extractor_model') ] + inference_manager.clear_inference_pool(__name__, model_names) + + +def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_set = create_static_model_set('full') + model_hash_set = {} + model_source_set = {} + + for voice_extractor_model in [ 'kim_vocal_1', 'kim_vocal_2', 'uvr_mdxnet' ]: + if state_manager.get_item('voice_extractor_model') == voice_extractor_model: + model_hash_set[voice_extractor_model] = model_set.get(voice_extractor_model).get('hashes').get('voice_extractor') + model_source_set[voice_extractor_model] = model_set.get(voice_extractor_model).get('sources').get('voice_extractor') + + return model_hash_set, model_source_set + + +def pre_check() -> bool: + model_hash_set, model_source_set = collect_model_downloads() + + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) + + +def batch_extract_voice(audio : Audio, chunk_size : int, step_size : int) -> Voice: + temp_voice = numpy.zeros((audio.shape[0], 2)).astype(numpy.float32) + temp_voice_chunk = numpy.zeros((audio.shape[0], 2)).astype(numpy.float32) + + for start in range(0, audio.shape[0], step_size): + end = min(start + chunk_size, audio.shape[0]) + temp_voice[start:end, ...] += extract_voice(audio[start:end, ...]) + temp_voice_chunk[start:end, ...] += 1 + + voice = temp_voice / temp_voice_chunk + return voice + + +def extract_voice(temp_audio_chunk : AudioChunk) -> VoiceChunk: + voice_extractor = get_inference_pool().get(state_manager.get_item('voice_extractor_model')) + voice_trim_size = 3840 + voice_chunk_size = (voice_extractor.get_inputs()[0].shape[3] - 1) * 1024 + temp_audio_chunk, audio_pad_size = prepare_audio_chunk(temp_audio_chunk.T, voice_chunk_size, voice_trim_size) + temp_audio_chunk = decompose_audio_chunk(temp_audio_chunk, voice_trim_size) + temp_audio_chunk = forward(temp_audio_chunk) + temp_audio_chunk = compose_audio_chunk(temp_audio_chunk, voice_trim_size) + temp_audio_chunk = normalize_audio_chunk(temp_audio_chunk, voice_chunk_size, voice_trim_size, audio_pad_size) + return temp_audio_chunk + + +def forward(temp_audio_chunk : AudioChunk) -> AudioChunk: + voice_extractor = get_inference_pool().get(state_manager.get_item('voice_extractor_model')) + + with thread_semaphore(): + temp_audio_chunk = voice_extractor.run(None, + { + 'input': temp_audio_chunk + })[0] + + return temp_audio_chunk + + +def prepare_audio_chunk(temp_audio_chunk : AudioChunk, chunk_size : int, audio_trim_size : int) -> Tuple[AudioChunk, int]: + audio_step_size = chunk_size - 2 * audio_trim_size + audio_pad_size = audio_step_size - temp_audio_chunk.shape[1] % audio_step_size + audio_chunk_size = temp_audio_chunk.shape[1] + audio_pad_size + temp_audio_chunk = temp_audio_chunk.astype(numpy.float32) / numpy.iinfo(numpy.int16).max + temp_audio_chunk = numpy.pad(temp_audio_chunk, ((0, 0), (audio_trim_size, audio_trim_size + audio_pad_size))) + temp_audio_chunks = [] + + for index in range(0, audio_chunk_size, audio_step_size): + temp_audio_chunks.append(temp_audio_chunk[:, index:index + chunk_size]) + + temp_audio_chunk = numpy.concatenate(temp_audio_chunks, axis = 0) + temp_audio_chunk = temp_audio_chunk.reshape((-1, chunk_size)) + return temp_audio_chunk, audio_pad_size + + +def decompose_audio_chunk(temp_audio_chunk : AudioChunk, audio_trim_size : int) -> AudioChunk: + audio_frame_size = 7680 + audio_frame_overlap = 6656 + audio_frame_total = 3072 + audio_bin_total = 256 + audio_channel_total = 4 + window = scipy.signal.windows.hann(audio_frame_size) + temp_audio_chunk = scipy.signal.stft(temp_audio_chunk, nperseg = audio_frame_size, noverlap = audio_frame_overlap, window = window)[2] + temp_audio_chunk = numpy.stack((numpy.real(temp_audio_chunk), numpy.imag(temp_audio_chunk)), axis = -1).transpose((0, 3, 1, 2)) + temp_audio_chunk = temp_audio_chunk.reshape(-1, 2, 2, audio_trim_size + 1, audio_bin_total).reshape(-1, audio_channel_total, audio_trim_size + 1, audio_bin_total) + temp_audio_chunk = temp_audio_chunk[:, :, :audio_frame_total] + temp_audio_chunk /= numpy.sqrt(1.0 / window.sum() ** 2) + return temp_audio_chunk + + +def compose_audio_chunk(temp_audio_chunk : AudioChunk, audio_trim_size : int) -> AudioChunk: + audio_frame_size = 7680 + audio_frame_overlap = 6656 + audio_frame_total = 3072 + audio_bin_total = 256 + window = scipy.signal.windows.hann(audio_frame_size) + temp_audio_chunk = numpy.pad(temp_audio_chunk, ((0, 0), (0, 0), (0, audio_trim_size + 1 - audio_frame_total), (0, 0))) + temp_audio_chunk = temp_audio_chunk.reshape(-1, 2, audio_trim_size + 1, audio_bin_total).transpose((0, 2, 3, 1)) + temp_audio_chunk = temp_audio_chunk[:, :, :, 0] + 1j * temp_audio_chunk[:, :, :, 1] + temp_audio_chunk = scipy.signal.istft(temp_audio_chunk, nperseg = audio_frame_size, noverlap = audio_frame_overlap, window = window)[1] + temp_audio_chunk *= numpy.sqrt(1.0 / window.sum() ** 2) + return temp_audio_chunk + + +def normalize_audio_chunk(temp_audio_chunk : AudioChunk, chunk_size : int, audio_trim_size : int, audio_pad_size : int) -> AudioChunk: + temp_audio_chunk = temp_audio_chunk.reshape((-1, 2, chunk_size)) + temp_audio_chunk = temp_audio_chunk[:, :, audio_trim_size:-audio_trim_size].transpose(1, 0, 2) + temp_audio_chunk = temp_audio_chunk.reshape(2, -1)[:, :-audio_pad_size].T + return temp_audio_chunk diff --git a/facefusion/workflows/__init__.py b/facefusion/workflows/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/facefusion/workflows/core.py b/facefusion/workflows/core.py new file mode 100644 index 0000000000000000000000000000000000000000..2dcdeef4d57096cb62ee114032542c1bb8d80bfc --- /dev/null +++ b/facefusion/workflows/core.py @@ -0,0 +1,8 @@ +from facefusion import logger, process_manager, translator + + +def is_process_stopping() -> bool: + if process_manager.is_stopping(): + process_manager.end() + logger.info(translator.get('processing_stopped'), __name__) + return process_manager.is_pending() diff --git a/facefusion/workflows/image_to_image.py b/facefusion/workflows/image_to_image.py new file mode 100644 index 0000000000000000000000000000000000000000..4abb765b76596a3b0133e92f4bc05d937a9fe6f8 --- /dev/null +++ b/facefusion/workflows/image_to_image.py @@ -0,0 +1,113 @@ +from functools import partial + +from facefusion import ffmpeg +from facefusion import logger, process_manager, state_manager, translator +from facefusion.audio import create_empty_audio_frame +from facefusion.content_analyser import analyse_image +from facefusion.filesystem import is_image +from facefusion.processors.core import get_processors_modules +from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_file_path +from facefusion.time_helper import calculate_end_time +from facefusion.types import ErrorCode +from facefusion.vision import conditional_merge_vision_mask, detect_image_resolution, extract_vision_mask, pack_resolution, read_static_image, read_static_images, restrict_image_resolution, scale_resolution, write_image +from facefusion.workflows.core import is_process_stopping + + +def process(start_time : float) -> ErrorCode: + tasks =\ + [ + setup, + prepare_image, + process_image, + partial(finalize_image, start_time) + ] + process_manager.start() + + for task in tasks: + error_code = task() # type:ignore[operator] + + if error_code > 0: + process_manager.end() + return error_code + + process_manager.end() + return 0 + + +def setup() -> ErrorCode: + if analyse_image(state_manager.get_item('target_path')): + return 3 + + logger.debug(translator.get('clearing_temp'), __name__) + clear_temp_directory(state_manager.get_item('target_path')) + logger.debug(translator.get('creating_temp'), __name__) + create_temp_directory(state_manager.get_item('target_path')) + return 0 + + +def prepare_image() -> ErrorCode: + output_image_resolution = scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_image_scale')) + temp_image_resolution = restrict_image_resolution(state_manager.get_item('target_path'), output_image_resolution) + + logger.info(translator.get('copying_image').format(resolution = pack_resolution(temp_image_resolution)), __name__) + if ffmpeg.copy_image(state_manager.get_item('target_path'), temp_image_resolution): + logger.debug(translator.get('copying_image_succeeded'), __name__) + else: + logger.error(translator.get('copying_image_failed'), __name__) + process_manager.end() + return 1 + return 0 + + +def process_image() -> ErrorCode: + temp_image_path = get_temp_file_path(state_manager.get_item('target_path')) + reference_vision_frame = read_static_image(temp_image_path) + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) + source_audio_frame = create_empty_audio_frame() + source_voice_frame = create_empty_audio_frame() + target_vision_frame = read_static_image(temp_image_path, 'rgba') + temp_vision_frame = target_vision_frame.copy() + temp_vision_mask = extract_vision_mask(temp_vision_frame) + + for processor_module in get_processors_modules(state_manager.get_item('processors')): + logger.info(translator.get('processing'), processor_module.__name__) + + 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_frame': target_vision_frame[:, :, :3], + 'temp_vision_frame': temp_vision_frame[:, :, :3], + 'temp_vision_mask': temp_vision_mask + }) + + processor_module.post_process() + + temp_vision_frame = conditional_merge_vision_mask(temp_vision_frame, temp_vision_mask) + write_image(temp_image_path, temp_vision_frame) + + if is_process_stopping(): + return 4 + return 0 + + +def finalize_image(start_time : float) -> ErrorCode: + output_image_resolution = scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_image_scale')) + + logger.info(translator.get('finalizing_image').format(resolution = pack_resolution(output_image_resolution)), __name__) + if ffmpeg.finalize_image(state_manager.get_item('target_path'), state_manager.get_item('output_path'), output_image_resolution): + logger.debug(translator.get('finalizing_image_succeeded'), __name__) + else: + logger.warn(translator.get('finalizing_image_skipped'), __name__) + + logger.debug(translator.get('clearing_temp'), __name__) + clear_temp_directory(state_manager.get_item('target_path')) + + if is_image(state_manager.get_item('output_path')): + logger.info(translator.get('processing_image_succeeded').format(seconds = calculate_end_time(start_time)), __name__) + else: + logger.error(translator.get('processing_image_failed'), __name__) + return 1 + return 0 diff --git a/facefusion/workflows/image_to_video.py b/facefusion/workflows/image_to_video.py new file mode 100644 index 0000000000000000000000000000000000000000..22ca7fb412f74fc86256ebd00c36bac7d2de5d29 --- /dev/null +++ b/facefusion/workflows/image_to_video.py @@ -0,0 +1,197 @@ +from concurrent.futures import ThreadPoolExecutor, as_completed +from functools import partial + +import numpy +from tqdm import tqdm + +from facefusion import ffmpeg +from facefusion import logger, process_manager, state_manager, translator, video_manager +from facefusion.audio import create_empty_audio_frame, get_audio_frame, get_voice_frame +from facefusion.common_helper import get_first +from facefusion.content_analyser import analyse_video +from facefusion.filesystem import filter_audio_paths, is_video +from facefusion.processors.core import get_processors_modules +from facefusion.temp_helper import clear_temp_directory, create_temp_directory, move_temp_file, resolve_temp_frame_paths +from facefusion.time_helper import calculate_end_time +from facefusion.types import ErrorCode +from facefusion.vision import conditional_merge_vision_mask, detect_video_resolution, extract_vision_mask, pack_resolution, read_static_image, read_static_images, read_static_video_frame, restrict_trim_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, write_image +from facefusion.workflows.core import is_process_stopping + + +def process(start_time : float) -> ErrorCode: + tasks =\ + [ + setup, + extract_frames, + process_video, + merge_frames, + restore_audio, + partial(finalize_video, start_time) + ] + process_manager.start() + + for task in tasks: + error_code = task() # type:ignore[operator] + + if error_code > 0: + process_manager.end() + return error_code + + process_manager.end() + return 0 + + +def setup() -> ErrorCode: + trim_frame_start, trim_frame_end = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) + + if analyse_video(state_manager.get_item('target_path'), trim_frame_start, trim_frame_end): + return 3 + + logger.debug(translator.get('clearing_temp'), __name__) + clear_temp_directory(state_manager.get_item('target_path')) + logger.debug(translator.get('creating_temp'), __name__) + create_temp_directory(state_manager.get_item('target_path')) + return 0 + + +def extract_frames() -> ErrorCode: + trim_frame_start, trim_frame_end = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) + output_video_resolution = scale_resolution(detect_video_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_video_scale')) + temp_video_resolution = restrict_video_resolution(state_manager.get_item('target_path'), output_video_resolution) + temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + logger.info(translator.get('extracting_frames').format(resolution=pack_resolution(temp_video_resolution), fps=temp_video_fps), __name__) + + if ffmpeg.extract_frames(state_manager.get_item('target_path'), temp_video_resolution, temp_video_fps, trim_frame_start, trim_frame_end): + logger.debug(translator.get('extracting_frames_succeeded'), __name__) + else: + if is_process_stopping(): + return 4 + logger.error(translator.get('extracting_frames_failed'), __name__) + return 1 + return 0 + + +def process_video() -> ErrorCode: + temp_frame_paths = resolve_temp_frame_paths(state_manager.get_item('target_path')) + + 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: + 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 = [] + + for frame_number, temp_frame_path in enumerate(temp_frame_paths): + 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() + + if not future.cancelled(): + future.result() + progress.update() + + 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: + trim_frame_start, trim_frame_end = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) + output_video_resolution = scale_resolution(detect_video_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_video_scale')) + temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + + logger.info(translator.get('merging_video').format(resolution = pack_resolution(output_video_resolution), fps = state_manager.get_item('output_video_fps')), __name__) + if ffmpeg.merge_video(state_manager.get_item('target_path'), temp_video_fps, output_video_resolution, state_manager.get_item('output_video_fps'), trim_frame_start, trim_frame_end): + logger.debug(translator.get('merging_video_succeeded'), __name__) + else: + if is_process_stopping(): + return 4 + logger.error(translator.get('merging_video_failed'), __name__) + return 1 + return 0 + + +def restore_audio() -> ErrorCode: + trim_frame_start, trim_frame_end = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) + + if state_manager.get_item('output_audio_volume') == 0: + logger.info(translator.get('skipping_audio'), __name__) + move_temp_file(state_manager.get_item('target_path'), state_manager.get_item('output_path')) + else: + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) + if source_audio_path: + if ffmpeg.replace_audio(state_manager.get_item('target_path'), source_audio_path, state_manager.get_item('output_path')): + video_manager.clear_video_pool() + logger.debug(translator.get('replacing_audio_succeeded'), __name__) + else: + video_manager.clear_video_pool() + if is_process_stopping(): + return 4 + logger.warn(translator.get('replacing_audio_skipped'), __name__) + move_temp_file(state_manager.get_item('target_path'), state_manager.get_item('output_path')) + else: + if ffmpeg.restore_audio(state_manager.get_item('target_path'), state_manager.get_item('output_path'), trim_frame_start, trim_frame_end): + video_manager.clear_video_pool() + logger.debug(translator.get('restoring_audio_succeeded'), __name__) + else: + video_manager.clear_video_pool() + if is_process_stopping(): + return 4 + logger.warn(translator.get('restoring_audio_skipped'), __name__) + move_temp_file(state_manager.get_item('target_path'), state_manager.get_item('output_path')) + return 0 + + +def process_temp_frame(temp_frame_path : str, frame_number : int) -> bool: + reference_vision_frame = read_static_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) + temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + target_vision_frame = read_static_image(temp_frame_path, 'rgba') + temp_vision_frame = target_vision_frame.copy() + temp_vision_mask = extract_vision_mask(temp_vision_frame) + + source_audio_frame = get_audio_frame(source_audio_path, temp_video_fps, frame_number) + source_voice_frame = get_voice_frame(source_audio_path, temp_video_fps, frame_number) + + if not numpy.any(source_audio_frame): + source_audio_frame = create_empty_audio_frame() + if not numpy.any(source_voice_frame): + source_voice_frame = create_empty_audio_frame() + + 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_frame': target_vision_frame[:, :, :3], + 'temp_vision_frame': temp_vision_frame[:, :, :3], + 'temp_vision_mask': temp_vision_mask + }) + + temp_vision_frame = conditional_merge_vision_mask(temp_vision_frame, temp_vision_mask) + return write_image(temp_frame_path, temp_vision_frame) + + +def finalize_video(start_time : float) -> ErrorCode: + logger.debug(translator.get('clearing_temp'), __name__) + clear_temp_directory(state_manager.get_item('target_path')) + + if is_video(state_manager.get_item('output_path')): + logger.info(translator.get('processing_video_succeeded').format(seconds = calculate_end_time(start_time)), __name__) + else: + logger.error(translator.get('processing_video_failed'), __name__) + return 1 + return 0 diff --git a/install.py b/install.py new file mode 100644 index 0000000000000000000000000000000000000000..476d588212802c71fe71d60fc54325fe9e458836 --- /dev/null +++ b/install.py @@ -0,0 +1,10 @@ +#!/usr/bin/env python3 + +import os + +os.environ['SYSTEM_VERSION_COMPAT'] = '0' + +from facefusion import installer + +if __name__ == '__main__': + installer.cli() diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 0000000000000000000000000000000000000000..c25ac312bdbb27ddc88bf7bfba8a5277e329d634 --- /dev/null +++ b/mypy.ini @@ -0,0 +1,7 @@ +[mypy] +check_untyped_defs = True +disallow_any_generics = True +disallow_untyped_calls = True +disallow_untyped_defs = True +ignore_missing_imports = True +strict_optional = False diff --git a/packages.txt b/packages.txt new file mode 100644 index 0000000000000000000000000000000000000000..134576534c667fc62b9a4cea4415ce4e7115dd8b --- /dev/null +++ b/packages.txt @@ -0,0 +1,3 @@ +ffmpeg +libsm6 +libxext6 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..10945b4253654a427c8b937d519d9405879f3bb3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +gradio-rangeslider==0.0.8 +gradio==5.44.1 +numpy==2.2.6 +onnx==1.19.1 +onnxruntime==1.23.2 +opencv-python==4.12.0.88 +psutil==7.1.3 +tqdm==4.67.1 +scipy==1.16.3 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/helper.py b/tests/helper.py new file mode 100644 index 0000000000000000000000000000000000000000..48a7e2368b2aee5ccf46f53d4d91fa9c34a0d7df --- /dev/null +++ b/tests/helper.py @@ -0,0 +1,44 @@ +import os +import tempfile + +from facefusion.filesystem import create_directory, is_directory, is_file, remove_directory +from facefusion.types import JobStatus + + +def is_test_job_file(file_path : str, job_status : JobStatus) -> bool: + return is_file(get_test_job_file(file_path, job_status)) + + +def get_test_job_file(file_path : str, job_status : JobStatus) -> str: + return os.path.join(get_test_jobs_directory(), job_status, file_path) + + +def get_test_jobs_directory() -> str: + return os.path.join(tempfile.gettempdir(), 'facefusion-test-jobs') + + +def get_test_example_file(file_path : str) -> str: + return os.path.join(get_test_examples_directory(), file_path) + + +def get_test_examples_directory() -> str: + return os.path.join(tempfile.gettempdir(), 'facefusion-test-examples') + + +def is_test_output_file(file_path : str) -> bool: + return is_file(get_test_output_file(file_path)) + + +def get_test_output_file(file_path : str) -> str: + return os.path.join(get_test_outputs_directory(), file_path) + + +def get_test_outputs_directory() -> str: + return os.path.join(tempfile.gettempdir(), 'facefusion-test-outputs') + + +def prepare_test_output_directory() -> bool: + test_outputs_directory = get_test_outputs_directory() + remove_directory(test_outputs_directory) + create_directory(test_outputs_directory) + return is_directory(test_outputs_directory) diff --git a/tests/test_audio.py b/tests/test_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..31160115b717263175ce8c8371dbf6ed73fa097e --- /dev/null +++ b/tests/test_audio.py @@ -0,0 +1,28 @@ +import subprocess + +import pytest + +from facefusion.audio import get_audio_frame, read_static_audio +from facefusion.download import conditional_download +from .helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.mp3' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), get_test_example_file('source.wav') ]) + + +def test_get_audio_frame() -> None: + assert hasattr(get_audio_frame(get_test_example_file('source.mp3'), 25), '__array_interface__') + assert hasattr(get_audio_frame(get_test_example_file('source.wav'), 25), '__array_interface__') + assert get_audio_frame('invalid', 25) is None + + +def test_read_static_audio() -> None: + assert len(read_static_audio(get_test_example_file('source.mp3'), 25)) == 280 + assert len(read_static_audio(get_test_example_file('source.wav'), 25)) == 280 + assert read_static_audio('invalid', 25) is None diff --git a/tests/test_cli_age_modifier.py b/tests/test_cli_age_modifier.py new file mode 100644 index 0000000000000000000000000000000000000000..8394f0319d0ccb0ffee4444a06032b71cd361bcc --- /dev/null +++ b/tests/test_cli_age_modifier.py @@ -0,0 +1,38 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_modify_age_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'age_modifier', '--age-modifier-direction', '100', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-age-face-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-age-face-to-image.jpg') is True + + +def test_modify_age_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'age_modifier', '--age-modifier-direction', '100', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-age-face-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-age-face-to-video.mp4') is True diff --git a/tests/test_cli_background_remover.py b/tests/test_cli_background_remover.py new file mode 100644 index 0000000000000000000000000000000000000000..f3051a72b1e28b96c588642b67203c6e0ab1f46e --- /dev/null +++ b/tests/test_cli_background_remover.py @@ -0,0 +1,39 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run(['ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg')]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_remove_background_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'background_remover', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-remove-background-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-remove-background-to-image.jpg') is True + + +def test_remove_background_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'background_remover', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-remove-background-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-remove-background-to-video.mp4') is True diff --git a/tests/test_cli_batch_runner.py b/tests/test_cli_batch_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..b615e53cf3254648bde2943b2daef18fdc4b3e78 --- /dev/null +++ b/tests/test_cli_batch_runner.py @@ -0,0 +1,45 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p-batch-1.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '2', get_test_example_file('target-240p-batch-2.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_batch_run_targets() -> None: + commands = [ sys.executable, 'facefusion.py', 'batch-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p-batch-*.jpg'), '-o', get_test_output_file('test-batch-run-targets-{index}.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-batch-run-targets-0.jpg') is True + assert is_test_output_file('test-batch-run-targets-1.jpg') is True + assert is_test_output_file('test-batch-run-targets-2.jpg') is False + + +def test_batch_run_sources_to_targets() -> None: + commands = [ sys.executable, 'facefusion.py', 'batch-run', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('target-240p-batch-*.jpg'), '-t', get_test_example_file('target-240p-batch-*.jpg'), '-o', get_test_output_file('test-batch-run-sources-to-targets-{index}.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-batch-run-sources-to-targets-0.jpg') is True + assert is_test_output_file('test-batch-run-sources-to-targets-1.jpg') is True + assert is_test_output_file('test-batch-run-sources-to-targets-2.jpg') is True + assert is_test_output_file('test-batch-run-sources-to-targets-3.jpg') is True + assert is_test_output_file('test-batch-run-sources-to-targets-4.jpg') is False diff --git a/tests/test_cli_expression_restorer.py b/tests/test_cli_expression_restorer.py new file mode 100644 index 0000000000000000000000000000000000000000..01f7dc417b5abf83b7239a96a76466cb3970084d --- /dev/null +++ b/tests/test_cli_expression_restorer.py @@ -0,0 +1,38 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_restore_expression_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'expression_restorer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-restore-expression-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-restore-expression-to-image.jpg') is True + + +def test_restore_expression_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'expression_restorer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-restore-expression-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-restore-expression-to-video.mp4') is True diff --git a/tests/test_cli_face_debugger.py b/tests/test_cli_face_debugger.py new file mode 100644 index 0000000000000000000000000000000000000000..d2b382c6d41b68f4e228e7a69577ecf0c5fcd0f5 --- /dev/null +++ b/tests/test_cli_face_debugger.py @@ -0,0 +1,39 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_debug_face_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-debug-face-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-debug-face-to-image.jpg') is True + + +def test_debug_face_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-debug-face-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-debug-face-to-video.mp4') is True diff --git a/tests/test_cli_face_editor.py b/tests/test_cli_face_editor.py new file mode 100644 index 0000000000000000000000000000000000000000..afd16c7fdb429b5fd72f7e390956258b9d8d5338 --- /dev/null +++ b/tests/test_cli_face_editor.py @@ -0,0 +1,39 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_edit_face_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_editor', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-edit-face-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-edit-face-to-image.jpg') is True + + +def test_edit_face_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_editor', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-edit-face-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-edit-face-to-video.mp4') is True diff --git a/tests/test_cli_face_enhancer.py b/tests/test_cli_face_enhancer.py new file mode 100644 index 0000000000000000000000000000000000000000..8d1fd2f89461ed01915eef072faff729d7f81deb --- /dev/null +++ b/tests/test_cli_face_enhancer.py @@ -0,0 +1,39 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_enhance_face_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_enhancer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-enhance-face-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-enhance-face-to-image.jpg') is True + + +def test_enhance_face_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-enhance-face-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-enhance-face-to-video.mp4') is True diff --git a/tests/test_cli_face_swapper.py b/tests/test_cli_face_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..bf5a178ff18a0c9e6c566facaeadccdf972f5ed8 --- /dev/null +++ b/tests/test_cli_face_swapper.py @@ -0,0 +1,39 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_swap_face_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_swapper', '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-swap-face-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-swap-face-to-image.jpg') is True + + +def test_swap_face_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_swapper', '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-swap-face-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-swap-face-to-video.mp4') is True diff --git a/tests/test_cli_frame_colorizer.py b/tests/test_cli_frame_colorizer.py new file mode 100644 index 0000000000000000000000000000000000000000..9dbe8c80669e528ea5d4ab98b5e40f2507c6dda4 --- /dev/null +++ b/tests/test_cli_frame_colorizer.py @@ -0,0 +1,40 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', '-vf', 'hue=s=0', get_test_example_file('target-240p-0sat.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'hue=s=0', get_test_example_file('target-240p-0sat.mp4') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_colorize_frame_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_colorizer', '-t', get_test_example_file('target-240p-0sat.jpg'), '-o', get_test_output_file('test_colorize-frame-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test_colorize-frame-to-image.jpg') is True + + +def test_colorize_frame_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_colorizer', '-t', get_test_example_file('target-240p-0sat.mp4'), '-o', get_test_output_file('test-colorize-frame-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-colorize-frame-to-video.mp4') is True diff --git a/tests/test_cli_frame_enhancer.py b/tests/test_cli_frame_enhancer.py new file mode 100644 index 0000000000000000000000000000000000000000..729259dad36c118bbf563b39fd154e4bae802f9b --- /dev/null +++ b/tests/test_cli_frame_enhancer.py @@ -0,0 +1,39 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_enhance_frame_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-enhance-frame-to-image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-enhance-frame-to-image.jpg') is True + + +def test_enhance_frame_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-enhance-frame-to-video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test-enhance-frame-to-video.mp4') is True diff --git a/tests/test_cli_job_manager.py b/tests/test_cli_job_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..ddad5cea061ea80df6b74e586b26bf9145590ef2 --- /dev/null +++ b/tests/test_cli_job_manager.py @@ -0,0 +1,209 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, count_step_total, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_job_file + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + + +@pytest.mark.skip() +def test_job_list() -> None: + pass + + +def test_job_create() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-create', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_job_file('test-job-create.json', 'drafted') is True + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-create', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + +def test_job_submit() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-submit', 'test-job-submit', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-submit', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-submit', 'test-job-submit', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_job_file('test-job-submit.json', 'queued') is True + assert subprocess.run(commands).returncode == 1 + + +def test_submit_all() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-submit-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-submit-all-1', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-submit-all-2', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit-all-1', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit-all-2', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-submit-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_job_file('test-job-submit-all-1.json', 'queued') is True + assert is_test_job_file('test-job-submit-all-2.json', 'queued') is True + assert subprocess.run(commands).returncode == 1 + + +def test_job_delete() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-delete', 'test-job-delete', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-delete', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-delete', 'test-job-delete', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_job_file('test-job-delete.json', 'drafted') is False + assert subprocess.run(commands).returncode == 1 + + +def test_job_delete_all() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-delete-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-delete-all-1', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-delete-all-2', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-delete-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_job_file('test-job-delete-all-1.json', 'drafted') is False + assert is_test_job_file('test-job-delete-all-2.json', 'drafted') is False + assert subprocess.run(commands).returncode == 1 + + +def test_job_add_step() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-add-step', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + + assert subprocess.run(commands).returncode == 1 + assert count_step_total('test-job-add-step') == 0 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-add-step', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-add-step', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert count_step_total('test-job-add-step') == 1 + + +def test_job_remix() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '0', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + + assert subprocess.run(commands).returncode == 1 + assert count_step_total('test-job-remix-step') == 0 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-remix-step', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-remix-step', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '0', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + + assert count_step_total('test-job-remix-step') == 1 + assert subprocess.run(commands).returncode == 0 + assert count_step_total('test-job-remix-step') == 2 + + commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '-1', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert count_step_total('test-job-remix-step') == 3 + + +def test_job_insert_step() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '0', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + + assert subprocess.run(commands).returncode == 1 + assert count_step_total('test-job-insert-step') == 0 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-insert-step', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-insert-step', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '0', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + + assert count_step_total('test-job-insert-step') == 1 + assert subprocess.run(commands).returncode == 0 + assert count_step_total('test-job-insert-step') == 2 + + commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '-1', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert count_step_total('test-job-insert-step') == 3 + + +def test_job_remove_step() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-remove-step', 'test-job-remove-step', '0', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-remove-step', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-remove-step', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] + subprocess.run(commands) + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-remove-step', 'test-job-remove-step', '0', '--jobs-path', get_test_jobs_directory() ] + + assert count_step_total('test-job-remove-step') == 2 + assert subprocess.run(commands).returncode == 0 + assert count_step_total('test-job-remove-step') == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-remove-step', 'test-job-remove-step', '-1', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 0 + assert subprocess.run(commands).returncode == 1 + assert count_step_total('test-job-remove-step') == 0 diff --git a/tests/test_cli_job_runner.py b/tests/test_cli_job_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..e5080c6242f2e2387b06a1d33b6c4844fccc3a21 --- /dev/null +++ b/tests/test_cli_job_runner.py @@ -0,0 +1,147 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs, move_job_file, set_steps_status +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_job_run() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-run', 'test-job-run', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-run', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-run.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-run', 'test-job-run', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-submit', 'test-job-run', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-run', 'test-job-run', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 0 + assert subprocess.run(commands).returncode == 1 + assert is_test_output_file('test-job-run.jpg') is True + + +def test_job_run_all() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-run-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-run-all-1', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-run-all-2', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-1', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-run-all-1.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-2', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-job-run-all-2.mp4'), '--trim-frame-end', '1' ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-2', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-job-run-all-2.mp4'), '--trim-frame-start', '0', '--trim-frame-end', '1' ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-run-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-submit-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-run-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 0 + assert subprocess.run(commands).returncode == 1 + assert is_test_output_file('test-job-run-all-1.jpg') is True + assert is_test_output_file('test-job-run-all-2.mp4') is True + + +def test_job_retry() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-retry', 'test-job-retry', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-retry', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-retry.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-retry', 'test-job-retry', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + set_steps_status('test-job-retry', 'failed') + move_job_file('test-job-retry', 'failed') + + commands = [ sys.executable, 'facefusion.py', 'job-retry', 'test-job-retry', '--jobs-path', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 0 + assert subprocess.run(commands).returncode == 1 + assert is_test_output_file('test-job-retry.jpg') is True + + +def test_job_retry_all() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-retry-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-retry-all-1', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-retry-all-2', '--jobs-path', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-1', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-retry-all-1.jpg') ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-2', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-job-retry-all-2.mp4'), '--trim-frame-end', '1' ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-2', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-job-retry-all-2.mp4'), '--trim-frame-start', '0', '--trim-frame-end', '1' ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-retry-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 1 + + set_steps_status('test-job-retry-all-1', 'failed') + set_steps_status('test-job-retry-all-2', 'failed') + move_job_file('test-job-retry-all-1', 'failed') + move_job_file('test-job-retry-all-2', 'failed') + + commands = [ sys.executable, 'facefusion.py', 'job-retry-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] + + assert subprocess.run(commands).returncode == 0 + assert subprocess.run(commands).returncode == 1 + assert is_test_output_file('test-job-retry-all-1.jpg') is True + assert is_test_output_file('test-job-retry-all-2.mp4') is True diff --git a/tests/test_cli_lip_syncer.py b/tests/test_cli_lip_syncer.py new file mode 100644 index 0000000000000000000000000000000000000000..d88c7a88e0c54c6ad77df2c8c3812c9deecbe0a2 --- /dev/null +++ b/tests/test_cli_lip_syncer.py @@ -0,0 +1,40 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.mp3', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def test_sync_lip_to_image() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'lip_syncer', '-s', get_test_example_file('source.mp3'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test_sync_lip_to_image.jpg') ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test_sync_lip_to_image.jpg') is True + + +def test_sync_lip_to_video() -> None: + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'lip_syncer', '-s', get_test_example_file('source.mp3'), '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test_sync_lip_to_video.mp4'), '--trim-frame-end', '1' ] + + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test_sync_lip_to_video.mp4') is True diff --git a/tests/test_cli_output_scale.py b/tests/test_cli_output_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..92f57053479ca60c70a993df7ce3679b64def405 --- /dev/null +++ b/tests/test_cli_output_scale.py @@ -0,0 +1,57 @@ +import subprocess +import sys + +import pytest + +from facefusion.download import conditional_download +from facefusion.jobs.job_manager import clear_jobs, init_jobs +from facefusion.types import Resolution, Scale +from facefusion.vision import detect_image_resolution, detect_video_resolution +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +@pytest.mark.parametrize('output_image_scale, output_image_resolution', +[ + (0.5, (212, 112)), + (1.0, (426, 226)), + (2.0, (852, 452)), + (8.0, (3408, 1808)) +]) +def test_output_image_scale(output_image_scale : Scale, output_image_resolution : Resolution) -> None: + output_file_path = get_test_output_file('test-output-image-scale-' + str(output_image_scale) + '.jpg') + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.jpg'), '-o', output_file_path, '--output-image-scale', str(output_image_scale) ] + + assert subprocess.run(commands).returncode == 0 + assert detect_image_resolution(output_file_path) == output_image_resolution + + +@pytest.mark.parametrize('output_video_scale, output_video_resolution', +[ + (0.5, (212, 112)), + (1.0, (426, 226)), + (2.0, (852, 452)), + (8.0, (3408, 1808)) +]) +def test_output_video_scale(output_video_scale : Scale, output_video_resolution : Resolution) -> None: + output_file_path = get_test_output_file('test-output-video-scale-' + str(output_video_scale) + '.mp4') + commands = [ sys.executable, 'facefusion.py', 'headless-run', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', output_file_path, '--trim-frame-end', '1', '--output-video-scale', str(output_video_scale) ] + + assert subprocess.run(commands).returncode == 0 + assert detect_video_resolution(output_file_path) == output_video_resolution diff --git a/tests/test_common_helper.py b/tests/test_common_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..1a720b6ffcb12b9fc7b28fd1a86d319018c31ef9 --- /dev/null +++ b/tests/test_common_helper.py @@ -0,0 +1,27 @@ +from facefusion.common_helper import calculate_float_step, calculate_int_step, create_float_metavar, create_float_range, create_int_metavar, create_int_range + + +def test_create_int_metavar() -> None: + assert create_int_metavar([ 1, 2, 3, 4, 5 ]) == '[1..5:1]' + + +def test_create_float_metavar() -> None: + assert create_float_metavar([ 0.1, 0.2, 0.3, 0.4, 0.5 ]) == '[0.1..0.5:0.1]' + + +def test_create_int_range() -> None: + assert create_int_range(0, 2, 1) == [ 0, 1, 2 ] + assert create_float_range(0, 1, 1) == [ 0, 1 ] + + +def test_create_float_range() -> None: + assert create_float_range(0.0, 1.0, 0.5) == [ 0.0, 0.5, 1.0 ] + assert create_float_range(0.0, 1.0, 0.05) == [ 0.0, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90, 0.95, 1.0 ] + + +def test_calc_int_step() -> None: + assert calculate_int_step([0, 1]) == 1 + + +def test_calc_float_step() -> None: + assert calculate_float_step([0.1, 0.2]) == 0.1 diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000000000000000000000000000000000000..48ae2c07e759dc655ea32c64bd1b178c53da53a7 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,85 @@ +from configparser import ConfigParser + +import pytest + +from facefusion import config + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + config.CONFIG_PARSER = ConfigParser() + config.CONFIG_PARSER.read_dict( + { + 'str': + { + 'valid': 'a', + 'unset': '' + }, + 'int': + { + 'valid': '1', + 'unset': '' + }, + 'float': + { + 'valid': '1.0', + 'unset': '' + }, + 'bool': + { + 'valid': 'True', + 'unset': '' + }, + 'str_list': + { + 'valid': 'a b c', + 'unset': '' + }, + 'int_list': + { + 'valid': '1 2 3', + 'unset': '' + } + }) + + +def test_get_str_value() -> None: + assert config.get_str_value('str', 'valid') == 'a' + assert config.get_str_value('str', 'unset', 'b') == 'b' + assert config.get_str_value('str', 'unset') is None + assert config.get_str_value('str', 'invalid') is None + + +def test_get_int_value() -> None: + assert config.get_int_value('int', 'valid') == 1 + assert config.get_int_value('int', 'unset', '1') == 1 + assert config.get_int_value('int', 'unset') is None + assert config.get_int_value('int', 'invalid') is None + + +def test_get_float_value() -> None: + assert config.get_float_value('float', 'valid') == 1.0 + assert config.get_float_value('float', 'unset', '1.0') == 1.0 + assert config.get_float_value('float', 'unset') is None + assert config.get_float_value('float', 'invalid') is None + + +def test_get_bool_value() -> None: + assert config.get_bool_value('bool', 'valid') is True + assert config.get_bool_value('bool', 'unset', 'False') is False + assert config.get_bool_value('bool', 'unset') is None + assert config.get_bool_value('bool', 'invalid') is None + + +def test_get_str_list() -> None: + assert config.get_str_list('str_list', 'valid') == [ 'a', 'b', 'c' ] + assert config.get_str_list('str_list', 'unset', 'c b a') == [ 'c', 'b', 'a' ] + assert config.get_str_list('str_list', 'unset') is None + assert config.get_str_list('str_list', 'invalid') is None + + +def test_get_int_list() -> None: + assert config.get_int_list('int_list', 'valid') == [ 1, 2, 3 ] + assert config.get_int_list('int_list', 'unset', '3 2 1') == [ 3, 2, 1 ] + assert config.get_int_list('int_list', 'unset') is None + assert config.get_int_list('int_list', 'invalid') is None diff --git a/tests/test_curl_builder.py b/tests/test_curl_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..54feaacc887ebf05894c683dac8017d653d22eb3 --- /dev/null +++ b/tests/test_curl_builder.py @@ -0,0 +1,17 @@ +from shutil import which + +from facefusion import metadata +from facefusion.curl_builder import chain, ping, run, set_timeout + + +def test_run() -> None: + user_agent = metadata.get('name') + '/' + metadata.get('version') + + assert run([]) == [ which('curl'), '--user-agent', user_agent, '--insecure', '--location', '--silent' ] + + +def test_chain() -> None: + assert chain( + ping(metadata.get('url')), + set_timeout(5) + ) == [ '-I', metadata.get('url'), '--connect-timeout', '5' ] diff --git a/tests/test_download.py b/tests/test_download.py new file mode 100644 index 0000000000000000000000000000000000000000..4186f73cee27e255bf9b9ae3f0a46fc07673bfc3 --- /dev/null +++ b/tests/test_download.py @@ -0,0 +1,18 @@ +from facefusion.download import get_static_download_size, ping_static_url, resolve_download_url_by_provider + + +def test_get_static_download_size() -> None: + assert get_static_download_size('https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/fairface.onnx') == 85170772 + assert get_static_download_size('https://huggingface.co/facefusion/models-3.0.0/resolve/main/fairface.onnx') == 85170772 + assert get_static_download_size('invalid') == 0 + + +def test_static_ping_url() -> None: + assert ping_static_url('https://github.com') is True + assert ping_static_url('https://huggingface.co') is True + assert ping_static_url('invalid') is False + + +def test_resolve_download_url_by_provider() -> None: + assert resolve_download_url_by_provider('github', 'models-3.0.0', 'fairface.onnx') == 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/fairface.onnx' + assert resolve_download_url_by_provider('huggingface', 'models-3.0.0', 'fairface.onnx') == 'https://huggingface.co/facefusion/models-3.0.0/resolve/main/fairface.onnx' diff --git a/tests/test_execution.py b/tests/test_execution.py new file mode 100644 index 0000000000000000000000000000000000000000..54b6f53660b99e3483ccdc5360981143ce5bfeed --- /dev/null +++ b/tests/test_execution.py @@ -0,0 +1,24 @@ +from facefusion.execution import create_inference_session_providers, get_available_execution_providers, has_execution_provider + + +def test_has_execution_provider() -> None: + assert has_execution_provider('cpu') is True + assert has_execution_provider('openvino') is False + + +def test_get_available_execution_providers() -> None: + assert 'cpu' in get_available_execution_providers() + + +def test_create_inference_session_providers() -> None: + inference_session_providers =\ + [ + ('CUDAExecutionProvider', + { + 'device_id': 1, + 'cudnn_conv_algo_search': 'EXHAUSTIVE' + }), + 'CPUExecutionProvider' + ] + + assert create_inference_session_providers(1, [ 'cpu', 'cuda' ]) == inference_session_providers diff --git a/tests/test_face_analyser.py b/tests/test_face_analyser.py new file mode 100644 index 0000000000000000000000000000000000000000..c0b0930b45674e76835a6d1b9b2b6a9ced66b982 --- /dev/null +++ b/tests/test_face_analyser.py @@ -0,0 +1,131 @@ +import subprocess + +import pytest + +from facefusion import face_classifier, face_detector, face_landmarker, face_recognizer, state_manager +from facefusion.download import conditional_download +from facefusion.face_analyser import get_many_faces +from facefusion.vision import read_static_image +from .helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.8:ih*0.8', get_test_example_file('source-80crop.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.7:ih*0.7', get_test_example_file('source-70crop.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.6:ih*0.6', get_test_example_file('source-60crop.jpg') ]) + state_manager.init_item('execution_device_ids', [ 0 ]) + state_manager.init_item('execution_providers', [ 'cpu' ]) + state_manager.init_item('download_providers', [ 'github' ]) + state_manager.init_item('face_detector_angles', [ 0 ]) + state_manager.init_item('face_detector_model', 'many') + state_manager.init_item('face_detector_score', 0.5) + state_manager.init_item('face_landmarker_model', 'many') + state_manager.init_item('face_landmarker_score', 0.5) + face_classifier.pre_check() + face_landmarker.pre_check() + face_recognizer.pre_check() + + +@pytest.fixture(autouse = True) +def before_each() -> None: + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_recognizer.clear_inference_pool() + + +def test_get_one_face_with_retinaface() -> None: + state_manager.init_item('face_detector_model', 'retinaface') + state_manager.init_item('face_detector_size', '320x320') + state_manager.init_item('face_detector_margin', (0, 0, 0, 0)) + face_detector.pre_check() + + source_paths =\ + [ + get_test_example_file('source.jpg'), + get_test_example_file('source-80crop.jpg'), + get_test_example_file('source-70crop.jpg'), + get_test_example_file('source-60crop.jpg') + ] + + for source_path in source_paths: + source_frame = read_static_image(source_path) + many_faces = get_many_faces([ source_frame ]) + + assert len(many_faces) == 1 + + +def test_get_one_face_with_scrfd() -> None: + state_manager.init_item('face_detector_model', 'scrfd') + state_manager.init_item('face_detector_size', '640x640') + state_manager.init_item('face_detector_margin', (0, 0, 0, 0)) + face_detector.pre_check() + + source_paths =\ + [ + get_test_example_file('source.jpg'), + get_test_example_file('source-80crop.jpg'), + get_test_example_file('source-70crop.jpg'), + get_test_example_file('source-60crop.jpg') + ] + + for source_path in source_paths: + source_frame = read_static_image(source_path) + many_faces = get_many_faces([ source_frame ]) + + assert len(many_faces) == 1 + + +def test_get_one_face_with_yoloface() -> None: + state_manager.init_item('face_detector_model', 'yoloface') + state_manager.init_item('face_detector_size', '640x640') + state_manager.init_item('face_detector_margin', (0, 0, 0, 0)) + face_detector.pre_check() + + source_paths =\ + [ + get_test_example_file('source.jpg'), + get_test_example_file('source-80crop.jpg'), + get_test_example_file('source-70crop.jpg'), + get_test_example_file('source-60crop.jpg') + ] + + for source_path in source_paths: + source_frame = read_static_image(source_path) + many_faces = get_many_faces([ source_frame ]) + + assert len(many_faces) == 1 + + +def test_get_one_face_with_yunet() -> None: + state_manager.init_item('face_detector_model', 'yunet') + state_manager.init_item('face_detector_size', '640x640') + state_manager.init_item('face_detector_margin', (0, 0, 0, 0)) + face_detector.pre_check() + + source_paths =\ + [ + get_test_example_file('source.jpg'), + get_test_example_file('source-80crop.jpg'), + get_test_example_file('source-70crop.jpg'), + get_test_example_file('source-60crop.jpg') + ] + + for source_path in source_paths: + source_frame = read_static_image(source_path) + many_faces = get_many_faces([ source_frame ]) + + assert len(many_faces) == 1 + + +def test_get_many_faces() -> None: + source_path = get_test_example_file('source.jpg') + source_frame = read_static_image(source_path) + many_faces = get_many_faces([ source_frame, source_frame, source_frame ]) + + assert len(many_faces) == 3 diff --git a/tests/test_ffmpeg.py b/tests/test_ffmpeg.py new file mode 100644 index 0000000000000000000000000000000000000000..0e6cd277d7b142cb0a5d669293e13ba94b2dddfc --- /dev/null +++ b/tests/test_ffmpeg.py @@ -0,0 +1,189 @@ +import os +import subprocess +import tempfile + +import pytest + +import facefusion.ffmpeg +from facefusion import 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 +from facefusion.filesystem import copy_file +from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_file_path, resolve_temp_frame_paths +from facefusion.types import EncoderSet +from .helper import get_test_example_file, get_test_examples_directory, get_test_output_file, prepare_test_output_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/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.mp3', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), get_test_example_file('source.wav') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=25', get_test_example_file('target-240p-25fps.mp4') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=30', get_test_example_file('target-240p-30fps.mp4') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=60', get_test_example_file('target-240p-60fps.mp4') ]) + + for output_video_format in [ 'avi', 'm4v', 'mkv', 'mov', 'mp4', 'webm', 'wmv' ]: + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), '-i', get_test_example_file('target-240p.mp4'), '-ar', '16000', get_test_example_file('target-240p-16khz.' + output_video_format) ]) + + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), '-i', get_test_example_file('target-240p.mp4'), '-ar', '48000', get_test_example_file('target-240p-48khz.mp4') ]) + state_manager.init_item('temp_path', tempfile.gettempdir()) + state_manager.init_item('temp_frame_format', 'png') + state_manager.init_item('output_audio_encoder', 'aac') + state_manager.init_item('output_audio_quality', 100) + state_manager.init_item('output_audio_volume', 100) + state_manager.init_item('output_video_encoder', 'libx264') + state_manager.init_item('output_video_quality', 100) + state_manager.init_item('output_video_preset', 'ultrafast') + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + prepare_test_output_directory() + + +def get_available_encoder_set() -> EncoderSet: + if os.getenv('CI'): + return\ + { + 'audio': [ 'aac' ], + 'video': [ 'libx264' ] + } + return facefusion.ffmpeg.get_available_encoder_set() + + +def test_get_available_encoder_set() -> None: + available_encoder_set = get_available_encoder_set() + + assert 'aac' in available_encoder_set.get('audio') + assert 'libx264' in available_encoder_set.get('video') + + +def test_extract_frames() -> None: + test_set =\ + [ + (get_test_example_file('target-240p-25fps.mp4'), 0, 270, 324), + (get_test_example_file('target-240p-25fps.mp4'), 224, 270, 55), + (get_test_example_file('target-240p-25fps.mp4'), 124, 224, 120), + (get_test_example_file('target-240p-25fps.mp4'), 0, 100, 120), + (get_test_example_file('target-240p-30fps.mp4'), 0, 324, 324), + (get_test_example_file('target-240p-30fps.mp4'), 224, 324, 100), + (get_test_example_file('target-240p-30fps.mp4'), 124, 224, 100), + (get_test_example_file('target-240p-30fps.mp4'), 0, 100, 100), + (get_test_example_file('target-240p-60fps.mp4'), 0, 648, 324), + (get_test_example_file('target-240p-60fps.mp4'), 224, 648, 212), + (get_test_example_file('target-240p-60fps.mp4'), 124, 224, 50), + (get_test_example_file('target-240p-60fps.mp4'), 0, 100, 50) + ] + + for target_path, trim_frame_start, trim_frame_end, frame_total in test_set: + create_temp_directory(target_path) + + assert extract_frames(target_path, (452, 240), 30.0, trim_frame_start, trim_frame_end) is True + assert len(resolve_temp_frame_paths(target_path)) == frame_total + + clear_temp_directory(target_path) + + +def test_merge_video() -> None: + target_paths =\ + [ + get_test_example_file('target-240p-16khz.avi'), + get_test_example_file('target-240p-16khz.m4v'), + get_test_example_file('target-240p-16khz.mkv'), + get_test_example_file('target-240p-16khz.mp4'), + get_test_example_file('target-240p-16khz.mov'), + get_test_example_file('target-240p-16khz.webm'), + get_test_example_file('target-240p-16khz.wmv') + ] + output_video_encoders = get_available_encoder_set().get('video') + + for target_path in target_paths: + for output_video_encoder in output_video_encoders: + state_manager.init_item('output_video_encoder', output_video_encoder) + create_temp_directory(target_path) + extract_frames(target_path, (452, 240), 25.0, 0, 1) + + assert merge_video(target_path, 25.0, (452, 240), 25.0, 0, 1) is True + + clear_temp_directory(target_path) + + state_manager.init_item('output_video_encoder', 'libx264') + + +def test_concat_video() -> None: + output_path = get_test_output_file('test-concat-video.mp4') + temp_output_paths =\ + [ + get_test_example_file('target-240p-16khz.mp4'), + get_test_example_file('target-240p-16khz.mp4') + ] + + assert concat_video(output_path, temp_output_paths) is True + + +def test_read_audio_buffer() -> None: + assert isinstance(read_audio_buffer(get_test_example_file('source.mp3'), 1, 16, 1), bytes) + assert isinstance(read_audio_buffer(get_test_example_file('source.wav'), 1, 16, 1), bytes) + assert read_audio_buffer(get_test_example_file('invalid.mp3'), 1, 16, 1) is None + + +def test_restore_audio() -> None: + test_set =\ + [ + (get_test_example_file('target-240p-16khz.avi'), get_test_output_file('target-240p-16khz.avi')), + (get_test_example_file('target-240p-16khz.m4v'), get_test_output_file('target-240p-16khz.m4v')), + (get_test_example_file('target-240p-16khz.mkv'), get_test_output_file('target-240p-16khz.mkv')), + (get_test_example_file('target-240p-16khz.mov'), get_test_output_file('target-240p-16khz.mov')), + (get_test_example_file('target-240p-16khz.mp4'), get_test_output_file('target-240p-16khz.mp4')), + (get_test_example_file('target-240p-48khz.mp4'), get_test_output_file('target-240p-48khz.mp4')), + (get_test_example_file('target-240p-16khz.webm'), get_test_output_file('target-240p-16khz.webm')), + (get_test_example_file('target-240p-16khz.wmv'), get_test_output_file('target-240p-16khz.wmv')) + ] + output_audio_encoders = get_available_encoder_set().get('audio') + + for target_path, output_path in test_set: + create_temp_directory(target_path) + + for output_audio_encoder in output_audio_encoders: + state_manager.init_item('output_audio_encoder', output_audio_encoder) + copy_file(target_path, get_temp_file_path(target_path)) + + assert restore_audio(target_path, output_path, 0, 270) is True + + clear_temp_directory(target_path) + + state_manager.init_item('output_audio_encoder', 'aac') + + +def test_replace_audio() -> None: + test_set =\ + [ + (get_test_example_file('target-240p-16khz.avi'), get_test_output_file('target-240p-16khz.avi')), + (get_test_example_file('target-240p-16khz.m4v'), get_test_output_file('target-240p-16khz.m4v')), + (get_test_example_file('target-240p-16khz.mkv'), get_test_output_file('target-240p-16khz.mkv')), + (get_test_example_file('target-240p-16khz.mov'), get_test_output_file('target-240p-16khz.mov')), + (get_test_example_file('target-240p-16khz.mp4'), get_test_output_file('target-240p-16khz.mp4')), + (get_test_example_file('target-240p-48khz.mp4'), get_test_output_file('target-240p-48khz.mp4')), + (get_test_example_file('target-240p-16khz.webm'), get_test_output_file('target-240p-16khz.webm')) + ] + output_audio_encoders = get_available_encoder_set().get('audio') + + for target_path, output_path in test_set: + create_temp_directory(target_path) + + for output_audio_encoder in output_audio_encoders: + state_manager.init_item('output_audio_encoder', output_audio_encoder) + copy_file(target_path, get_temp_file_path(target_path)) + + assert replace_audio(target_path, get_test_example_file('source.mp3'), output_path) is True + assert replace_audio(target_path, get_test_example_file('source.wav'), output_path) is True + + clear_temp_directory(target_path) + + state_manager.init_item('output_audio_encoder', 'aac') diff --git a/tests/test_ffmpeg_builder.py b/tests/test_ffmpeg_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..15ea737e7b1fab268b7d52ba641798f387887875 --- /dev/null +++ b/tests/test_ffmpeg_builder.py @@ -0,0 +1,110 @@ +from shutil import which + +from facefusion import ffmpeg_builder +from facefusion.ffmpeg_builder import chain, concat, keep_video_alpha, run, select_frame_range, set_audio_quality, set_audio_sample_size, set_stream_mode, set_video_encoder, set_video_fps, set_video_quality + + +def test_run() -> None: + assert run([]) == [ which('ffmpeg'), '-loglevel', 'error' ] + + +def test_chain() -> None: + assert chain( + ffmpeg_builder.set_input('input.mp4'), + ffmpeg_builder.set_output('output.mp4') + ) == [ '-i', 'input.mp4', 'output.mp4' ] + assert chain( + ffmpeg_builder.set_video_encoder('libx264'), + ffmpeg_builder.set_video_fps(30), + ffmpeg_builder.set_audio_encoder('aac') + ) == [ '-c:v', 'libx264', '-vf', 'fps=30', '-c:a', 'aac' ] + + +def test_concat() -> None: + assert concat( + set_video_encoder('libvpx-vp9'), + set_video_fps(30) + ) == [ '-c:v', 'libvpx-vp9', '-vf', 'fps=30' ] + assert concat( + set_video_encoder('libvpx-vp9'), + set_video_fps(30), + keep_video_alpha('libvpx-vp9') + ) == [ '-c:v', 'libvpx-vp9', '-vf', 'fps=30,format=yuva420p' ] + assert concat( + select_frame_range(0, 100, 30), + keep_video_alpha('libvpx-vp9') + ) == [ '-vf', 'trim=start_frame=0:end_frame=100,fps=30,format=yuva420p' ] + + +def test_set_stream_mode() -> None: + assert set_stream_mode('udp') == [ '-f', 'mpegts' ] + assert set_stream_mode('v4l2') == [ '-f', 'v4l2' ] + + +def test_select_frame_range() -> None: + assert select_frame_range(0, None, 30) == [ '-vf', 'trim=start_frame=0,fps=30' ] + assert select_frame_range(None, 100, 30) == [ '-vf', 'trim=end_frame=100,fps=30' ] + assert select_frame_range(0, 100, 30) == [ '-vf', 'trim=start_frame=0:end_frame=100,fps=30' ] + assert select_frame_range(None, None, 30) == [ '-vf', 'fps=30' ] + + +def test_set_audio_sample_size() -> None: + assert set_audio_sample_size(16) == [ '-f', 's16le' ] + assert set_audio_sample_size(32) == [ '-f', 's32le' ] + + +def test_set_audio_quality() -> None: + assert set_audio_quality('aac', 0) == [ '-q:a', '0.1' ] + assert set_audio_quality('aac', 50) == [ '-q:a', '1.0' ] + assert set_audio_quality('aac', 100) == [ '-q:a', '2.0' ] + assert set_audio_quality('libmp3lame', 0) == [ '-q:a', '9' ] + assert set_audio_quality('libmp3lame', 50) == [ '-q:a', '4' ] + assert set_audio_quality('libmp3lame', 100) == [ '-q:a', '0' ] + assert set_audio_quality('libopus', 0) == [ '-b:a', '64k' ] + assert set_audio_quality('libopus', 50) == [ '-b:a', '160k' ] + assert set_audio_quality('libopus', 100) == [ '-b:a', '256k' ] + assert set_audio_quality('libvorbis', 0) == [ '-q:a', '-1.0' ] + assert set_audio_quality('libvorbis', 50) == [ '-q:a', '4.5' ] + assert set_audio_quality('libvorbis', 100) == [ '-q:a', '10.0' ] + assert set_audio_quality('flac', 0) == [] + assert set_audio_quality('flac', 50) == [] + assert set_audio_quality('flac', 100) == [] + + +def test_set_video_quality() -> None: + assert set_video_quality('libx264', 0) == [ '-crf', '51' ] + assert set_video_quality('libx264', 50) == [ '-crf', '26' ] + assert set_video_quality('libx264', 100) == [ '-crf', '0' ] + assert set_video_quality('libx264rgb', 0) == [ '-crf', '51' ] + assert set_video_quality('libx264rgb', 50) == [ '-crf', '26' ] + assert set_video_quality('libx264rgb', 100) == [ '-crf', '0' ] + assert set_video_quality('libx265', 0) == [ '-crf', '51' ] + assert set_video_quality('libx265', 50) == [ '-crf', '26' ] + assert set_video_quality('libx265', 100) == [ '-crf', '0' ] + assert set_video_quality('libvpx-vp9', 0) == [ '-crf', '63' ] + assert set_video_quality('libvpx-vp9', 50) == [ '-crf', '32' ] + assert set_video_quality('libvpx-vp9', 100) == [ '-crf', '0' ] + assert set_video_quality('h264_nvenc', 0) == [ '-cq' , '51' ] + assert set_video_quality('h264_nvenc', 50) == [ '-cq' , '26' ] + assert set_video_quality('h264_nvenc', 100) == [ '-cq' , '0' ] + assert set_video_quality('hevc_nvenc', 0) == [ '-cq' , '51' ] + assert set_video_quality('hevc_nvenc', 50) == [ '-cq' , '26' ] + assert set_video_quality('hevc_nvenc', 100) == [ '-cq' , '0' ] + assert set_video_quality('h264_amf', 0) == [ '-qp_i', '51', '-qp_p', '51', '-qp_b', '51' ] + assert set_video_quality('h264_amf', 50) == [ '-qp_i', '26', '-qp_p', '26', '-qp_b', '26' ] + assert set_video_quality('h264_amf', 100) == [ '-qp_i', '0', '-qp_p', '0', '-qp_b', '0' ] + assert set_video_quality('hevc_amf', 0) == [ '-qp_i', '51', '-qp_p', '51', '-qp_b', '51' ] + assert set_video_quality('hevc_amf', 50) == [ '-qp_i', '26', '-qp_p', '26', '-qp_b', '26' ] + assert set_video_quality('hevc_amf', 100) == [ '-qp_i', '0', '-qp_p', '0', '-qp_b', '0' ] + assert set_video_quality('h264_qsv', 0) == [ '-qp', '51' ] + assert set_video_quality('h264_qsv', 50) == [ '-qp', '26' ] + assert set_video_quality('h264_qsv', 100) == [ '-qp', '0' ] + assert set_video_quality('hevc_qsv', 0) == [ '-qp', '51' ] + assert set_video_quality('hevc_qsv', 50) == [ '-qp', '26' ] + assert set_video_quality('hevc_qsv', 100) == [ '-qp', '0' ] + assert set_video_quality('h264_videotoolbox', 0) == [ '-b:v', '1024k' ] + assert set_video_quality('h264_videotoolbox', 50) == [ '-b:v', '25768k' ] + assert set_video_quality('h264_videotoolbox', 100) == [ '-b:v', '50512k' ] + assert set_video_quality('hevc_videotoolbox', 0) == [ '-b:v', '1024k' ] + assert set_video_quality('hevc_videotoolbox', 50) == [ '-b:v', '25768k' ] + assert set_video_quality('hevc_videotoolbox', 100) == [ '-b:v', '50512k' ] diff --git a/tests/test_filesystem.py b/tests/test_filesystem.py new file mode 100644 index 0000000000000000000000000000000000000000..df52896340a1140478ae3b856b5fa53e032a83d2 --- /dev/null +++ b/tests/test_filesystem.py @@ -0,0 +1,135 @@ +import os.path + +import pytest + +from facefusion.download import conditional_download +from facefusion.filesystem import create_directory, filter_audio_paths, filter_image_paths, get_file_extension, get_file_format, get_file_size, has_audio, has_image, has_video, in_directory, is_audio, is_directory, is_file, is_image, is_video, remove_directory, resolve_file_paths, same_file_extension +from .helper import get_test_example_file, get_test_examples_directory, get_test_outputs_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.mp3', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + + +def test_get_file_size() -> None: + assert get_file_size(get_test_example_file('source.jpg')) == 549458 + assert get_file_size('invalid') == 0 + + +def test_get_file_extension() -> None: + assert get_file_extension('source.jpg') == '.jpg' + assert get_file_extension('source.mp3') == '.mp3' + assert get_file_extension('invalid') is None + + +def test_get_file_format() -> None: + assert get_file_format('source.jpg') == 'jpeg' + assert get_file_format('source.jpeg') == 'jpeg' + assert get_file_format('source.mp3') == 'mp3' + assert get_file_format('invalid') is None + + +def test_same_file_extension() -> None: + assert same_file_extension('source.jpg', 'source.jpg') is True + assert same_file_extension('source.jpg', 'source.mp3') is False + assert same_file_extension('invalid', 'invalid') is False + + +def test_is_file() -> None: + assert is_file(get_test_example_file('source.jpg')) is True + assert is_file(get_test_examples_directory()) is False + assert is_file('invalid') is False + + +def test_is_audio() -> None: + assert is_audio(get_test_example_file('source.mp3')) is True + assert is_audio(get_test_example_file('source.jpg')) is False + assert is_audio('invalid') is False + + +def test_has_audio() -> None: + assert has_audio([ get_test_example_file('source.mp3') ]) is True + assert has_audio([ get_test_example_file('source.mp3'), get_test_example_file('source.jpg') ]) is True + assert has_audio([ get_test_example_file('source.jpg'), get_test_example_file('source.jpg') ]) is False + assert has_audio([ 'invalid' ]) is False + + +def test_is_image() -> None: + assert is_image(get_test_example_file('source.jpg')) is True + assert is_image(get_test_example_file('target-240p.mp4')) is False + assert is_image('invalid') is False + + +def test_has_image() -> None: + assert has_image([ get_test_example_file('source.jpg') ]) is True + assert has_image([ get_test_example_file('source.jpg'), get_test_example_file('source.mp3') ]) is True + assert has_image([ get_test_example_file('source.mp3'), get_test_example_file('source.mp3') ]) is False + assert has_image([ 'invalid' ]) is False + + +def test_is_video() -> None: + assert is_video(get_test_example_file('target-240p.mp4')) is True + assert is_video(get_test_example_file('source.jpg')) is False + assert is_video('invalid') is False + + +def test_has_video() -> None: + assert has_video([ get_test_example_file('target-240p.mp4') ]) is True + assert has_video([ get_test_example_file('target-240p.mp4'), get_test_example_file('source.mp3') ]) is True + assert has_video([ get_test_example_file('source.mp3'), get_test_example_file('source.mp3') ]) is False + assert has_video([ 'invalid' ]) is False + + +def test_filter_audio_paths() -> None: + assert filter_audio_paths([ get_test_example_file('source.jpg'), get_test_example_file('source.mp3') ]) == [ get_test_example_file('source.mp3') ] + assert filter_audio_paths([ get_test_example_file('source.jpg'), get_test_example_file('source.jpg') ]) == [] + assert filter_audio_paths([ 'invalid' ]) == [] + + +def test_filter_image_paths() -> None: + assert filter_image_paths([ get_test_example_file('source.jpg'), get_test_example_file('source.mp3') ]) == [ get_test_example_file('source.jpg') ] + assert filter_image_paths([ get_test_example_file('source.mp3'), get_test_example_file('source.mp3') ]) == [] + assert filter_audio_paths([ 'invalid' ]) == [] + + +def test_resolve_file_paths() -> None: + file_paths = resolve_file_paths(get_test_examples_directory()) + + for file_path in file_paths: + assert file_path == get_test_example_file(file_path) + + assert resolve_file_paths('invalid') == [] + + +def test_create_directory() -> None: + create_directory_path = os.path.join(get_test_outputs_directory(), 'create_directory') + + assert create_directory(create_directory_path) is True + assert create_directory(get_test_example_file('source.jpg')) is False + + +def test_remove_directory() -> None: + remove_directory_path = os.path.join(get_test_outputs_directory(), 'remove_directory') + create_directory(remove_directory_path) + + assert remove_directory(remove_directory_path) is True + assert remove_directory(get_test_example_file('source.jpg')) is False + assert remove_directory('invalid') is False + + +def test_is_directory() -> None: + assert is_directory(get_test_examples_directory()) is True + assert is_directory(get_test_example_file('source.jpg')) is False + assert is_directory('invalid') is False + + +def test_in_directory() -> None: + assert in_directory(get_test_example_file('source.jpg')) is True + assert in_directory('source.jpg') is False + assert in_directory('invalid') is False diff --git a/tests/test_inference_manager.py b/tests/test_inference_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..58d7dfbb2fc71679c544c06bd63cd3ac5996e072 --- /dev/null +++ b/tests/test_inference_manager.py @@ -0,0 +1,32 @@ +from unittest.mock import patch + +import pytest +from onnxruntime import InferenceSession + +from facefusion import content_analyser, state_manager +from facefusion.inference_manager import INFERENCE_POOL_SET, get_inference_pool + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + state_manager.init_item('execution_device_ids', [ 0 ]) + state_manager.init_item('execution_providers', [ 'cpu' ]) + state_manager.init_item('download_providers', [ 'github' ]) + content_analyser.pre_check() + + +def test_get_inference_pool() -> None: + model_names = [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ] + _, model_source_set = content_analyser.collect_model_downloads() + + with patch('facefusion.inference_manager.detect_app_context', return_value = 'cli'): + get_inference_pool('facefusion.content_analyser', model_names, model_source_set) + + assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.nsfw_1.nsfw_2.nsfw_3.0.cpu').get('nsfw_1'), InferenceSession) + + with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'): + get_inference_pool('facefusion.content_analyser', model_names, model_source_set) + + assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.nsfw_1.nsfw_2.nsfw_3.0.cpu').get('nsfw_1'), InferenceSession) + + assert INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.nsfw_1.nsfw_2.nsfw_3.0.cpu').get('nsfw_1') == INFERENCE_POOL_SET.get('ui').get('facefusion.content_analyser.nsfw_1.nsfw_2.nsfw_3.0.cpu').get('nsfw_1') diff --git a/tests/test_job_helper.py b/tests/test_job_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..f6f798fe8acb5133392e5cf8d16aee1bfbb17d59 --- /dev/null +++ b/tests/test_job_helper.py @@ -0,0 +1,9 @@ +import os + +from facefusion.jobs.job_helper import get_step_output_path + + +def test_get_step_output_path() -> None: + assert get_step_output_path('test-job', 0, 'test.mp4') == 'test-test-job-0.mp4' + assert get_step_output_path('test-job', 0, 'test/test.mp4') == os.path.join('test', 'test-test-job-0.mp4') + assert get_step_output_path('test-job', 0, 'invalid') is None diff --git a/tests/test_job_list.py b/tests/test_job_list.py new file mode 100644 index 0000000000000000000000000000000000000000..3b262b1daf645e6a4caeb5d46a582de36214d149 --- /dev/null +++ b/tests/test_job_list.py @@ -0,0 +1,24 @@ +from time import sleep + +import pytest + +from facefusion.jobs.job_list import compose_job_list +from facefusion.jobs.job_manager import clear_jobs, create_job, init_jobs +from .helper import get_test_jobs_directory + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + + +def test_compose_job_list() -> None: + create_job('job-test-compose-job-list-1') + sleep(0.5) + create_job('job-test-compose-job-list-2') + job_headers, job_contents = compose_job_list('drafted') + + assert job_headers == [ 'job id', 'steps', 'date created', 'date updated', 'job status' ] + assert job_contents[0] == [ 'job-test-compose-job-list-1', 0, 'just now', None, 'drafted' ] + assert job_contents[1] == [ 'job-test-compose-job-list-2', 0, 'just now', None, 'drafted' ] diff --git a/tests/test_job_manager.py b/tests/test_job_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..65c258ba913c6c9d325c6b18ce67335ba75ea14a --- /dev/null +++ b/tests/test_job_manager.py @@ -0,0 +1,386 @@ +from time import sleep + +import pytest + +from facefusion.jobs.job_helper import get_step_output_path +from facefusion.jobs.job_manager import add_step, clear_jobs, count_step_total, create_job, delete_job, delete_jobs, find_job_ids, find_jobs, get_steps, init_jobs, insert_step, move_job_file, remix_step, remove_step, set_step_status, set_steps_status, submit_job, submit_jobs +from .helper import get_test_jobs_directory + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + + +def test_create_job() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + + assert create_job('job-test-create-job') is True + assert create_job('job-test-create-job') is False + + add_step('job-test-submit-job', args_1) + submit_job('job-test-create-job') + + assert create_job('job-test-create-job') is False + + +def test_submit_job() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + + assert submit_job('job-invalid') is False + + create_job('job-test-submit-job') + + assert submit_job('job-test-submit-job') is False + + add_step('job-test-submit-job', args_1) + + assert submit_job('job-test-submit-job') is True + assert submit_job('job-test-submit-job') is False + + +def test_submit_jobs() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + args_2 =\ + { + 'source_path': 'source-2.jpg', + 'target_path': 'target-2.jpg', + 'output_path': 'output-2.jpg' + } + halt_on_error = True + + assert submit_jobs(halt_on_error) is False + + create_job('job-test-submit-jobs-1') + create_job('job-test-submit-jobs-2') + + assert submit_jobs(halt_on_error) is False + + add_step('job-test-submit-jobs-1', args_1) + add_step('job-test-submit-jobs-2', args_2) + + assert submit_jobs(halt_on_error) is True + assert submit_jobs(halt_on_error) is False + + +def test_delete_job() -> None: + assert delete_job('job-invalid') is False + + create_job('job-test-delete-job') + + assert delete_job('job-test-delete-job') is True + assert delete_job('job-test-delete-job') is False + + +def test_delete_jobs() -> None: + halt_on_error = True + + assert delete_jobs(halt_on_error) is False + + create_job('job-test-delete-jobs-1') + create_job('job-test-delete-jobs-2') + + assert delete_jobs(halt_on_error) is True + + +def test_find_jobs() -> None: + create_job('job-test-find-jobs-1') + sleep(0.5) + create_job('job-test-find-jobs-2') + + assert 'job-test-find-jobs-1' in find_jobs('drafted') + assert 'job-test-find-jobs-2' in find_jobs('drafted') + assert not find_jobs('queued') + + move_job_file('job-test-find-jobs-1', 'queued') + + assert 'job-test-find-jobs-2' in find_jobs('drafted') + assert 'job-test-find-jobs-1' in find_jobs('queued') + + +def test_find_job_ids() -> None: + create_job('job-test-find-job-ids-1') + sleep(0.5) + create_job('job-test-find-job-ids-2') + sleep(0.5) + create_job('job-test-find-job-ids-3') + + assert find_job_ids('drafted') == [ 'job-test-find-job-ids-1', 'job-test-find-job-ids-2', 'job-test-find-job-ids-3' ] + assert find_job_ids('queued') == [] + assert find_job_ids('completed') == [] + assert find_job_ids('failed') == [] + + move_job_file('job-test-find-job-ids-1', 'queued') + move_job_file('job-test-find-job-ids-2', 'queued') + move_job_file('job-test-find-job-ids-3', 'queued') + + assert find_job_ids('drafted') == [] + assert find_job_ids('queued') == [ 'job-test-find-job-ids-1', 'job-test-find-job-ids-2', 'job-test-find-job-ids-3' ] + assert find_job_ids('completed') == [] + assert find_job_ids('failed') == [] + + move_job_file('job-test-find-job-ids-1', 'completed') + + assert find_job_ids('drafted') == [] + assert find_job_ids('queued') == [ 'job-test-find-job-ids-2', 'job-test-find-job-ids-3' ] + assert find_job_ids('completed') == [ 'job-test-find-job-ids-1' ] + assert find_job_ids('failed') == [] + + move_job_file('job-test-find-job-ids-2', 'failed') + + assert find_job_ids('drafted') == [] + assert find_job_ids('queued') == [ 'job-test-find-job-ids-3' ] + assert find_job_ids('completed') == [ 'job-test-find-job-ids-1' ] + assert find_job_ids('failed') == [ 'job-test-find-job-ids-2' ] + + move_job_file('job-test-find-job-ids-3', 'completed') + + assert find_job_ids('drafted') == [] + assert find_job_ids('queued') == [] + assert find_job_ids('completed') == [ 'job-test-find-job-ids-1', 'job-test-find-job-ids-3' ] + assert find_job_ids('failed') == [ 'job-test-find-job-ids-2' ] + + +def test_add_step() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + args_2 =\ + { + 'source_path': 'source-2.jpg', + 'target_path': 'target-2.jpg', + 'output_path': 'output-2.jpg' + } + + assert add_step('job-invalid', args_1) is False + + create_job('job-test-add-step') + + assert add_step('job-test-add-step', args_1) is True + assert add_step('job-test-add-step', args_2) is True + + steps = get_steps('job-test-add-step') + + assert steps[0].get('args') == args_1 + assert steps[1].get('args') == args_2 + assert count_step_total('job-test-add-step') == 2 + + +def test_remix_step() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + args_2 =\ + { + 'source_path': 'source-2.jpg', + 'target_path': 'target-2.jpg', + 'output_path': 'output-2.jpg' + } + + assert remix_step('job-invalid', 0, args_1) is False + + create_job('job-test-remix-step') + add_step('job-test-remix-step', args_1) + add_step('job-test-remix-step', args_2) + + assert remix_step('job-test-remix-step', 99, args_1) is False + assert remix_step('job-test-remix-step', 0, args_2) is True + assert remix_step('job-test-remix-step', -1, args_2) is True + + steps = get_steps('job-test-remix-step') + + assert steps[0].get('args') == args_1 + assert steps[1].get('args') == args_2 + assert steps[2].get('args').get('source_path') == args_2.get('source_path') + assert steps[2].get('args').get('target_path') == get_step_output_path('job-test-remix-step', 0, args_1.get('output_path')) + assert steps[2].get('args').get('output_path') == args_2.get('output_path') + assert steps[3].get('args').get('source_path') == args_2.get('source_path') + assert steps[3].get('args').get('target_path') == get_step_output_path('job-test-remix-step', 2, args_2.get('output_path')) + assert steps[3].get('args').get('output_path') == args_2.get('output_path') + assert count_step_total('job-test-remix-step') == 4 + + +def test_insert_step() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + args_2 =\ + { + 'source_path': 'source-2.jpg', + 'target_path': 'target-2.jpg', + 'output_path': 'output-2.jpg' + } + args_3 =\ + { + 'source_path': 'source-3.jpg', + 'target_path': 'target-3.jpg', + 'output_path': 'output-3.jpg' + } + + assert insert_step('job-invalid', 0, args_1) is False + + create_job('job-test-insert-step') + add_step('job-test-insert-step', args_1) + add_step('job-test-insert-step', args_1) + + assert insert_step('job-test-insert-step', 99, args_1) is False + assert insert_step('job-test-insert-step', 0, args_2) is True + assert insert_step('job-test-insert-step', -1, args_3) is True + + steps = get_steps('job-test-insert-step') + + assert steps[0].get('args') == args_2 + assert steps[1].get('args') == args_1 + assert steps[2].get('args') == args_3 + assert steps[3].get('args') == args_1 + assert count_step_total('job-test-insert-step') == 4 + + +def test_remove_step() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + args_2 =\ + { + 'source_path': 'source-2.jpg', + 'target_path': 'target-2.jpg', + 'output_path': 'output-2.jpg' + } + args_3 =\ + { + 'source_path': 'source-3.jpg', + 'target_path': 'target-3.jpg', + 'output_path': 'output-3.jpg' + } + + assert remove_step('job-invalid', 0) is False + + create_job('job-test-remove-step') + add_step('job-test-remove-step', args_1) + add_step('job-test-remove-step', args_2) + add_step('job-test-remove-step', args_1) + add_step('job-test-remove-step', args_3) + + assert remove_step('job-test-remove-step', 99) is False + assert remove_step('job-test-remove-step', 0) is True + assert remove_step('job-test-remove-step', -1) is True + + steps = get_steps('job-test-remove-step') + + assert steps[0].get('args') == args_2 + assert steps[1].get('args') == args_1 + assert count_step_total('job-test-remove-step') == 2 + + +def test_get_steps() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + args_2 =\ + { + 'source_path': 'source-2.jpg', + 'target_path': 'target-2.jpg', + 'output_path': 'output-2.jpg' + } + + assert get_steps('job-invalid') == [] + + create_job('job-test-get-steps') + add_step('job-test-get-steps', args_1) + add_step('job-test-get-steps', args_2) + steps = get_steps('job-test-get-steps') + + assert steps[0].get('args') == args_1 + assert steps[1].get('args') == args_2 + assert count_step_total('job-test-get-steps') == 2 + + +def test_set_step_status() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + args_2 =\ + { + 'source_path': 'source-2.jpg', + 'target_path': 'target-2.jpg', + 'output_path': 'output-2.jpg' + } + + assert set_step_status('job-invalid', 0, 'completed') is False + + create_job('job-test-set-step-status') + add_step('job-test-set-step-status', args_1) + add_step('job-test-set-step-status', args_2) + + assert set_step_status('job-test-set-step-status', 99, 'completed') is False + assert set_step_status('job-test-set-step-status', 0, 'completed') is True + assert set_step_status('job-test-set-step-status', 1, 'failed') is True + + steps = get_steps('job-test-set-step-status') + + assert steps[0].get('status') == 'completed' + assert steps[1].get('status') == 'failed' + assert count_step_total('job-test-set-step-status') == 2 + + +def test_set_steps_status() -> None: + args_1 =\ + { + 'source_path': 'source-1.jpg', + 'target_path': 'target-1.jpg', + 'output_path': 'output-1.jpg' + } + args_2 =\ + { + 'source_path': 'source-2.jpg', + 'target_path': 'target-2.jpg', + 'output_path': 'output-2.jpg' + } + + assert set_steps_status('job-invalid', 'queued') is False + + create_job('job-test-set-steps-status') + add_step('job-test-set-steps-status', args_1) + add_step('job-test-set-steps-status', args_2) + + assert set_steps_status('job-test-set-steps-status', 'queued') is True + + steps = get_steps('job-test-set-steps-status') + + assert steps[0].get('status') == 'queued' + assert steps[1].get('status') == 'queued' + assert count_step_total('job-test-set-steps-status') == 2 diff --git a/tests/test_job_runner.py b/tests/test_job_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..fac36d5a72422299c1bf505cd87ef6272316d7f3 --- /dev/null +++ b/tests/test_job_runner.py @@ -0,0 +1,276 @@ +import subprocess + +import pytest + +from facefusion.download import conditional_download +from facefusion.filesystem import copy_file +from facefusion.jobs.job_manager import add_step, clear_jobs, create_job, init_jobs, move_job_file, submit_job, submit_jobs +from facefusion.jobs.job_runner import collect_output_set, finalize_steps, retry_job, retry_jobs, run_job, run_jobs, run_steps +from facefusion.types import Args +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_jobs(get_test_jobs_directory()) + init_jobs(get_test_jobs_directory()) + prepare_test_output_directory() + + +def process_step(job_id : str, step_index : int, step_args : Args) -> bool: + return copy_file(step_args.get('target_path'), step_args.get('output_path')) + + +def test_run_job() -> None: + args_1 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-1.mp4') + } + args_2 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-2.mp4') + } + args_3 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.jpg'), + 'output_path': get_test_output_file('output-3.jpg') + } + + assert run_job('job-invalid', process_step) is False + + create_job('job-test-run-job') + add_step('job-test-run-job', args_1) + add_step('job-test-run-job', args_2) + add_step('job-test-run-job', args_2) + add_step('job-test-run-job', args_3) + + assert run_job('job-test-run-job', process_step) is False + + submit_job('job-test-run-job') + + assert run_job('job-test-run-job', process_step) is True + + +def test_run_jobs() -> None: + args_1 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-1.mp4') + } + args_2 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-2.mp4') + } + args_3 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.jpg'), + 'output_path': get_test_output_file('output-3.jpg') + } + halt_on_error = True + + assert run_jobs(process_step, halt_on_error) is False + + create_job('job-test-run-jobs-1') + create_job('job-test-run-jobs-2') + add_step('job-test-run-jobs-1', args_1) + add_step('job-test-run-jobs-1', args_1) + add_step('job-test-run-jobs-2', args_2) + add_step('job-test-run-jobs-3', args_3) + + assert run_jobs(process_step, halt_on_error) is False + + submit_jobs(halt_on_error) + + assert run_jobs(process_step, halt_on_error) is True + + +def test_retry_job() -> None: + args_1 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-1.mp4') + } + + assert retry_job('job-invalid', process_step) is False + + create_job('job-test-retry-job') + add_step('job-test-retry-job', args_1) + submit_job('job-test-retry-job') + + assert retry_job('job-test-retry-job', process_step) is False + + move_job_file('job-test-retry-job', 'failed') + + assert retry_job('job-test-retry-job', process_step) is True + + +def test_retry_jobs() -> None: + args_1 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-1.mp4') + } + args_2 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-2.mp4') + } + args_3 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.jpg'), + 'output_path': get_test_output_file('output-3.jpg') + } + halt_on_error = True + + assert retry_jobs(process_step, halt_on_error) is False + + create_job('job-test-retry-jobs-1') + create_job('job-test-retry-jobs-2') + add_step('job-test-retry-jobs-1', args_1) + add_step('job-test-retry-jobs-1', args_1) + add_step('job-test-retry-jobs-2', args_2) + add_step('job-test-retry-jobs-3', args_3) + + assert retry_jobs(process_step, halt_on_error) is False + + move_job_file('job-test-retry-jobs-1', 'failed') + move_job_file('job-test-retry-jobs-2', 'failed') + + assert retry_jobs(process_step, halt_on_error) is True + + +def test_run_steps() -> None: + args_1 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-1.mp4') + } + args_2 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-2.mp4') + } + args_3 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.jpg'), + 'output_path': get_test_output_file('output-3.jpg') + } + + assert run_steps('job-invalid', process_step) is False + + create_job('job-test-run-steps') + add_step('job-test-run-steps', args_1) + add_step('job-test-run-steps', args_1) + add_step('job-test-run-steps', args_2) + add_step('job-test-run-steps', args_3) + + assert run_steps('job-test-run-steps', process_step) is True + + +def test_finalize_steps() -> None: + args_1 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-1.mp4') + } + args_2 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-2.mp4') + } + args_3 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.jpg'), + 'output_path': get_test_output_file('output-3.jpg') + } + + create_job('job-test-finalize-steps') + add_step('job-test-finalize-steps', args_1) + add_step('job-test-finalize-steps', args_1) + add_step('job-test-finalize-steps', args_2) + add_step('job-test-finalize-steps', args_3) + + copy_file(args_1.get('target_path'), get_test_output_file('output-1-job-test-finalize-steps-0.mp4')) + copy_file(args_1.get('target_path'), get_test_output_file('output-1-job-test-finalize-steps-1.mp4')) + copy_file(args_2.get('target_path'), get_test_output_file('output-2-job-test-finalize-steps-2.mp4')) + copy_file(args_3.get('target_path'), get_test_output_file('output-3-job-test-finalize-steps-3.jpg')) + + assert finalize_steps('job-test-finalize-steps') is True + assert is_test_output_file('output-1.mp4') is True + assert is_test_output_file('output-2.mp4') is True + assert is_test_output_file('output-3.jpg') is True + + +def test_collect_output_set() -> None: + args_1 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-1.mp4') + } + args_2 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.mp4'), + 'output_path': get_test_output_file('output-2.mp4') + } + args_3 =\ + { + 'source_path': get_test_example_file('source.jpg'), + 'target_path': get_test_example_file('target-240p.jpg'), + 'output_path': get_test_output_file('output-3.jpg') + } + + create_job('job-test-collect-output-set') + add_step('job-test-collect-output-set', args_1) + add_step('job-test-collect-output-set', args_1) + add_step('job-test-collect-output-set', args_2) + add_step('job-test-collect-output-set', args_3) + + output_set =\ + { + get_test_output_file('output-1.mp4'): + [ + get_test_output_file('output-1-job-test-collect-output-set-0.mp4'), + get_test_output_file('output-1-job-test-collect-output-set-1.mp4') + ], + get_test_output_file('output-2.mp4'): + [ + get_test_output_file('output-2-job-test-collect-output-set-2.mp4') + ], + get_test_output_file('output-3.jpg'): + [ + get_test_output_file('output-3-job-test-collect-output-set-3.jpg') + ] + } + + assert collect_output_set('job-test-collect-output-set') == output_set diff --git a/tests/test_json.py b/tests/test_json.py new file mode 100644 index 0000000000000000000000000000000000000000..4e0b366ddf76ed23adb9ecb916c524b497c5adb6 --- /dev/null +++ b/tests/test_json.py @@ -0,0 +1,19 @@ +import tempfile + +from facefusion.json import read_json, write_json + + +def test_read_json() -> None: + _, json_path = tempfile.mkstemp(suffix = '.json') + + assert not read_json(json_path) + + write_json(json_path, {}) + + assert read_json(json_path) == {} + + +def test_write_json() -> None: + _, json_path = tempfile.mkstemp(suffix = '.json') + + assert write_json(json_path, {}) diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 0000000000000000000000000000000000000000..244a25e6984ef4f1b98943d98f2888632570c9ff --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,8 @@ +from facefusion.common_helper import is_linux, is_macos +from facefusion.memory import limit_system_memory + + +def test_limit_system_memory() -> None: + assert limit_system_memory(4) is True + if is_linux() or is_macos(): + assert limit_system_memory(1024) is False diff --git a/tests/test_normalizer.py b/tests/test_normalizer.py new file mode 100644 index 0000000000000000000000000000000000000000..e15e776d1be55c5000fff89e9acb60ccfb1ac458 --- /dev/null +++ b/tests/test_normalizer.py @@ -0,0 +1,24 @@ +from facefusion.normalizer import normalize_color, normalize_fps, normalize_space + + +def test_normalize_color() -> None: + assert normalize_color([ 0 ]) == (0, 0, 0, 255) + assert normalize_color([ 0, 128 ]) == (0, 128, 0, 255) + assert normalize_color([ 0, 128, 255 ]) == (0, 128, 255, 255) + assert normalize_color([ 0, 128, 255, 0 ]) == (0, 128, 255, 0) + assert normalize_color(None) is None + + +def test_normalize_space() -> None: + assert normalize_space([ 0, 0, 0, 0 ]) == (0, 0, 0, 0) + assert normalize_space([ 1 ]) == (1, 1, 1, 1) + assert normalize_space([ 1, 2 ]) == (1, 2, 1, 2) + assert normalize_space([ 1, 2, 3 ]) == (1, 2, 3, 2) + assert normalize_space(None) is None + + +def test_normalize_fps() -> None: + assert normalize_fps(0.0) == 1.0 + assert normalize_fps(25.0) == 25.0 + assert normalize_fps(61.0) == 60.0 + assert normalize_fps(None) is None diff --git a/tests/test_process_manager.py b/tests/test_process_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..d3887e0f5dd46ecd2fc1a02eeffb69c7ed792e8d --- /dev/null +++ b/tests/test_process_manager.py @@ -0,0 +1,22 @@ +from facefusion.process_manager import end, is_pending, is_processing, is_stopping, set_process_state, start, stop + + +def test_start() -> None: + set_process_state('pending') + start() + + assert is_processing() + + +def test_stop() -> None: + set_process_state('processing') + stop() + + assert is_stopping() + + +def test_end() -> None: + set_process_state('processing') + end() + + assert is_pending() diff --git a/tests/test_program_helper.py b/tests/test_program_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..c6ee214a4ef6e095f6d3ffb10d0ff5354df72ba6 --- /dev/null +++ b/tests/test_program_helper.py @@ -0,0 +1,52 @@ +from argparse import ArgumentParser + +from facefusion.program_helper import find_argument_group, validate_actions, validate_args + + +def test_find_argument_group() -> None: + program = ArgumentParser() + program.add_argument_group('test-1') + program.add_argument_group('test-2') + + assert find_argument_group(program, 'test-1') + assert find_argument_group(program, 'test-2') + assert find_argument_group(program, 'test-3') is None + + +def test_validate_args() -> None: + program = ArgumentParser() + program.add_argument('--test-1', default = 'test_1', choices = [ 'test_1', 'test_2' ]) + + assert validate_args(program) is True + + subparsers = program.add_subparsers() + sub_program = subparsers.add_parser('sub-command') + sub_program.add_argument('--test-2', default = 'test_2', choices = [ 'test_1', 'test_2' ]) + + assert validate_args(program) is True + + for action in sub_program._actions: + if action.dest == 'test_2': + action.default = 'test_3' + + assert validate_args(program) is False + + +def test_validate_actions() -> None: + program = ArgumentParser() + program.add_argument('--test-1', default = 'test_1', choices = [ 'test_1', 'test_2' ]) + program.add_argument('--test-2', default = 'test_2', choices= [ 'test_1', 'test_2' ], nargs = '+') + + assert validate_actions(program) is True + + args =\ + { + 'test_1': 'test_2', + 'test_2': [ 'test_1', 'test_3' ] + } + + for action in program._actions: + if action.dest in args: + action.default = args[action.dest] + + assert validate_actions(program) is False diff --git a/tests/test_sanitizer.py b/tests/test_sanitizer.py new file mode 100644 index 0000000000000000000000000000000000000000..4c1e436c0b2d3ec1e0018fa22ef7da49e3d23214 --- /dev/null +++ b/tests/test_sanitizer.py @@ -0,0 +1,8 @@ +from facefusion.sanitizer import sanitize_int_range + + +def test_sanitize_int_range() -> None: + assert sanitize_int_range(0, [ 0, 1, 2 ]) == 0 + assert sanitize_int_range(2, [0, 1, 2]) == 2 + assert sanitize_int_range(-1, [ 0, 1 ]) == 0 + assert sanitize_int_range(3, [ 0, 1 ]) == 0 diff --git a/tests/test_state_manager.py b/tests/test_state_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..67b4d47bdfd0f5627c1f7784f41c4a1d5f1f6bf1 --- /dev/null +++ b/tests/test_state_manager.py @@ -0,0 +1,35 @@ +from typing import Union + +import pytest + +from facefusion.processors.types import ProcessorState +from facefusion.state_manager import STATE_SET, get_item, init_item, set_item +from facefusion.types import AppContext, State + + +def get_state(app_context : AppContext) -> Union[State, ProcessorState]: + return STATE_SET.get(app_context) + + +def clear_state(app_context : AppContext) -> None: + STATE_SET[app_context] = {} #type:ignore[typeddict-item] + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_state('cli') + clear_state('ui') + + +def test_init_item() -> None: + init_item('video_memory_strategy', 'tolerant') + + assert get_state('cli').get('video_memory_strategy') == 'tolerant' + assert get_state('ui').get('video_memory_strategy') == 'tolerant' + + +def test_get_item_and_set_item() -> None: + set_item('video_memory_strategy', 'tolerant') + + assert get_item('video_memory_strategy') == 'tolerant' + assert get_state('ui').get('video_memory_strategy') is None diff --git a/tests/test_temp_helper.py b/tests/test_temp_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..bb0d3b50c70c9ce42c868020d140320f0c20f9c2 --- /dev/null +++ b/tests/test_temp_helper.py @@ -0,0 +1,34 @@ +import os.path +import tempfile + +import pytest + +from facefusion import state_manager +from facefusion.download import conditional_download +from facefusion.temp_helper import get_temp_directory_path, get_temp_file_path, get_temp_frames_pattern +from .helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + state_manager.init_item('temp_path', tempfile.gettempdir()) + state_manager.init_item('temp_frame_format', 'png') + + +def test_get_temp_file_path() -> None: + temp_directory = tempfile.gettempdir() + assert get_temp_file_path(get_test_example_file('target-240p.mp4')) == os.path.join(temp_directory, 'facefusion', 'target-240p', 'temp.mp4') + + +def test_get_temp_directory_path() -> None: + temp_directory = tempfile.gettempdir() + assert get_temp_directory_path(get_test_example_file('target-240p.mp4')) == os.path.join(temp_directory, 'facefusion', 'target-240p') + + +def test_get_temp_frames_pattern() -> None: + temp_directory = tempfile.gettempdir() + assert get_temp_frames_pattern(get_test_example_file('target-240p.mp4'), '%04d') == os.path.join(temp_directory, 'facefusion', 'target-240p', '%04d.png') diff --git a/tests/test_time_helper.py b/tests/test_time_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..17d87b8400dcf5bedaf8b866fb80f2d88d49726c --- /dev/null +++ b/tests/test_time_helper.py @@ -0,0 +1,15 @@ +from datetime import datetime, timedelta + +from facefusion.time_helper import describe_time_ago + + +def get_time_ago(days : int, hours : int, minutes : int) -> datetime: + previous_time = datetime.now() - timedelta(days = days, hours = hours, minutes = minutes) + return previous_time.astimezone() + + +def test_describe_time_ago() -> None: + assert describe_time_ago(get_time_ago(0, 0, 0)) == 'just now' + assert describe_time_ago(get_time_ago(0, 0, 10)) == '10 minutes ago' + assert describe_time_ago(get_time_ago(0, 5, 10)) == '5 hours and 10 minutes ago' + assert describe_time_ago(get_time_ago(1, 5, 10)) == '1 days, 5 hours and 10 minutes ago' diff --git a/tests/test_translator.py b/tests/test_translator.py new file mode 100644 index 0000000000000000000000000000000000000000..f5fa619cab52a7869c8845ca4d6954f2cddfec04 --- /dev/null +++ b/tests/test_translator.py @@ -0,0 +1,14 @@ +from facefusion import translator +from facefusion.locales import LOCALES + + +def test_load() -> None: + translator.load(LOCALES, __name__) + + assert __name__ in translator.LOCALE_POOL_SET + + +def test_get() -> None: + assert translator.get('conda_not_activated') == 'conda is not activated' + assert translator.get('help.skip_conda') == 'skip the conda environment check' + assert translator.get('invalid') is None diff --git a/tests/test_vision.py b/tests/test_vision.py new file mode 100644 index 0000000000000000000000000000000000000000..bff7e500d41010b058875a0fee625736c9beb62d --- /dev/null +++ b/tests/test_vision.py @@ -0,0 +1,169 @@ +import subprocess + +import pytest + +from facefusion.download import conditional_download +from facefusion.vision import calculate_histogram_difference, count_trim_frame_total, 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_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, unpack_resolution, write_image +from .helper import get_test_example_file, get_test_examples_directory, get_test_output_file, prepare_test_output_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-1080p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('目标-240p.webp') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-1080p.mp4'), '-vframes', '1', get_test_example_file('target-1080p.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', '-vf', 'hue=s=0', get_test_example_file('target-240p-0sat.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', '-vf', 'transpose=0', get_test_example_file('target-240p-90deg.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-1080p.mp4'), '-vframes', '1', '-vf', 'transpose=0', get_test_example_file('target-1080p-90deg.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=25', get_test_example_file('target-240p-25fps.mp4') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=30', get_test_example_file('target-240p-30fps.mp4') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=60', get_test_example_file('target-240p-60fps.mp4') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'transpose=0', get_test_example_file('target-240p-90deg.mp4') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-1080p.mp4'), '-vf', 'transpose=0', get_test_example_file('target-1080p-90deg.mp4') ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + prepare_test_output_directory() + + +def test_read_image() -> None: + assert read_image(get_test_example_file('target-240p.jpg')).shape == (226, 426, 3) + assert read_image(get_test_example_file('目标-240p.webp')).shape == (226, 426, 3) + assert read_image('invalid') is None + + +def test_write_image() -> None: + vision_frame = read_image(get_test_example_file('target-240p.jpg')) + + assert write_image(get_test_output_file('target-240p.jpg'), vision_frame) is True + assert write_image(get_test_output_file('目标-240p.webp'), vision_frame) is True + + +def test_detect_image_resolution() -> None: + assert detect_image_resolution(get_test_example_file('target-240p.jpg')) == (426, 226) + assert detect_image_resolution(get_test_example_file('target-240p-90deg.jpg')) == (226, 426) + assert detect_image_resolution(get_test_example_file('target-1080p.jpg')) == (2048, 1080) + assert detect_image_resolution(get_test_example_file('target-1080p-90deg.jpg')) == (1080, 2048) + assert detect_image_resolution('invalid') is None + + +def test_restrict_image_resolution() -> None: + assert restrict_image_resolution(get_test_example_file('target-1080p.jpg'), (426, 226)) == (426, 226) + assert restrict_image_resolution(get_test_example_file('target-1080p.jpg'), (2048, 1080)) == (2048, 1080) + assert restrict_image_resolution(get_test_example_file('target-1080p.jpg'), (4096, 2160)) == (2048, 1080) + + +def test_read_video_frame() -> None: + assert hasattr(read_video_frame(get_test_example_file('target-240p-25fps.mp4')), '__array_interface__') + assert read_video_frame('invalid') is None + + +def test_count_video_frame_total() -> None: + assert count_video_frame_total(get_test_example_file('target-240p-25fps.mp4')) == 270 + assert count_video_frame_total(get_test_example_file('target-240p-30fps.mp4')) == 324 + assert count_video_frame_total(get_test_example_file('target-240p-60fps.mp4')) == 648 + assert count_video_frame_total('invalid') == 0 + + +def test_predict_video_frame_total() -> None: + assert predict_video_frame_total(get_test_example_file('target-240p-25fps.mp4'), 12.5, 0, 100) == 50 + assert predict_video_frame_total(get_test_example_file('target-240p-25fps.mp4'), 25, 0, 100) == 100 + assert predict_video_frame_total(get_test_example_file('target-240p-25fps.mp4'), 25, 0, 200) == 200 + assert predict_video_frame_total('invalid', 25, 0, 100) == 0 + + +def test_detect_video_fps() -> None: + assert detect_video_fps(get_test_example_file('target-240p-25fps.mp4')) == 25.0 + assert detect_video_fps(get_test_example_file('target-240p-30fps.mp4')) == 30.0 + assert detect_video_fps(get_test_example_file('target-240p-60fps.mp4')) == 60.0 + assert detect_video_fps('invalid') is None + + +def test_restrict_video_fps() -> None: + assert restrict_video_fps(get_test_example_file('target-1080p.mp4'), 20.0) == 20.0 + assert restrict_video_fps(get_test_example_file('target-1080p.mp4'), 25.0) == 25.0 + assert restrict_video_fps(get_test_example_file('target-1080p.mp4'), 60.0) == 25.0 + + +def test_detect_video_duration() -> None: + assert detect_video_duration(get_test_example_file('target-240p.mp4')) == 10.8 + assert detect_video_duration('invalid') == 0 + + +def test_count_trim_frame_total() -> None: + assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), 0, 200) == 200 + assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), 70, 270) == 200 + assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), -10, None) == 270 + assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), None, -10) == 0 + assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), 280, None) == 0 + assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), None, 280) == 270 + assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), None, None) == 270 + + +def test_restrict_trim_frame() -> None: + assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), 0, 200) == (0, 200) + assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), 70, 270) == (70, 270) + assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), -10, None) == (0, 270) + assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), None, -10) == (0, 0) + assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), 280, None) == (270, 270) + assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), None, 280) == (0, 270) + assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), None, None) == (0, 270) + + +def test_detect_video_resolution() -> None: + assert detect_video_resolution(get_test_example_file('target-240p.mp4')) == (426, 226) + assert detect_video_resolution(get_test_example_file('target-240p-90deg.mp4')) == (226, 426) + assert detect_video_resolution(get_test_example_file('target-1080p.mp4')) == (2048, 1080) + assert detect_video_resolution(get_test_example_file('target-1080p-90deg.mp4')) == (1080, 2048) + assert detect_video_resolution('invalid') is None + + +def test_restrict_video_resolution() -> None: + assert restrict_video_resolution(get_test_example_file('target-1080p.mp4'), (426, 226)) == (426, 226) + assert restrict_video_resolution(get_test_example_file('target-1080p.mp4'), (2048, 1080)) == (2048, 1080) + assert restrict_video_resolution(get_test_example_file('target-1080p.mp4'), (4096, 2160)) == (2048, 1080) + + +def test_scale_resolution() -> None: + assert scale_resolution((426, 226), 0.5) == (212, 112) + assert scale_resolution((2048, 1080), 1.0) == (2048, 1080) + assert scale_resolution((4096, 2160), 2.0) == (8192, 4320) + + +def test_normalize_resolution() -> None: + assert normalize_resolution((2.5, 2.5)) == (2, 2) + assert normalize_resolution((3.0, 3.0)) == (4, 4) + assert normalize_resolution((6.5, 6.5)) == (6, 6) + + +def test_pack_resolution() -> None: + assert pack_resolution((1, 1)) == '0x0' + assert pack_resolution((2, 2)) == '2x2' + + +def test_unpack_resolution() -> None: + assert unpack_resolution('0x0') == (0, 0) + assert unpack_resolution('2x2') == (2, 2) + + +def test_calc_histogram_difference() -> None: + source_vision_frame = read_image(get_test_example_file('target-240p.jpg')) + target_vision_frame = read_image(get_test_example_file('target-240p-0sat.jpg')) + + assert calculate_histogram_difference(source_vision_frame, source_vision_frame) == 1.0 + assert calculate_histogram_difference(source_vision_frame, target_vision_frame) < 0.5 + + +def test_match_frame_color() -> None: + source_vision_frame = read_image(get_test_example_file('target-240p.jpg')) + target_vision_frame = read_image(get_test_example_file('target-240p-0sat.jpg')) + output_vision_frame = match_frame_color(source_vision_frame, target_vision_frame) + + assert calculate_histogram_difference(source_vision_frame, output_vision_frame) > 0.5