diff --git a/.flake8 b/.flake8 index 1ca6ddab..eac4a0d7 100644 --- a/.flake8 +++ b/.flake8 @@ -1,3 +1,6 @@ [flake8] -select = E3, E4, F -per-file-ignores = facefusion/core.py:E402 +select = E3, E4, 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/.github/preview.png b/.github/preview.png old mode 100644 new mode 100755 index c311234b..ee316111 Binary files a/.github/preview.png and b/.github/preview.png differ diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9fe3701c..82e1b0cb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,9 +13,12 @@ jobs: with: python-version: '3.10' - run: pip install flake8 + - run: pip install flake8-import-order - run: pip install mypy - - run: flake8 run.py facefusion tests - - run: mypy run.py facefusion tests + - run: flake8 facefusion.py install.py + - run: flake8 facefusion tests + - run: mypy facefusion.py install.py + - run: mypy facefusion tests test: strategy: matrix: @@ -24,7 +27,7 @@ jobs: steps: - name: Checkout uses: actions/checkout@v4 - - name: Set up FFMpeg + - name: Set up FFmpeg uses: FedericoCarboni/setup-ffmpeg@v3 - name: Set up Python 3.10 uses: actions/setup-python@v5 @@ -33,3 +36,23 @@ jobs: - 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.10 + uses: actions/setup-python@v5 + with: + python-version: '3.10' + - 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 index 66381e31..40bebe79 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ .assets +.caches +.jobs .idea .vscode diff --git a/.install/LICENSE.md b/.install/LICENSE.md deleted file mode 100644 index ef02a652..00000000 --- a/.install/LICENSE.md +++ /dev/null @@ -1,3 +0,0 @@ -CC BY-NC license - -Copyright (c) 2024 Henry Ruhs diff --git a/.install/facefusion.ico b/.install/facefusion.ico deleted file mode 100644 index a4fc026b..00000000 Binary files a/.install/facefusion.ico and /dev/null differ diff --git a/.install/facefusion.nsi b/.install/facefusion.nsi deleted file mode 100644 index f840f877..00000000 --- a/.install/facefusion.nsi +++ /dev/null @@ -1,183 +0,0 @@ -!include MUI2.nsh -!include nsDialogs.nsh -!include LogicLib.nsh - -RequestExecutionLevel admin -ManifestDPIAware true - -Name 'FaceFusion 2.6.1' -OutFile 'FaceFusion_2.6.1.exe' - -!define MUI_ICON 'facefusion.ico' - -!insertmacro MUI_PAGE_DIRECTORY -Page custom InstallPage PostInstallPage -!insertmacro MUI_PAGE_INSTFILES -!insertmacro MUI_LANGUAGE English - -Var UseDefault -Var UseCuda -Var UseDirectMl -Var UseOpenVino - -Function .onInit - StrCpy $INSTDIR 'C:\FaceFusion' -FunctionEnd - -Function InstallPage - nsDialogs::Create 1018 - !insertmacro MUI_HEADER_TEXT 'Choose Your Accelerator' 'Choose your accelerator based on the graphics card.' - - ${NSD_CreateRadioButton} 0 40u 100% 10u 'Default' - Pop $UseDefault - - ${NSD_CreateRadioButton} 0 55u 100% 10u 'CUDA (NVIDIA)' - Pop $UseCuda - - ${NSD_CreateRadioButton} 0 70u 100% 10u 'DirectML (AMD, Intel, NVIDIA)' - Pop $UseDirectMl - - ${NSD_CreateRadioButton} 0 85u 100% 10u 'OpenVINO (Intel)' - Pop $UseOpenVino - - ${NSD_Check} $UseDefault - - nsDialogs::Show -FunctionEnd - -Function PostInstallPage - ${NSD_GetState} $UseDefault $UseDefault - ${NSD_GetState} $UseCuda $UseCuda - ${NSD_GetState} $UseDirectMl $UseDirectMl - ${NSD_GetState} $UseOpenVino $UseOpenVino -FunctionEnd - -Function Destroy - ${If} ${Silent} - Quit - ${Else} - Abort - ${EndIf} -FunctionEnd - -Section 'Prepare Your Platform' - DetailPrint 'Install GIT' - inetc::get 'https://github.com/git-for-windows/git/releases/download/v2.45.2.windows.1/Git-2.45.2-64-bit.exe' '$TEMP\Git.exe' - ExecWait '$TEMP\Git.exe /CURRENTUSER /VERYSILENT /DIR=$LOCALAPPDATA\Programs\Git' $0 - Delete '$TEMP\Git.exe' - - ${If} $0 > 0 - DetailPrint 'Git installation aborted with error code $0' - Call Destroy - ${EndIf} - - DetailPrint 'Uninstall Conda' - ExecWait '$LOCALAPPDATA\Programs\Miniconda3\Uninstall-Miniconda3.exe /S _?=$LOCALAPPDATA\Programs\Miniconda3' - RMDir /r '$LOCALAPPDATA\Programs\Miniconda3' - - DetailPrint 'Install Conda' - inetc::get 'https://repo.anaconda.com/miniconda/Miniconda3-py310_24.3.0-0-Windows-x86_64.exe' '$TEMP\Miniconda3.exe' - ExecWait '$TEMP\Miniconda3.exe /InstallationType=JustMe /AddToPath=1 /S /D=$LOCALAPPDATA\Programs\Miniconda3' $1 - Delete '$TEMP\Miniconda3.exe' - - ${If} $1 > 0 - DetailPrint 'Conda installation aborted with error code $1' - Call Destroy - ${EndIf} -SectionEnd - -Section 'Download Your Copy' - SetOutPath $INSTDIR - - DetailPrint 'Download Your Copy' - RMDir /r $INSTDIR - nsExec::Exec '$LOCALAPPDATA\Programs\Git\cmd\git.exe clone https://github.com/facefusion/facefusion --branch 2.6.1 .' -SectionEnd - -Section 'Setup Your Environment' - DetailPrint 'Setup Your Environment' - nsExec::Exec '$LOCALAPPDATA\Programs\Miniconda3\Scripts\conda.exe init --all' - nsExec::Exec '$LOCALAPPDATA\Programs\Miniconda3\Scripts\conda.exe create --name facefusion python=3.10 --yes' -SectionEnd - -Section 'Create Install Batch' - SetOutPath $INSTDIR - - FileOpen $0 install-ffmpeg.bat w - FileOpen $1 install-accelerator.bat w - FileOpen $2 install-application.bat w - - FileWrite $0 '@echo off && conda activate facefusion && conda install conda-forge::ffmpeg=7.0.1 --yes' - ${If} $UseCuda == 1 - FileWrite $1 '@echo off && conda activate facefusion && conda install cudatoolkit=11.8 cudnn=8.9.2.26 conda-forge::gputil=1.4.0 conda-forge::zlib-wapi --yes' - FileWrite $2 '@echo off && conda activate facefusion && python install.py --onnxruntime cuda-11.8' - ${ElseIf} $UseDirectMl == 1 - FileWrite $2 '@echo off && conda activate facefusion && python install.py --onnxruntime directml' - ${ElseIf} $UseOpenVino == 1 - FileWrite $1 '@echo off && conda activate facefusion && conda install conda-forge::openvino=2023.1.0 --yes' - FileWrite $2 '@echo off && conda activate facefusion && python install.py --onnxruntime openvino' - ${Else} - FileWrite $2 '@echo off && conda activate facefusion && python install.py --onnxruntime default' - ${EndIf} - - FileClose $0 - FileClose $1 - FileClose $2 -SectionEnd - -Section 'Install Your FFmpeg' - SetOutPath $INSTDIR - - DetailPrint 'Install Your FFmpeg' - nsExec::ExecToLog 'install-ffmpeg.bat' -SectionEnd - -Section 'Install Your Accelerator' - SetOutPath $INSTDIR - - DetailPrint 'Install Your Accelerator' - nsExec::ExecToLog 'install-accelerator.bat' -SectionEnd - -Section 'Install The Application' - SetOutPath $INSTDIR - - DetailPrint 'Install The Application' - nsExec::ExecToLog 'install-application.bat' -SectionEnd - -Section 'Create Run Batch' - SetOutPath $INSTDIR - FileOpen $0 run.bat w - FileWrite $0 '@echo off && conda activate facefusion && python run.py %*' - FileClose $0 -SectionEnd - -Section 'Register The Application' - DetailPrint 'Register The Application' - - CreateDirectory $SMPROGRAMS\FaceFusion - CreateShortcut '$SMPROGRAMS\FaceFusion\FaceFusion.lnk' $INSTDIR\run.bat '--open-browser' $INSTDIR\.install\facefusion.ico - CreateShortcut '$SMPROGRAMS\FaceFusion\FaceFusion Benchmark.lnk' $INSTDIR\run.bat '--ui-layouts benchmark --open-browser' $INSTDIR\.install\facefusion.ico - CreateShortcut '$SMPROGRAMS\FaceFusion\FaceFusion Webcam.lnk' $INSTDIR\run.bat '--ui-layouts webcam --open-browser' $INSTDIR\.install\facefusion.ico - - CreateShortcut $DESKTOP\FaceFusion.lnk $INSTDIR\run.bat '--open-browser' $INSTDIR\.install\facefusion.ico - - WriteUninstaller $INSTDIR\Uninstall.exe - - WriteRegStr HKLM SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\FaceFusion DisplayName 'FaceFusion' - WriteRegStr HKLM SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\FaceFusion DisplayVersion '2.6.0' - WriteRegStr HKLM SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\FaceFusion Publisher 'Henry Ruhs' - WriteRegStr HKLM SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\FaceFusion InstallLocation $INSTDIR - WriteRegStr HKLM SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\FaceFusion UninstallString $INSTDIR\uninstall.exe -SectionEnd - -Section 'Uninstall' - nsExec::Exec '$LOCALAPPDATA\Programs\Miniconda3\Scripts\conda.exe env remove --name facefusion --yes' - - Delete $DESKTOP\FaceFusion.lnk - RMDir /r $SMPROGRAMS\FaceFusion - RMDir /r $INSTDIR - - DeleteRegKey HKLM SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\FaceFusion -SectionEnd diff --git a/README.md b/README.md index 66fd662f..a524ae9d 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,10 @@ FaceFusion ========== -> Next generation face swapper and enhancer. +> Industry leading face manipulation platform. [![Build Status](https://img.shields.io/github/actions/workflow/status/facefusion/facefusion/ci.yml.svg?branch=master)](https://github.com/facefusion/facefusion/actions?query=workflow:ci) +[![Coverage Status](https://img.shields.io/coveralls/facefusion/facefusion.svg)](https://coveralls.io/r/facefusion/facefusion) ![License](https://img.shields.io/badge/license-MIT-green) @@ -16,7 +17,7 @@ Preview Installation ------------ -Be aware, the [installation](https://docs.facefusion.io/installation) needs technical skills and is not recommended for beginners. In case you are not comfortable using a terminal, our [Windows Installer](https://buymeacoffee.com/henryruhs/e/251939) can have you up and running in minutes. +Be aware, the [installation](https://docs.facefusion.io/installation) needs technical skills and is not recommended for beginners. In case you are not comfortable using a terminal, our [Windows Installer](https://windows-installer.facefusion.io) and [macOS Installer](https://macos-installer.facefusion.io) get you started. Usage @@ -25,85 +26,30 @@ Usage Run the command: ``` -python run.py [options] +python facefusion.py [commands] [options] options: - -h, --help show this help message and exit - -c CONFIG_PATH, --config CONFIG_PATH choose the config file to override defaults - -s SOURCE_PATHS, --source SOURCE_PATHS choose single or multiple source images or audios - -t TARGET_PATH, --target TARGET_PATH choose single target image or video - -o OUTPUT_PATH, --output OUTPUT_PATH specify the output file or directory - -v, --version show program's version number and exit + -h, --help show this help message and exit + -v, --version show program's version number and exit -misc: - --force-download force automate downloads and exit - --skip-download omit automate downloads and remote lookups - --headless run the program without a user interface - --log-level {error,warn,info,debug} adjust the message severity displayed in the terminal - -execution: - --execution-device-id EXECUTION_DEVICE_ID specify the device used for processing - --execution-providers EXECUTION_PROVIDERS [EXECUTION_PROVIDERS ...] accelerate the model inference using different providers (choices: cpu, ...) - --execution-thread-count [1-128] specify the amount of parallel threads while processing - --execution-queue-count [1-32] specify the amount of frames each thread is processing - -memory: - --video-memory-strategy {strict,moderate,tolerant} balance fast frame processing and low VRAM usage - --system-memory-limit [0-128] limit the available RAM that can be used while processing - -face analyser: - --face-analyser-order {left-right,right-left,top-bottom,bottom-top,small-large,large-small,best-worst,worst-best} specify the order in which the face analyser detects faces - --face-analyser-age {child,teen,adult,senior} filter the detected faces based on their age - --face-analyser-gender {female,male} filter the detected faces based on their gender - --face-detector-model {many,retinaface,scrfd,yoloface,yunet} choose the model responsible for detecting the face - --face-detector-size FACE_DETECTOR_SIZE specify the size of the frame provided to the face detector - --face-detector-score [0.0-0.95] filter the detected faces base on the confidence score - --face-landmarker-score [0.0-0.95] filter the detected landmarks base on the confidence score - -face selector: - --face-selector-mode {many,one,reference} use reference based tracking or simple matching - --reference-face-position REFERENCE_FACE_POSITION specify the position used to create the reference face - --reference-face-distance [0.0-1.45] specify the desired similarity between the reference face and target face - --reference-frame-number REFERENCE_FRAME_NUMBER specify the frame used to create the reference face - -face mask: - --face-mask-types FACE_MASK_TYPES [FACE_MASK_TYPES ...] mix and match different face mask types (choices: box, occlusion, region) - --face-mask-blur [0.0-0.95] specify the degree of blur applied the box mask - --face-mask-padding FACE_MASK_PADDING [FACE_MASK_PADDING ...] apply top, right, bottom and left padding to the box mask - --face-mask-regions FACE_MASK_REGIONS [FACE_MASK_REGIONS ...] choose the facial features used for the region mask (choices: skin, left-eyebrow, right-eyebrow, left-eye, right-eye, glasses, nose, mouth, upper-lip, lower-lip) - -frame extraction: - --trim-frame-start TRIM_FRAME_START specify the the start frame of the target video - --trim-frame-end TRIM_FRAME_END specify the the end frame of the target video - --temp-frame-format {bmp,jpg,png} specify the temporary resources format - --keep-temp keep the temporary resources after processing - -output creation: - --output-image-quality [0-100] specify the image quality which translates to the compression factor - --output-image-resolution OUTPUT_IMAGE_RESOLUTION specify the image output resolution based on the target image - --output-video-encoder {libx264,libx265,libvpx-vp9,h264_nvenc,hevc_nvenc,h264_amf,hevc_amf} specify the encoder use for the video compression - --output-video-preset {ultrafast,superfast,veryfast,faster,fast,medium,slow,slower,veryslow} balance fast video processing and video file size - --output-video-quality [0-100] specify the video quality which translates to the compression factor - --output-video-resolution OUTPUT_VIDEO_RESOLUTION specify the video output resolution based on the target video - --output-video-fps OUTPUT_VIDEO_FPS specify the video output fps based on the target video - --skip-audio omit the audio from the target video - -frame processors: - --frame-processors FRAME_PROCESSORS [FRAME_PROCESSORS ...] load a single or multiple frame processors. (choices: face_debugger, face_enhancer, face_swapper, frame_colorizer, frame_enhancer, lip_syncer, ...) - --face-debugger-items FACE_DEBUGGER_ITEMS [FACE_DEBUGGER_ITEMS ...] load a single or multiple frame processors (choices: bounding-box, face-landmark-5, face-landmark-5/68, face-landmark-68, face-landmark-68/5, face-mask, face-detector-score, face-landmarker-score, age, gender) - --face-enhancer-model {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} choose the model responsible for enhancing the face - --face-enhancer-blend [0-100] blend the enhanced into the previous face - --face-swapper-model {blendswap_256,inswapper_128,inswapper_128_fp16,simswap_256,simswap_512_unofficial,uniface_256} choose the model responsible for swapping the face - --frame-colorizer-model {ddcolor,ddcolor_artistic,deoldify,deoldify_artistic,deoldify_stable} choose the model responsible for colorizing the frame - --frame-colorizer-blend [0-100] blend the colorized into the previous frame - --frame-colorizer-size {192x192,256x256,384x384,512x512} specify the size of the frame provided to the frame colorizer - --frame-enhancer-model {clear_reality_x4,lsdir_x4,nomos8k_sc_x4,real_esrgan_x2,real_esrgan_x2_fp16,real_esrgan_x4,real_esrgan_x4_fp16,real_hatgan_x4,span_kendata_x4,ultra_sharp_x4} choose the model responsible for enhancing the frame - --frame-enhancer-blend [0-100] blend the enhanced into the previous frame - --lip-syncer-model {wav2lip_gan} choose the model responsible for syncing the lips - -uis: - --open-browser open the browser once the program is ready - --ui-layouts UI_LAYOUTS [UI_LAYOUTS ...] launch a single or multiple UI layouts (choices: benchmark, default, webcam, ...) +commands: + run run the program + headless-run run the program in headless mode + force-download force automate downloads and exit + 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-list list jobs by status + 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 ``` diff --git a/facefusion.ico b/facefusion.ico new file mode 100755 index 00000000..24ce6004 Binary files /dev/null and b/facefusion.ico differ diff --git a/facefusion.ini b/facefusion.ini index dcb08f29..bba59217 100644 --- a/facefusion.ini +++ b/facefusion.ini @@ -1,40 +1,31 @@ -[general] +[paths] +jobs_path = source_paths = target_path = output_path = -[misc] -force_download = -skip_download = -headless = -log_level = - -[execution] -execution_device_id = -execution_providers = -execution_thread_count = -execution_queue_count = - -[memory] -video_memory_strategy = -system_memory_limit = - -[face_analyser] -face_analyser_order = -face_analyser_age = -face_analyser_gender = +[face_detector] face_detector_model = +face_detector_angles = face_detector_size = face_detector_score = + +[face_landmarker] +face_landmarker_model = face_landmarker_score = [face_selector] 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_mask] +[face_masker] face_mask_types = face_mask_blur = face_mask_padding = @@ -49,6 +40,7 @@ keep_temp = [output_creation] output_image_quality = output_image_resolution = +output_audio_encoder = output_video_encoder = output_video_preset = output_video_quality = @@ -56,12 +48,32 @@ output_video_resolution = output_video_fps = skip_audio = -[frame_processors] -frame_processors = +[processors] +processors = +age_modifier_model = +age_modifier_direction = +expression_restorer_model = +expression_restorer_factor = 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_swapper_model = +face_swapper_pixel_boost = frame_colorizer_model = frame_colorizer_blend = frame_colorizer_size = @@ -72,3 +84,18 @@ lip_syncer_model = [uis] open_browser = ui_layouts = +ui_workflow = + +[execution] +execution_device_id = +execution_providers = +execution_thread_count = +execution_queue_count = + +[memory] +video_memory_strategy = +system_memory_limit = + +[misc] +skip_download = +log_level = diff --git a/run.py b/facefusion.py similarity index 65% rename from run.py rename to facefusion.py index 1c2a8bd9..98a865c7 100755 --- a/run.py +++ b/facefusion.py @@ -1,5 +1,9 @@ #!/usr/bin/env python3 +import os + +os.environ['OMP_NUM_THREADS'] = '1' + from facefusion import core if __name__ == '__main__': diff --git a/facefusion/app_context.py b/facefusion/app_context.py new file mode 100644 index 00000000..f1a273af --- /dev/null +++ b/facefusion/app_context.py @@ -0,0 +1,16 @@ +import os +import sys + +from facefusion.typing 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 00000000..8ab29d34 --- /dev/null +++ b/facefusion/args.py @@ -0,0 +1,118 @@ +from facefusion import state_manager +from facefusion.filesystem import is_image, is_video, list_directory +from facefusion.jobs import job_store +from facefusion.normalizer import normalize_fps, normalize_padding +from facefusion.processors.core import get_processors_modules +from facefusion.typing import ApplyStateItem, Args +from facefusion.vision import create_image_resolutions, create_video_resolutions, detect_image_resolution, detect_video_fps, detect_video_resolution, pack_resolution + + +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 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('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')) + # 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_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 + state_manager.init_item('face_selector_mode', args.get('face_selector_mode')) + state_manager.init_item('face_selector_order', args.get('face_selector_order')) + state_manager.init_item('face_selector_gender', args.get('face_selector_gender')) + state_manager.init_item('face_selector_race', args.get('face_selector_race')) + state_manager.init_item('face_selector_age_start', args.get('face_selector_age_start')) + state_manager.init_item('face_selector_age_end', args.get('face_selector_age_end')) + state_manager.init_item('reference_face_position', args.get('reference_face_position')) + state_manager.init_item('reference_face_distance', args.get('reference_face_distance')) + state_manager.init_item('reference_frame_number', args.get('reference_frame_number')) + # face masker + apply_state_item('face_mask_types', args.get('face_mask_types')) + apply_state_item('face_mask_blur', args.get('face_mask_blur')) + apply_state_item('face_mask_padding', normalize_padding(args.get('face_mask_padding'))) + apply_state_item('face_mask_regions', args.get('face_mask_regions')) + # 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')) + if is_image(args.get('target_path')): + output_image_resolution = detect_image_resolution(args.get('target_path')) + output_image_resolutions = create_image_resolutions(output_image_resolution) + if args.get('output_image_resolution') in output_image_resolutions: + apply_state_item('output_image_resolution', args.get('output_image_resolution')) + else: + apply_state_item('output_image_resolution', pack_resolution(output_image_resolution)) + apply_state_item('output_audio_encoder', args.get('output_audio_encoder')) + 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')) + if is_video(args.get('target_path')): + output_video_resolution = detect_video_resolution(args.get('target_path')) + output_video_resolutions = create_video_resolutions(output_video_resolution) + if args.get('output_video_resolution') in output_video_resolutions: + apply_state_item('output_video_resolution', args.get('output_video_resolution')) + else: + apply_state_item('output_video_resolution', pack_resolution(output_video_resolution)) + 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) + apply_state_item('skip_audio', args.get('skip_audio')) + # processors + available_processors = list_directory('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 + if args.get('command') == 'run': + 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_id', args.get('execution_device_id')) + apply_state_item('execution_providers', args.get('execution_providers')) + apply_state_item('execution_thread_count', args.get('execution_thread_count')) + apply_state_item('execution_queue_count', args.get('execution_queue_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('skip_download', args.get('skip_download')) + apply_state_item('log_level', args.get('log_level')) + # 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 index de800502..abe52c88 100644 --- a/facefusion/audio.py +++ b/facefusion/audio.py @@ -1,11 +1,13 @@ -from typing import Optional, Any, List from functools import lru_cache +from typing import Any, List, Optional + import numpy import scipy +from numpy._typing import NDArray -from facefusion.filesystem import is_audio from facefusion.ffmpeg import read_audio_buffer -from facefusion.typing import Fps, Audio, AudioFrame, Spectrogram, MelFilterBank +from facefusion.filesystem import is_audio +from facefusion.typing import Audio, AudioFrame, Fps, Mel, MelFilterBank, Spectrogram from facefusion.voice_extractor import batch_extract_voice @@ -36,8 +38,8 @@ def read_static_voice(audio_path : str, fps : Fps) -> Optional[List[AudioFrame]] def read_voice(audio_path : str, fps : Fps) -> Optional[List[AudioFrame]]: sample_rate = 48000 channel_total = 2 - chunk_size = 1024 * 240 - step_size = 1024 * 180 + chunk_size = 240 * 1024 + step_size = 180 * 1024 if is_audio(audio_path): audio_buffer = read_audio_buffer(audio_path, sample_rate, channel_total) @@ -73,7 +75,7 @@ def create_empty_audio_frame() -> AudioFrame: return audio_frame -def prepare_audio(audio : numpy.ndarray[Any, Any]) -> Audio: +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) @@ -81,7 +83,7 @@ def prepare_audio(audio : numpy.ndarray[Any, Any]) -> Audio: return audio -def prepare_voice(audio : numpy.ndarray[Any, Any]) -> Audio: +def prepare_voice(audio : Audio) -> Audio: sample_rate = 48000 resample_rate = 16000 @@ -94,7 +96,7 @@ def convert_hertz_to_mel(hertz : float) -> float: return 2595 * numpy.log10(1 + hertz / 700) -def convert_mel_to_hertz(mel : numpy.ndarray[Any, Any]) -> numpy.ndarray[Any, Any]: +def convert_mel_to_hertz(mel : Mel) -> NDArray[Any]: return 700 * (10 ** (mel / 2595) - 1) diff --git a/facefusion/choices.py b/facefusion/choices.py index e5587b85..15a8cc5e 100755 --- a/facefusion/choices.py +++ b/facefusion/choices.py @@ -1,37 +1,64 @@ -from typing import List, Dict +import logging +from typing import List, Sequence -from facefusion.typing import VideoMemoryStrategy, FaceSelectorMode, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, FaceDetectorModel, FaceMaskType, FaceMaskRegion, TempFrameFormat, OutputVideoEncoder, OutputVideoPreset -from facefusion.common_helper import create_int_range, create_float_range +from facefusion.common_helper import create_float_range, create_int_range +from facefusion.typing import Angle, ExecutionProviderSet, FaceDetectorSet, FaceLandmarkerModel, FaceMaskRegion, FaceMaskType, FaceSelectorMode, FaceSelectorOrder, Gender, JobStatus, LogLevelSet, OutputAudioEncoder, OutputVideoEncoder, OutputVideoPreset, Race, Score, TempFrameFormat, UiWorkflow, VideoMemoryStrategy video_memory_strategies : List[VideoMemoryStrategy] = [ 'strict', 'moderate', 'tolerant' ] -face_analyser_orders : List[FaceAnalyserOrder] = [ 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best' ] -face_analyser_ages : List[FaceAnalyserAge] = [ 'child', 'teen', 'adult', 'senior' ] -face_analyser_genders : List[FaceAnalyserGender] = [ 'female', 'male' ] -face_detector_set : Dict[FaceDetectorModel, List[str]] =\ + +face_detector_set : FaceDetectorSet =\ { 'many': [ '640x640' ], 'retinaface': [ '160x160', '320x320', '480x480', '512x512', '640x640' ], 'scrfd': [ '160x160', '320x320', '480x480', '512x512', '640x640' ], - 'yoloface': [ '640x640' ], - 'yunet': [ '160x160', '320x320', '480x480', '512x512', '640x640', '768x768', '960x960', '1024x1024' ] + 'yoloface': [ '640x640' ] } +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_mask_types : List[FaceMaskType] = [ 'box', 'occlusion', 'region' ] face_mask_regions : List[FaceMaskRegion] = [ 'skin', 'left-eyebrow', 'right-eyebrow', 'left-eye', 'right-eye', 'glasses', 'nose', 'mouth', 'upper-lip', 'lower-lip' ] temp_frame_formats : List[TempFrameFormat] = [ 'bmp', 'jpg', 'png' ] -output_video_encoders : List[OutputVideoEncoder] = [ 'libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc', 'h264_amf', 'hevc_amf' ] +output_audio_encoders : List[OutputAudioEncoder] = [ 'aac', 'libmp3lame', 'libopus', 'libvorbis' ] +output_video_encoders : List[OutputVideoEncoder] = [ 'libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc', 'h264_amf', 'hevc_amf', 'h264_videotoolbox', 'hevc_videotoolbox' ] output_video_presets : List[OutputVideoPreset] = [ 'ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow' ] image_template_sizes : List[float] = [ 0.25, 0.5, 0.75, 1, 1.5, 2, 2.5, 3, 3.5, 4 ] video_template_sizes : List[int] = [ 240, 360, 480, 540, 720, 1080, 1440, 2160, 4320 ] -execution_thread_count_range : List[int] = create_int_range(1, 128, 1) -execution_queue_count_range : List[int] = create_int_range(1, 32, 1) -system_memory_limit_range : List[int] = create_int_range(0, 128, 1) -face_detector_score_range : List[float] = create_float_range(0.0, 1.0, 0.05) -face_landmarker_score_range : List[float] = create_float_range(0.0, 1.0, 0.05) -face_mask_blur_range : List[float] = create_float_range(0.0, 1.0, 0.05) -face_mask_padding_range : List[int] = create_int_range(0, 100, 1) -reference_face_distance_range : List[float] = create_float_range(0.0, 1.5, 0.05) -output_image_quality_range : List[int] = create_int_range(0, 100, 1) -output_video_quality_range : List[int] = create_int_range(0, 100, 1) +log_level_set : LogLevelSet =\ +{ + 'error': logging.ERROR, + 'warn': logging.WARNING, + 'info': logging.INFO, + 'debug': logging.DEBUG +} + +execution_provider_set : ExecutionProviderSet =\ +{ + 'cpu': 'CPUExecutionProvider', + 'coreml': 'CoreMLExecutionProvider', + 'cuda': 'CUDAExecutionProvider', + 'directml': 'DmlExecutionProvider', + 'openvino': 'OpenVINOExecutionProvider', + 'rocm': 'ROCMExecutionProvider', + 'tensorrt': 'TensorrtExecutionProvider' +} + +ui_workflows : List[UiWorkflow] = [ 'instant_runner', 'job_runner', 'job_manager' ] +job_statuses : List[JobStatus] = [ 'drafted', 'queued', 'completed', 'failed' ] + +execution_thread_count_range : Sequence[int] = create_int_range(1, 32, 1) +execution_queue_count_range : Sequence[int] = create_int_range(1, 4, 1) +system_memory_limit_range : Sequence[int] = create_int_range(0, 128, 4) +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.5, 0.05) +output_image_quality_range : Sequence[int] = create_int_range(0, 100, 1) +output_video_quality_range : Sequence[int] = create_int_range(0, 100, 1) diff --git a/facefusion/common_helper.py b/facefusion/common_helper.py index d37419bc..e8d9e4ba 100644 --- a/facefusion/common_helper.py +++ b/facefusion/common_helper.py @@ -1,12 +1,28 @@ -from typing import List, Any import platform +from typing import Any, Sequence -def create_metavar(ranges : List[Any]) -> str: - return '[' + str(ranges[0]) + '-' + str(ranges[-1]) + ']' +def is_linux() -> bool: + return platform.system().lower() == 'linux' -def create_int_range(start : int, end : int, step : int) -> List[int]: +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(calc_int_step(int_range)) + ']' + + +def create_float_metavar(float_range : Sequence[float]) -> str: + return '[' + str(float_range[0]) + '..' + str(float_range[-1]) + ':' + str(calc_float_step(float_range)) + ']' + + +def create_int_range(start : int, end : int, step : int) -> Sequence[int]: int_range = [] current = start @@ -16,7 +32,7 @@ def create_int_range(start : int, end : int, step : int) -> List[int]: return int_range -def create_float_range(start : float, end : float, step : float) -> List[float]: +def create_float_range(start : float, end : float, step : float) -> Sequence[float]: float_range = [] current = start @@ -26,21 +42,17 @@ def create_float_range(start : float, end : float, step : float) -> List[float]: return float_range -def is_linux() -> bool: - return to_lower_case(platform.system()) == 'linux' +def calc_int_step(int_range : Sequence[int]) -> int: + return int_range[1] - int_range[0] -def is_macos() -> bool: - return to_lower_case(platform.system()) == 'darwin' - - -def is_windows() -> bool: - return to_lower_case(platform.system()) == 'windows' - - -def to_lower_case(__string__ : Any) -> str: - return str(__string__).lower() +def calc_float_step(float_range : Sequence[float]) -> float: + return round(float_range[1] - float_range[0], 2) def get_first(__list__ : Any) -> Any: return next(iter(__list__), None) + + +def get_last(__list__ : Any) -> Any: + return next(reversed(__list__), None) diff --git a/facefusion/config.py b/facefusion/config.py index c8407647..928052b0 100644 --- a/facefusion/config.py +++ b/facefusion/config.py @@ -1,7 +1,7 @@ from configparser import ConfigParser -from typing import Any, Optional, List +from typing import Any, List, Optional -import facefusion.globals +from facefusion import state_manager CONFIG = None @@ -11,7 +11,7 @@ def get_config() -> ConfigParser: if CONFIG is None: CONFIG = ConfigParser() - CONFIG.read(facefusion.globals.config_path, encoding = 'utf-8') + CONFIG.read(state_manager.get_item('config_path'), encoding = 'utf-8') return CONFIG diff --git a/facefusion/content_analyser.py b/facefusion/content_analyser.py index 5cc5a184..998e560c 100644 --- a/facefusion/content_analyser.py +++ b/facefusion/content_analyser.py @@ -1,27 +1,38 @@ -from typing import Any from functools import lru_cache -from time import sleep + import cv2 import numpy -import onnxruntime from tqdm import tqdm -import facefusion.globals -from facefusion import process_manager, wording -from facefusion.thread_helper import thread_lock, conditional_thread_semaphore -from facefusion.typing import VisionFrame, ModelSet, Fps -from facefusion.execution import apply_execution_provider_options -from facefusion.vision import get_video_frame, count_video_frame_total, read_image, detect_video_fps -from facefusion.filesystem import resolve_relative_path, is_file -from facefusion.download import conditional_download +from facefusion import inference_manager, state_manager, wording +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.filesystem import resolve_relative_path +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.typing import Fps, InferencePool, ModelOptions, ModelSet, VisionFrame +from facefusion.vision import count_video_frame_total, detect_video_fps, get_video_frame, read_image -CONTENT_ANALYSER = None -MODELS : ModelSet =\ +MODEL_SET : ModelSet =\ { 'open_nsfw': { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/open_nsfw.onnx', - 'path': resolve_relative_path('../.assets/models/open_nsfw.onnx') + 'hashes': + { + 'content_analyser': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/open_nsfw.hash', + 'path': resolve_relative_path('../.assets/models/open_nsfw.hash') + } + }, + 'sources': + { + 'content_analyser': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/open_nsfw.onnx', + 'path': resolve_relative_path('../.assets/models/open_nsfw.onnx') + } + }, + 'size': (224, 224), + 'mean': [ 104, 117, 123 ] } } PROBABILITY_LIMIT = 0.80 @@ -29,34 +40,25 @@ RATE_LIMIT = 10 STREAM_COUNTER = 0 -def get_content_analyser() -> Any: - global CONTENT_ANALYSER - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if CONTENT_ANALYSER is None: - model_path = MODELS.get('open_nsfw').get('path') - CONTENT_ANALYSER = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return CONTENT_ANALYSER +def get_inference_pool() -> InferencePool: + model_sources = get_model_options().get('sources') + return inference_manager.get_inference_pool(__name__, model_sources) -def clear_content_analyser() -> None: - global CONTENT_ANALYSER +def clear_inference_pool() -> None: + inference_manager.clear_inference_pool(__name__) - CONTENT_ANALYSER = None + +def get_model_options() -> ModelOptions: + return MODEL_SET.get('open_nsfw') def pre_check() -> bool: download_directory_path = resolve_relative_path('../.assets/models') - model_url = MODELS.get('open_nsfw').get('url') - model_path = MODELS.get('open_nsfw').get('path') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, [ model_url ]) - process_manager.end() - return is_file(model_path) + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) def analyse_stream(vision_frame : VisionFrame, video_fps : Fps) -> bool: @@ -69,19 +71,29 @@ def analyse_stream(vision_frame : VisionFrame, video_fps : Fps) -> bool: def analyse_frame(vision_frame : VisionFrame) -> bool: - content_analyser = get_content_analyser() vision_frame = prepare_frame(vision_frame) - with conditional_thread_semaphore(facefusion.globals.execution_providers): - probability = content_analyser.run(None, - { - content_analyser.get_inputs()[0].name: vision_frame - })[0][0][1] + probability = forward(vision_frame) + return probability > PROBABILITY_LIMIT +def forward(vision_frame : VisionFrame) -> float: + content_analyser = get_inference_pool().get('content_analyser') + + with conditional_thread_semaphore(): + probability = content_analyser.run(None, + { + 'input': vision_frame + })[0][0][1] + + return probability + + def prepare_frame(vision_frame : VisionFrame) -> VisionFrame: - vision_frame = cv2.resize(vision_frame, (224, 224)).astype(numpy.float32) - vision_frame -= numpy.array([ 104, 117, 123 ]).astype(numpy.float32) + model_size = get_model_options().get('size') + model_mean = get_model_options().get('mean') + vision_frame = cv2.resize(vision_frame, model_size).astype(numpy.float32) + vision_frame -= numpy.array(model_mean).astype(numpy.float32) vision_frame = numpy.expand_dims(vision_frame, axis = 0) return vision_frame @@ -100,7 +112,7 @@ def analyse_video(video_path : str, start_frame : int, end_frame : int) -> bool: rate = 0.0 counter = 0 - with tqdm(total = len(frame_range), desc = wording.get('analysing'), unit = 'frame', ascii = ' =', disable = facefusion.globals.log_level in [ 'warn', 'error' ]) as progress: + with tqdm(total = len(frame_range), desc = wording.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: frame = get_video_frame(video_path, frame_number) diff --git a/facefusion/core.py b/facefusion/core.py index e69139c7..3bf2033f 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -1,437 +1,445 @@ -import os - -os.environ['OMP_NUM_THREADS'] = '1' - +import shutil import signal import sys -import warnings -import shutil +from time import time + import numpy -import onnxruntime -from time import sleep, time -from argparse import ArgumentParser, HelpFormatter -import facefusion.choices -import facefusion.globals -from facefusion.face_analyser import get_one_face, get_average_face -from facefusion.face_store import get_reference_faces, append_reference_face -from facefusion import face_analyser, face_masker, content_analyser, config, process_manager, metadata, logger, wording, voice_extractor +from facefusion import content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, logger, process_manager, state_manager, voice_extractor, wording +from facefusion.args import apply_args, collect_job_args, reduce_step_args +from facefusion.common_helper import get_first from facefusion.content_analyser import analyse_image, analyse_video -from facefusion.processors.frame.core import get_frame_processors_modules, load_frame_processor_module -from facefusion.common_helper import create_metavar, get_first -from facefusion.execution import encode_execution_providers, decode_execution_providers -from facefusion.normalizer import normalize_output_path, normalize_padding, normalize_fps +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.exit_helper import conditional_exit, graceful_exit, hard_exit +from facefusion.face_analyser import get_average_face, get_many_faces, get_one_face +from facefusion.face_selector import sort_and_filter_faces +from facefusion.face_store import append_reference_face, clear_reference_faces, get_reference_faces +from facefusion.ffmpeg import copy_image, extract_frames, finalize_image, merge_video, replace_audio, restore_audio +from facefusion.filesystem import filter_audio_paths, is_image, is_video, list_directory, resolve_relative_path +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.statistics import conditional_log_statistics -from facefusion.download import conditional_download -from facefusion.filesystem import get_temp_frame_paths, get_temp_file_path, create_temp, move_temp, clear_temp, is_image, is_video, filter_audio_paths, resolve_relative_path, list_directory -from facefusion.ffmpeg import extract_frames, merge_video, copy_image, finalize_image, restore_audio, replace_audio -from facefusion.vision import read_image, read_static_images, detect_image_resolution, restrict_video_fps, create_image_resolutions, get_video_frame, detect_video_resolution, detect_video_fps, restrict_video_resolution, restrict_image_resolution, create_video_resolutions, pack_resolution, unpack_resolution - -onnxruntime.set_default_logger_severity(3) -warnings.filterwarnings('ignore', category = UserWarning, module = 'gradio') +from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_file_path, get_temp_frame_paths, move_temp_file +from facefusion.typing import Args, ErrorCode +from facefusion.vision import get_video_frame, pack_resolution, read_image, read_static_images, restrict_image_resolution, restrict_video_fps, restrict_video_resolution, unpack_resolution def cli() -> None: - signal.signal(signal.SIGINT, lambda signal_number, frame: destroy()) - program = ArgumentParser(formatter_class = lambda prog: HelpFormatter(prog, max_help_position = 200), add_help = False) - # general - program.add_argument('-c', '--config', help = wording.get('help.config'), dest = 'config_path', default = 'facefusion.ini') - apply_config(program) - program.add_argument('-s', '--source', help = wording.get('help.source'), action = 'append', dest = 'source_paths', default = config.get_str_list('general.source_paths')) - program.add_argument('-t', '--target', help = wording.get('help.target'), dest = 'target_path', default = config.get_str_value('general.target_path')) - program.add_argument('-o', '--output', help = wording.get('help.output'), dest = 'output_path', default = config.get_str_value('general.output_path')) - program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version') - # misc - group_misc = program.add_argument_group('misc') - group_misc.add_argument('--force-download', help = wording.get('help.force_download'), action = 'store_true', default = config.get_bool_value('misc.force_download')) - group_misc.add_argument('--skip-download', help = wording.get('help.skip_download'), action = 'store_true', default = config.get_bool_value('misc.skip_download')) - group_misc.add_argument('--headless', help = wording.get('help.headless'), action = 'store_true', default = config.get_bool_value('misc.headless')) - group_misc.add_argument('--log-level', help = wording.get('help.log_level'), default = config.get_str_value('misc.log_level', 'info'), choices = logger.get_log_levels()) - # execution - execution_providers = encode_execution_providers(onnxruntime.get_available_providers()) - group_execution = program.add_argument_group('execution') - group_execution.add_argument('--execution-device-id', help = wording.get('help.execution_device_id'), default = config.get_str_value('execution.face_detector_size', '0')) - group_execution.add_argument('--execution-providers', help = wording.get('help.execution_providers').format(choices = ', '.join(execution_providers)), default = config.get_str_list('execution.execution_providers', 'cpu'), choices = execution_providers, nargs = '+', metavar = 'EXECUTION_PROVIDERS') - group_execution.add_argument('--execution-thread-count', help = wording.get('help.execution_thread_count'), type = int, default = config.get_int_value('execution.execution_thread_count', '4'), choices = facefusion.choices.execution_thread_count_range, metavar = create_metavar(facefusion.choices.execution_thread_count_range)) - group_execution.add_argument('--execution-queue-count', help = wording.get('help.execution_queue_count'), type = int, default = config.get_int_value('execution.execution_queue_count', '1'), choices = facefusion.choices.execution_queue_count_range, metavar = create_metavar(facefusion.choices.execution_queue_count_range)) - # memory - group_memory = program.add_argument_group('memory') - group_memory.add_argument('--video-memory-strategy', help = wording.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 = wording.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_metavar(facefusion.choices.system_memory_limit_range)) - # face analyser - group_face_analyser = program.add_argument_group('face analyser') - group_face_analyser.add_argument('--face-analyser-order', help = wording.get('help.face_analyser_order'), default = config.get_str_value('face_analyser.face_analyser_order', 'left-right'), choices = facefusion.choices.face_analyser_orders) - group_face_analyser.add_argument('--face-analyser-age', help = wording.get('help.face_analyser_age'), default = config.get_str_value('face_analyser.face_analyser_age'), choices = facefusion.choices.face_analyser_ages) - group_face_analyser.add_argument('--face-analyser-gender', help = wording.get('help.face_analyser_gender'), default = config.get_str_value('face_analyser.face_analyser_gender'), choices = facefusion.choices.face_analyser_genders) - group_face_analyser.add_argument('--face-detector-model', help = wording.get('help.face_detector_model'), default = config.get_str_value('face_analyser.face_detector_model', 'yoloface'), choices = facefusion.choices.face_detector_set.keys()) - group_face_analyser.add_argument('--face-detector-size', help = wording.get('help.face_detector_size'), default = config.get_str_value('face_analyser.face_detector_size', '640x640')) - group_face_analyser.add_argument('--face-detector-score', help = wording.get('help.face_detector_score'), type = float, default = config.get_float_value('face_analyser.face_detector_score', '0.5'), choices = facefusion.choices.face_detector_score_range, metavar = create_metavar(facefusion.choices.face_detector_score_range)) - group_face_analyser.add_argument('--face-landmarker-score', help = wording.get('help.face_landmarker_score'), type = float, default = config.get_float_value('face_analyser.face_landmarker_score', '0.5'), choices = facefusion.choices.face_landmarker_score_range, metavar = create_metavar(facefusion.choices.face_landmarker_score_range)) - # face selector - group_face_selector = program.add_argument_group('face selector') - group_face_selector.add_argument('--face-selector-mode', help = wording.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('--reference-face-position', help = wording.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 = wording.get('help.reference_face_distance'), type = float, default = config.get_float_value('face_selector.reference_face_distance', '0.6'), choices = facefusion.choices.reference_face_distance_range, metavar = create_metavar(facefusion.choices.reference_face_distance_range)) - group_face_selector.add_argument('--reference-frame-number', help = wording.get('help.reference_frame_number'), type = int, default = config.get_int_value('face_selector.reference_frame_number', '0')) - # face mask - group_face_mask = program.add_argument_group('face mask') - group_face_mask.add_argument('--face-mask-types', help = wording.get('help.face_mask_types').format(choices = ', '.join(facefusion.choices.face_mask_types)), default = config.get_str_list('face_mask.face_mask_types', 'box'), choices = facefusion.choices.face_mask_types, nargs = '+', metavar = 'FACE_MASK_TYPES') - group_face_mask.add_argument('--face-mask-blur', help = wording.get('help.face_mask_blur'), type = float, default = config.get_float_value('face_mask.face_mask_blur', '0.3'), choices = facefusion.choices.face_mask_blur_range, metavar = create_metavar(facefusion.choices.face_mask_blur_range)) - group_face_mask.add_argument('--face-mask-padding', help = wording.get('help.face_mask_padding'), type = int, default = config.get_int_list('face_mask.face_mask_padding', '0 0 0 0'), nargs = '+') - group_face_mask.add_argument('--face-mask-regions', help = wording.get('help.face_mask_regions').format(choices = ', '.join(facefusion.choices.face_mask_regions)), default = config.get_str_list('face_mask.face_mask_regions', ' '.join(facefusion.choices.face_mask_regions)), choices = facefusion.choices.face_mask_regions, nargs = '+', metavar = 'FACE_MASK_REGIONS') - # frame extraction - group_frame_extraction = program.add_argument_group('frame extraction') - group_frame_extraction.add_argument('--trim-frame-start', help = wording.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 = wording.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 = wording.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 = wording.get('help.keep_temp'), action = 'store_true', default = config.get_bool_value('frame_extraction.keep_temp')) - # output creation - group_output_creation = program.add_argument_group('output creation') - group_output_creation.add_argument('--output-image-quality', help = wording.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_metavar(facefusion.choices.output_image_quality_range)) - group_output_creation.add_argument('--output-image-resolution', help = wording.get('help.output_image_resolution'), default = config.get_str_value('output_creation.output_image_resolution')) - group_output_creation.add_argument('--output-video-encoder', help = wording.get('help.output_video_encoder'), default = config.get_str_value('output_creation.output_video_encoder', 'libx264'), choices = facefusion.choices.output_video_encoders) - group_output_creation.add_argument('--output-video-preset', help = wording.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 = wording.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_metavar(facefusion.choices.output_video_quality_range)) - group_output_creation.add_argument('--output-video-resolution', help = wording.get('help.output_video_resolution'), default = config.get_str_value('output_creation.output_video_resolution')) - group_output_creation.add_argument('--output-video-fps', help = wording.get('help.output_video_fps'), type = float, default = config.get_str_value('output_creation.output_video_fps')) - group_output_creation.add_argument('--skip-audio', help = wording.get('help.skip_audio'), action = 'store_true', default = config.get_bool_value('output_creation.skip_audio')) - # frame processors - available_frame_processors = list_directory('facefusion/processors/frame/modules') - program = ArgumentParser(parents = [ program ], formatter_class = program.formatter_class, add_help = True) - group_frame_processors = program.add_argument_group('frame processors') - group_frame_processors.add_argument('--frame-processors', help = wording.get('help.frame_processors').format(choices = ', '.join(available_frame_processors)), default = config.get_str_list('frame_processors.frame_processors', 'face_swapper'), nargs = '+') - for frame_processor in available_frame_processors: - frame_processor_module = load_frame_processor_module(frame_processor) - frame_processor_module.register_args(group_frame_processors) - # uis - available_ui_layouts = list_directory('facefusion/uis/layouts') - group_uis = program.add_argument_group('uis') - group_uis.add_argument('--open-browser', help=wording.get('help.open_browser'), action = 'store_true', default = config.get_bool_value('uis.open_browser')) - group_uis.add_argument('--ui-layouts', help = wording.get('help.ui_layouts').format(choices = ', '.join(available_ui_layouts)), default = config.get_str_list('uis.ui_layouts', 'default'), nargs = '+') - run(program) + signal.signal(signal.SIGINT, lambda signal_number, frame: graceful_exit(0)) + program = create_program() + if validate_args(program): + args = vars(program.parse_args()) + apply_args(args, state_manager.init_item) -def apply_config(program : ArgumentParser) -> None: - known_args = program.parse_known_args() - facefusion.globals.config_path = get_first(known_args).config_path - - -def validate_args(program : ArgumentParser) -> None: - try: - for action in program._actions: - if action.default: - if isinstance(action.default, list): - for default in action.default: - program._check_value(action, default) - else: - program._check_value(action, action.default) - except Exception as exception: - program.error(str(exception)) - - -def apply_args(program : ArgumentParser) -> None: - args = program.parse_args() - # general - facefusion.globals.source_paths = args.source_paths - facefusion.globals.target_path = args.target_path - facefusion.globals.output_path = args.output_path - # misc - facefusion.globals.force_download = args.force_download - facefusion.globals.skip_download = args.skip_download - facefusion.globals.headless = args.headless - facefusion.globals.log_level = args.log_level - # execution - facefusion.globals.execution_device_id = args.execution_device_id - facefusion.globals.execution_providers = decode_execution_providers(args.execution_providers) - facefusion.globals.execution_thread_count = args.execution_thread_count - facefusion.globals.execution_queue_count = args.execution_queue_count - # memory - facefusion.globals.video_memory_strategy = args.video_memory_strategy - facefusion.globals.system_memory_limit = args.system_memory_limit - # face analyser - facefusion.globals.face_analyser_order = args.face_analyser_order - facefusion.globals.face_analyser_age = args.face_analyser_age - facefusion.globals.face_analyser_gender = args.face_analyser_gender - facefusion.globals.face_detector_model = args.face_detector_model - if args.face_detector_size in facefusion.choices.face_detector_set[args.face_detector_model]: - facefusion.globals.face_detector_size = args.face_detector_size - else: - facefusion.globals.face_detector_size = '640x640' - facefusion.globals.face_detector_score = args.face_detector_score - facefusion.globals.face_landmarker_score = args.face_landmarker_score - # face selector - facefusion.globals.face_selector_mode = args.face_selector_mode - facefusion.globals.reference_face_position = args.reference_face_position - facefusion.globals.reference_face_distance = args.reference_face_distance - facefusion.globals.reference_frame_number = args.reference_frame_number - # face mask - facefusion.globals.face_mask_types = args.face_mask_types - facefusion.globals.face_mask_blur = args.face_mask_blur - facefusion.globals.face_mask_padding = normalize_padding(args.face_mask_padding) - facefusion.globals.face_mask_regions = args.face_mask_regions - # frame extraction - facefusion.globals.trim_frame_start = args.trim_frame_start - facefusion.globals.trim_frame_end = args.trim_frame_end - facefusion.globals.temp_frame_format = args.temp_frame_format - facefusion.globals.keep_temp = args.keep_temp - # output creation - facefusion.globals.output_image_quality = args.output_image_quality - if is_image(args.target_path): - output_image_resolution = detect_image_resolution(args.target_path) - output_image_resolutions = create_image_resolutions(output_image_resolution) - if args.output_image_resolution in output_image_resolutions: - facefusion.globals.output_image_resolution = args.output_image_resolution + if state_manager.get_item('command'): + logger.init(state_manager.get_item('log_level')) + route(args) else: - facefusion.globals.output_image_resolution = pack_resolution(output_image_resolution) - facefusion.globals.output_video_encoder = args.output_video_encoder - facefusion.globals.output_video_preset = args.output_video_preset - facefusion.globals.output_video_quality = args.output_video_quality - if is_video(args.target_path): - output_video_resolution = detect_video_resolution(args.target_path) - output_video_resolutions = create_video_resolutions(output_video_resolution) - if args.output_video_resolution in output_video_resolutions: - facefusion.globals.output_video_resolution = args.output_video_resolution - else: - facefusion.globals.output_video_resolution = pack_resolution(output_video_resolution) - if args.output_video_fps or is_video(args.target_path): - facefusion.globals.output_video_fps = normalize_fps(args.output_video_fps) or detect_video_fps(args.target_path) - facefusion.globals.skip_audio = args.skip_audio - # frame processors - available_frame_processors = list_directory('facefusion/processors/frame/modules') - facefusion.globals.frame_processors = args.frame_processors - for frame_processor in available_frame_processors: - frame_processor_module = load_frame_processor_module(frame_processor) - frame_processor_module.apply_args(program) - # uis - facefusion.globals.open_browser = args.open_browser - facefusion.globals.ui_layouts = args.ui_layouts + program.print_help() -def run(program : ArgumentParser) -> None: - validate_args(program) - apply_args(program) - logger.init(facefusion.globals.log_level) - - if facefusion.globals.system_memory_limit > 0: - limit_system_memory(facefusion.globals.system_memory_limit) - if facefusion.globals.force_download: - force_download() - return - if not pre_check() or not content_analyser.pre_check() or not face_analyser.pre_check() or not face_masker.pre_check() or not voice_extractor.pre_check(): - return - for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - if not frame_processor_module.pre_check(): - return - if facefusion.globals.headless: - conditional_process() - else: +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() + return conditional_exit(error_code) + if state_manager.get_item('command') in [ 'job-create', 'job-submit', 'job-submit-all', 'job-delete', 'job-delete-all', 'job-add-step', 'job-remix-step', 'job-insert-step', 'job-remove-step', 'job-list' ]: + 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 not pre_check(): + return conditional_exit(2) + if state_manager.get_item('command') == 'run': import facefusion.uis.core as ui - for ui_layout in ui.get_ui_layouts_modules(facefusion.globals.ui_layouts): + if not common_pre_check() or not processors_pre_check(): + return conditional_exit(2) + for ui_layout in ui.get_ui_layouts_modules(state_manager.get_item('ui_layouts')): if not ui_layout.pre_check(): - return + return conditional_exit(2) ui.launch() - - -def destroy() -> None: - process_manager.stop() - while process_manager.is_processing(): - sleep(0.5) - if facefusion.globals.target_path: - clear_temp(facefusion.globals.target_path) - sys.exit(0) + if state_manager.get_item('command') == 'headless-run': + if not job_manager.init_jobs(state_manager.get_item('jobs_path')): + hard_exit(1) + error_core = process_headless(args) + hard_exit(error_core) + 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, 9): - logger.error(wording.get('python_not_supported').format(version = '3.9'), __name__.upper()) + logger.error(wording.get('python_not_supported').format(version = '3.9'), __name__) + return False + if not shutil.which('curl'): + logger.error(wording.get('curl_not_installed'), __name__) return False if not shutil.which('ffmpeg'): - logger.error(wording.get('ffmpeg_not_installed'), __name__.upper()) + logger.error(wording.get('ffmpeg_not_installed'), __name__) return False return True -def conditional_process() -> None: +def common_pre_check() -> bool: + modules =\ + [ + content_analyser, + face_classifier, + face_detector, + face_landmarker, + face_masker, + face_recognizer, + voice_extractor + ] + + return all(module.pre_check() for module in modules) + + +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 conditional_process() -> ErrorCode: start_time = time() - for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - while not frame_processor_module.post_check(): - logger.disable() - sleep(0.5) - logger.enable() - if not frame_processor_module.pre_process('output'): - return + for processor_module in get_processors_modules(state_manager.get_item('processors')): + if not processor_module.pre_process('output'): + return 2 conditional_append_reference_faces() - if is_image(facefusion.globals.target_path): - process_image(start_time) - if is_video(facefusion.globals.target_path): - process_video(start_time) + if is_image(state_manager.get_item('target_path')): + return process_image(start_time) + if is_video(state_manager.get_item('target_path')): + return process_video(start_time) + return 0 def conditional_append_reference_faces() -> None: - if 'reference' in facefusion.globals.face_selector_mode and not get_reference_faces(): - source_frames = read_static_images(facefusion.globals.source_paths) - source_face = get_average_face(source_frames) - if is_video(facefusion.globals.target_path): - reference_frame = get_video_frame(facefusion.globals.target_path, facefusion.globals.reference_frame_number) + if 'reference' in state_manager.get_item('face_selector_mode') and not get_reference_faces(): + source_frames = read_static_images(state_manager.get_item('source_paths')) + source_faces = get_many_faces(source_frames) + source_face = get_average_face(source_faces) + if is_video(state_manager.get_item('target_path')): + reference_frame = get_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) else: - reference_frame = read_image(facefusion.globals.target_path) - reference_face = get_one_face(reference_frame, facefusion.globals.reference_face_position) + reference_frame = read_image(state_manager.get_item('target_path')) + reference_faces = sort_and_filter_faces(get_many_faces([ reference_frame ])) + reference_face = get_one_face(reference_faces, state_manager.get_item('reference_face_position')) append_reference_face('origin', reference_face) + if source_face and reference_face: - for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - abstract_reference_frame = frame_processor_module.get_reference_frame(source_face, reference_face, reference_frame) + for processor_module in get_processors_modules(state_manager.get_item('processors')): + abstract_reference_frame = processor_module.get_reference_frame(source_face, reference_face, reference_frame) if numpy.any(abstract_reference_frame): - reference_frame = abstract_reference_frame - reference_face = get_one_face(reference_frame, facefusion.globals.reference_face_position) - append_reference_face(frame_processor_module.__name__, reference_face) + abstract_reference_faces = sort_and_filter_faces(get_many_faces([ abstract_reference_frame ])) + abstract_reference_face = get_one_face(abstract_reference_faces, state_manager.get_item('reference_face_position')) + append_reference_face(processor_module.__name__, abstract_reference_face) -def force_download() -> None: +def force_download() -> ErrorCode: download_directory_path = resolve_relative_path('../.assets/models') - available_frame_processors = list_directory('facefusion/processors/frame/modules') - model_list =\ + available_processors = list_directory('facefusion/processors/modules') + common_modules =\ [ - content_analyser.MODELS, - face_analyser.MODELS, - face_masker.MODELS, - voice_extractor.MODELS + content_analyser, + face_classifier, + face_detector, + face_landmarker, + face_recognizer, + face_masker, + voice_extractor ] + processor_modules = get_processors_modules(available_processors) - for frame_processor_module in get_frame_processors_modules(available_frame_processors): - if hasattr(frame_processor_module, 'MODELS'): - model_list.append(frame_processor_module.MODELS) - model_urls = [ models[model].get('url') for models in model_list for model in models ] - conditional_download(download_directory_path, model_urls) + for module in common_modules + processor_modules: + if hasattr(module, 'MODEL_SET'): + for model in module.MODEL_SET.values(): + model_hashes = model.get('hashes') + model_sources = model.get('sources') + + if model_hashes and model_sources: + if not conditional_download_hashes(download_directory_path, model_hashes) or not conditional_download_sources(download_directory_path, model_sources): + return 1 + + return 0 -def process_image(start_time : float) -> None: - normed_output_path = normalize_output_path(facefusion.globals.target_path, facefusion.globals.output_path) - if analyse_image(facefusion.globals.target_path): - return +def route_job_manager(args : Args) -> ErrorCode: + if state_manager.get_item('command') == 'job-create': + if job_manager.create_job(state_manager.get_item('job_id')): + logger.info(wording.get('job_created').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.error(wording.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(wording.get('job_submitted').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.error(wording.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(): + logger.info(wording.get('job_all_submitted'), __name__) + return 0 + logger.error(wording.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(wording.get('job_deleted').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.error(wording.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(): + logger.info(wording.get('job_all_deleted'), __name__) + return 0 + logger.error(wording.get('job_all_not_deleted'), __name__) + return 1 + if state_manager.get_item('command') == 'job-list': + job_headers, job_contents = compose_job_list(state_manager.get_item('job_status')) + + if job_contents: + logger.table(job_headers, job_contents) + return 0 + 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(wording.get('job_step_added').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.error(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.get('processing_job_succeed').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.info(wording.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(wording.get('running_jobs'), __name__) + if job_runner.run_jobs(process_step): + logger.info(wording.get('processing_jobs_succeed'), __name__) + return 0 + logger.info(wording.get('processing_jobs_failed'), __name__) + return 1 + if state_manager.get_item('command') == 'job-retry': + logger.info(wording.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(wording.get('processing_job_succeed').format(job_id = state_manager.get_item('job_id')), __name__) + return 0 + logger.info(wording.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(wording.get('retrying_jobs'), __name__) + if job_runner.retry_jobs(process_step): + logger.info(wording.get('processing_jobs_succeed'), __name__) + return 0 + logger.info(wording.get('processing_jobs_failed'), __name__) + return 1 + return 2 + + +def process_step(job_id : str, step_index : int, step_args : Args) -> bool: + clear_reference_faces() + 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(wording.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 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_image(start_time : float) -> ErrorCode: + if analyse_image(state_manager.get_item('target_path')): + return 3 # clear temp - logger.debug(wording.get('clearing_temp'), __name__.upper()) - clear_temp(facefusion.globals.target_path) + logger.debug(wording.get('clearing_temp'), __name__) + clear_temp_directory(state_manager.get_item('target_path')) # create temp - logger.debug(wording.get('creating_temp'), __name__.upper()) - create_temp(facefusion.globals.target_path) + logger.debug(wording.get('creating_temp'), __name__) + create_temp_directory(state_manager.get_item('target_path')) # copy image process_manager.start() - temp_image_resolution = pack_resolution(restrict_image_resolution(facefusion.globals.target_path, unpack_resolution(facefusion.globals.output_image_resolution))) - logger.info(wording.get('copying_image').format(resolution = temp_image_resolution), __name__.upper()) - if copy_image(facefusion.globals.target_path, temp_image_resolution): - logger.debug(wording.get('copying_image_succeed'), __name__.upper()) + temp_image_resolution = pack_resolution(restrict_image_resolution(state_manager.get_item('target_path'), unpack_resolution(state_manager.get_item('output_image_resolution')))) + logger.info(wording.get('copying_image').format(resolution = temp_image_resolution), __name__) + if copy_image(state_manager.get_item('target_path'), temp_image_resolution): + logger.debug(wording.get('copying_image_succeed'), __name__) else: - logger.error(wording.get('copying_image_failed'), __name__.upper()) - return + logger.error(wording.get('copying_image_failed'), __name__) + process_manager.end() + return 1 # process image - temp_file_path = get_temp_file_path(facefusion.globals.target_path) - for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - logger.info(wording.get('processing'), frame_processor_module.NAME) - frame_processor_module.process_image(facefusion.globals.source_paths, temp_file_path, temp_file_path) - frame_processor_module.post_process() + temp_file_path = get_temp_file_path(state_manager.get_item('target_path')) + for processor_module in get_processors_modules(state_manager.get_item('processors')): + logger.info(wording.get('processing'), processor_module.__name__) + processor_module.process_image(state_manager.get_item('source_paths'), temp_file_path, temp_file_path) + processor_module.post_process() if is_process_stopping(): - return + process_manager.end() + return 4 # finalize image - logger.info(wording.get('finalizing_image').format(resolution = facefusion.globals.output_image_resolution), __name__.upper()) - if finalize_image(facefusion.globals.target_path, normed_output_path, facefusion.globals.output_image_resolution): - logger.debug(wording.get('finalizing_image_succeed'), __name__.upper()) + logger.info(wording.get('finalizing_image').format(resolution = state_manager.get_item('output_image_resolution')), __name__) + if finalize_image(state_manager.get_item('target_path'), state_manager.get_item('output_path'), state_manager.get_item('output_image_resolution')): + logger.debug(wording.get('finalizing_image_succeed'), __name__) else: - logger.warn(wording.get('finalizing_image_skipped'), __name__.upper()) + logger.warn(wording.get('finalizing_image_skipped'), __name__) # clear temp - logger.debug(wording.get('clearing_temp'), __name__.upper()) - clear_temp(facefusion.globals.target_path) + logger.debug(wording.get('clearing_temp'), __name__) + clear_temp_directory(state_manager.get_item('target_path')) # validate image - if is_image(normed_output_path): + if is_image(state_manager.get_item('output_path')): seconds = '{:.2f}'.format((time() - start_time) % 60) - logger.info(wording.get('processing_image_succeed').format(seconds = seconds), __name__.upper()) + logger.info(wording.get('processing_image_succeed').format(seconds = seconds), __name__) conditional_log_statistics() else: - logger.error(wording.get('processing_image_failed'), __name__.upper()) + logger.error(wording.get('processing_image_failed'), __name__) + process_manager.end() + return 1 process_manager.end() + return 0 -def process_video(start_time : float) -> None: - normed_output_path = normalize_output_path(facefusion.globals.target_path, facefusion.globals.output_path) - if analyse_video(facefusion.globals.target_path, facefusion.globals.trim_frame_start, facefusion.globals.trim_frame_end): - return +def process_video(start_time : float) -> ErrorCode: + if analyse_video(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')): + return 3 # clear temp - logger.debug(wording.get('clearing_temp'), __name__.upper()) - clear_temp(facefusion.globals.target_path) + logger.debug(wording.get('clearing_temp'), __name__) + clear_temp_directory(state_manager.get_item('target_path')) # create temp - logger.debug(wording.get('creating_temp'), __name__.upper()) - create_temp(facefusion.globals.target_path) + logger.debug(wording.get('creating_temp'), __name__) + create_temp_directory(state_manager.get_item('target_path')) # extract frames process_manager.start() - temp_video_resolution = pack_resolution(restrict_video_resolution(facefusion.globals.target_path, unpack_resolution(facefusion.globals.output_video_resolution))) - temp_video_fps = restrict_video_fps(facefusion.globals.target_path, facefusion.globals.output_video_fps) - logger.info(wording.get('extracting_frames').format(resolution = temp_video_resolution, fps = temp_video_fps), __name__.upper()) - if extract_frames(facefusion.globals.target_path, temp_video_resolution, temp_video_fps): - logger.debug(wording.get('extracting_frames_succeed'), __name__.upper()) + temp_video_resolution = pack_resolution(restrict_video_resolution(state_manager.get_item('target_path'), unpack_resolution(state_manager.get_item('output_video_resolution')))) + temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + logger.info(wording.get('extracting_frames').format(resolution = temp_video_resolution, fps = temp_video_fps), __name__) + if extract_frames(state_manager.get_item('target_path'), temp_video_resolution, temp_video_fps): + logger.debug(wording.get('extracting_frames_succeed'), __name__) else: if is_process_stopping(): - return - logger.error(wording.get('extracting_frames_failed'), __name__.upper()) - return + process_manager.end() + return 4 + logger.error(wording.get('extracting_frames_failed'), __name__) + process_manager.end() + return 1 # process frames - temp_frame_paths = get_temp_frame_paths(facefusion.globals.target_path) + temp_frame_paths = get_temp_frame_paths(state_manager.get_item('target_path')) if temp_frame_paths: - for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - logger.info(wording.get('processing'), frame_processor_module.NAME) - frame_processor_module.process_video(facefusion.globals.source_paths, temp_frame_paths) - frame_processor_module.post_process() + for processor_module in get_processors_modules(state_manager.get_item('processors')): + logger.info(wording.get('processing'), processor_module.__name__) + processor_module.process_video(state_manager.get_item('source_paths'), temp_frame_paths) + processor_module.post_process() if is_process_stopping(): - return + return 4 else: - logger.error(wording.get('temp_frames_not_found'), __name__.upper()) - return + logger.error(wording.get('temp_frames_not_found'), __name__) + process_manager.end() + return 1 # merge video - logger.info(wording.get('merging_video').format(resolution = facefusion.globals.output_video_resolution, fps = facefusion.globals.output_video_fps), __name__.upper()) - if merge_video(facefusion.globals.target_path, facefusion.globals.output_video_resolution, facefusion.globals.output_video_fps): - logger.debug(wording.get('merging_video_succeed'), __name__.upper()) + logger.info(wording.get('merging_video').format(resolution = state_manager.get_item('output_video_resolution'), fps = state_manager.get_item('output_video_fps')), __name__) + if merge_video(state_manager.get_item('target_path'), state_manager.get_item('output_video_resolution'), state_manager.get_item('output_video_fps')): + logger.debug(wording.get('merging_video_succeed'), __name__) else: if is_process_stopping(): - return - logger.error(wording.get('merging_video_failed'), __name__.upper()) - return + process_manager.end() + return 4 + logger.error(wording.get('merging_video_failed'), __name__) + process_manager.end() + return 1 # handle audio - if facefusion.globals.skip_audio: - logger.info(wording.get('skipping_audio'), __name__.upper()) - move_temp(facefusion.globals.target_path, normed_output_path) + if state_manager.get_item('skip_audio'): + logger.info(wording.get('skipping_audio'), __name__) + move_temp_file(state_manager.get_item('target_path'), state_manager.get_item('output_path')) else: - if 'lip_syncer' in facefusion.globals.frame_processors: - source_audio_path = get_first(filter_audio_paths(facefusion.globals.source_paths)) - if source_audio_path and replace_audio(facefusion.globals.target_path, source_audio_path, normed_output_path): - logger.debug(wording.get('restoring_audio_succeed'), __name__.upper()) + if 'lip_syncer' in state_manager.get_item('processors'): + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) + if source_audio_path and replace_audio(state_manager.get_item('target_path'), source_audio_path, state_manager.get_item('output_path')): + logger.debug(wording.get('restoring_audio_succeed'), __name__) else: if is_process_stopping(): - return - logger.warn(wording.get('restoring_audio_skipped'), __name__.upper()) - move_temp(facefusion.globals.target_path, normed_output_path) + process_manager.end() + return 4 + logger.warn(wording.get('restoring_audio_skipped'), __name__) + move_temp_file(state_manager.get_item('target_path'), state_manager.get_item('output_path')) else: - if restore_audio(facefusion.globals.target_path, normed_output_path, facefusion.globals.output_video_fps): - logger.debug(wording.get('restoring_audio_succeed'), __name__.upper()) + if restore_audio(state_manager.get_item('target_path'), state_manager.get_item('output_path'), state_manager.get_item('output_video_fps')): + logger.debug(wording.get('restoring_audio_succeed'), __name__) else: if is_process_stopping(): - return - logger.warn(wording.get('restoring_audio_skipped'), __name__.upper()) - move_temp(facefusion.globals.target_path, normed_output_path) + process_manager.end() + return 4 + logger.warn(wording.get('restoring_audio_skipped'), __name__) + move_temp_file(state_manager.get_item('target_path'), state_manager.get_item('output_path')) # clear temp - logger.debug(wording.get('clearing_temp'), __name__.upper()) - clear_temp(facefusion.globals.target_path) + logger.debug(wording.get('clearing_temp'), __name__) + clear_temp_directory(state_manager.get_item('target_path')) # validate video - if is_video(normed_output_path): + if is_video(state_manager.get_item('output_path')): seconds = '{:.2f}'.format((time() - start_time)) - logger.info(wording.get('processing_video_succeed').format(seconds = seconds), __name__.upper()) + logger.info(wording.get('processing_video_succeed').format(seconds = seconds), __name__) conditional_log_statistics() else: - logger.error(wording.get('processing_video_failed'), __name__.upper()) + logger.error(wording.get('processing_video_failed'), __name__) + process_manager.end() + return 1 process_manager.end() + return 0 def is_process_stopping() -> bool: if process_manager.is_stopping(): process_manager.end() - logger.info(wording.get('processing_stopped'), __name__.upper()) + logger.info(wording.get('processing_stopped'), __name__) return process_manager.is_pending() diff --git a/facefusion/date_helper.py b/facefusion/date_helper.py new file mode 100644 index 00000000..c60e2f61 --- /dev/null +++ b/facefusion/date_helper.py @@ -0,0 +1,28 @@ +from datetime import datetime, timedelta +from typing import Optional, Tuple + +from facefusion import wording + + +def get_current_date_time() -> datetime: + return datetime.now().astimezone() + + +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 wording.get('time_ago_days').format(days = days, hours = hours, minutes = minutes) + if timedelta(hours = 1) < time_ago: + return wording.get('time_ago_hours').format(hours = hours, minutes = minutes) + if timedelta(minutes = 1) < time_ago: + return wording.get('time_ago_minutes').format(minutes = minutes) + return wording.get('time_ago_now') diff --git a/facefusion/download.py b/facefusion/download.py index 1e417573..43c92a0e 100644 --- a/facefusion/download.py +++ b/facefusion/download.py @@ -1,15 +1,19 @@ import os -import subprocess +import shutil import ssl +import subprocess import urllib.request -from typing import List from functools import lru_cache +from typing import List, Tuple +from urllib.parse import urlparse + from tqdm import tqdm -import facefusion.globals -from facefusion import wording +from facefusion import logger, process_manager, state_manager, wording from facefusion.common_helper import is_macos -from facefusion.filesystem import get_file_size, is_file +from facefusion.filesystem import get_file_size, is_file, remove_file +from facefusion.hash_helper import validate_hash +from facefusion.typing import DownloadSet if is_macos(): ssl._create_default_https_context = ssl._create_unverified_context @@ -17,28 +21,30 @@ if is_macos(): def conditional_download(download_directory_path : str, urls : List[str]) -> None: for url in urls: - download_file_path = os.path.join(download_directory_path, os.path.basename(url)) + 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_download_size(url) + if initial_size < download_size: - with tqdm(total = download_size, initial = initial_size, desc = wording.get('downloading'), unit = 'B', unit_scale = True, unit_divisor = 1024, ascii = ' =', disable = facefusion.globals.log_level in [ 'warn', 'error' ]) as progress: - subprocess.Popen([ 'curl', '--create-dirs', '--silent', '--insecure', '--location', '--continue-at', '-', '--output', download_file_path, url ]) + with tqdm(total = download_size, initial = initial_size, desc = wording.get('downloading'), unit = 'B', unit_scale = True, unit_divisor = 1024, ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: + subprocess.Popen([ shutil.which('curl'), '--create-dirs', '--silent', '--insecure', '--location', '--continue-at', '-', '--output', download_file_path, url ]) current_size = initial_size + + progress.set_postfix(file = 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) - if download_size and not is_download_done(url, download_file_path): - os.remove(download_file_path) - conditional_download(download_directory_path, [ url ]) @lru_cache(maxsize = None) def get_download_size(url : str) -> int: try: response = urllib.request.urlopen(url, timeout = 10) - return int(response.getheader('Content-Length')) - except (OSError, ValueError): + content_length = response.headers.get('Content-Length') + return int(content_length) + except (OSError, TypeError, ValueError): return 0 @@ -46,3 +52,80 @@ def is_download_done(url : str, file_path : str) -> bool: if is_file(file_path): return get_download_size(url) == get_file_size(file_path) return False + + +def conditional_download_hashes(download_directory_path : str, hashes : DownloadSet) -> bool: + hash_paths = [ hashes.get(hash_key).get('path') for hash_key in hashes.keys() ] + + process_manager.check() + if not state_manager.get_item('skip_download'): + _, invalid_hash_paths = validate_hash_paths(hash_paths) + if invalid_hash_paths: + for index in hashes: + if hashes.get(index).get('path') in invalid_hash_paths: + invalid_hash_url = hashes.get(index).get('url') + 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, _ = os.path.splitext(os.path.basename(valid_hash_path)) + logger.debug(wording.get('validating_hash_succeed').format(hash_file_name = valid_hash_file_name), __name__) + for invalid_hash_path in invalid_hash_paths: + invalid_hash_file_name, _ = os.path.splitext(os.path.basename(invalid_hash_path)) + logger.error(wording.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(download_directory_path : str, sources : DownloadSet) -> bool: + source_paths = [ sources.get(source_key).get('path') for source_key in sources.keys() ] + + process_manager.check() + if not state_manager.get_item('skip_download'): + _, invalid_source_paths = validate_source_paths(source_paths) + if invalid_source_paths: + for index in sources: + if sources.get(index).get('path') in invalid_source_paths: + invalid_source_url = sources.get(index).get('url') + 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, _ = os.path.splitext(os.path.basename(valid_source_path)) + logger.debug(wording.get('validating_source_succeed').format(source_file_name = valid_source_file_name), __name__) + for invalid_source_path in invalid_source_paths: + invalid_source_file_name, _ = os.path.splitext(os.path.basename(invalid_source_path)) + logger.error(wording.get('validating_source_failed').format(source_file_name = invalid_source_file_name), __name__) + + if remove_file(invalid_source_path): + logger.error(wording.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 diff --git a/facefusion/execution.py b/facefusion/execution.py index 42b0ce10..771923e4 100644 --- a/facefusion/execution.py +++ b/facefusion/execution.py @@ -1,28 +1,40 @@ -from typing import List, Any -from functools import lru_cache import subprocess import xml.etree.ElementTree as ElementTree -import onnxruntime +from functools import lru_cache +from typing import Any, List -from facefusion.typing import ExecutionDevice, ValueAndUnit +from onnxruntime import get_available_providers, set_default_logger_severity + +from facefusion.choices import execution_provider_set +from facefusion.typing import ExecutionDevice, ExecutionProviderKey, ExecutionProviderSet, ExecutionProviderValue, ValueAndUnit + +set_default_logger_severity(3) -def encode_execution_providers(execution_providers : List[str]) -> List[str]: - return [ execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers ] +def get_execution_provider_choices() -> List[ExecutionProviderKey]: + return list(get_available_execution_provider_set().keys()) -def decode_execution_providers(execution_providers : List[str]) -> List[str]: - available_execution_providers = onnxruntime.get_available_providers() - encoded_execution_providers = encode_execution_providers(available_execution_providers) - - return [ execution_provider for execution_provider, encoded_execution_provider in zip(available_execution_providers, encoded_execution_providers) if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers) ] +def has_execution_provider(execution_provider_key : ExecutionProviderKey) -> bool: + return execution_provider_key in get_execution_provider_choices() -def has_execution_provider(execution_provider : str) -> bool: - return execution_provider in onnxruntime.get_available_providers() +def get_available_execution_provider_set() -> ExecutionProviderSet: + available_execution_providers = get_available_providers() + available_execution_provider_set : ExecutionProviderSet = {} + + for execution_provider_key, execution_provider_value in execution_provider_set.items(): + if execution_provider_value in available_execution_providers: + available_execution_provider_set[execution_provider_key] = execution_provider_value + return available_execution_provider_set -def apply_execution_provider_options(execution_device_id : str, execution_providers : List[str]) -> List[Any]: +def extract_execution_providers(execution_provider_keys : List[ExecutionProviderKey]) -> List[ExecutionProviderValue]: + return [ execution_provider_set[execution_provider_key] for execution_provider_key in execution_provider_keys if execution_provider_key in execution_provider_set ] + + +def create_execution_providers(execution_device_id : str, execution_provider_keys : List[ExecutionProviderKey]) -> List[Any]: + execution_providers = extract_execution_providers(execution_provider_keys) execution_providers_with_options : List[Any] = [] for execution_provider in execution_providers: @@ -32,19 +44,33 @@ def apply_execution_provider_options(execution_device_id : str, execution_provid 'device_id': execution_device_id, 'cudnn_conv_algo_search': 'EXHAUSTIVE' if use_exhaustive() else 'DEFAULT' })) - elif execution_provider == 'OpenVINOExecutionProvider': + if execution_provider == 'TensorrtExecutionProvider': execution_providers_with_options.append((execution_provider, { 'device_id': execution_device_id, - 'device_type': execution_device_id + '_FP32' + '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 })) - elif execution_provider in [ 'DmlExecutionProvider', 'ROCMExecutionProvider' ]: + if execution_provider == 'OpenVINOExecutionProvider': + execution_providers_with_options.append((execution_provider, + { + 'device_type': 'GPU.' + execution_device_id, + 'precision': 'FP32' + })) + if execution_provider in [ 'DmlExecutionProvider', 'ROCMExecutionProvider' ]: execution_providers_with_options.append((execution_provider, { 'device_id': execution_device_id })) - else: + if execution_provider == 'CoreMLExecutionProvider': execution_providers_with_options.append(execution_provider) + + if 'CPUExecutionProvider' in execution_providers: + execution_providers_with_options.append('CPUExecutionProvider') + return execution_providers_with_options @@ -67,6 +93,7 @@ def detect_static_execution_devices() -> List[ExecutionDevice]: def detect_execution_devices() -> List[ExecutionDevice]: execution_devices : List[ExecutionDevice] = [] + try: output, _ = run_nvidia_smi().communicate() root_element = ElementTree.fromstring(output) @@ -105,8 +132,8 @@ def create_value_and_unit(text : str) -> ValueAndUnit: value, unit = text.split() value_and_unit : ValueAndUnit =\ { - 'value': value, - 'unit': unit + 'value': int(value), + 'unit': str(unit) } return value_and_unit diff --git a/facefusion/exit_helper.py b/facefusion/exit_helper.py new file mode 100644 index 00000000..e2991b43 --- /dev/null +++ b/facefusion/exit_helper.py @@ -0,0 +1,24 @@ +import sys +from time import sleep + +from facefusion import process_manager, state_manager +from facefusion.temp_helper import clear_temp_directory +from facefusion.typing import ErrorCode + + +def hard_exit(error_code : ErrorCode) -> None: + sys.exit(error_code) + + +def conditional_exit(error_code : ErrorCode) -> None: + if state_manager.get_item('command') == 'headless-run': + hard_exit(error_code) + + +def graceful_exit(error_code : ErrorCode) -> None: + 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 index c9fee571..8e870d62 100644 --- a/facefusion/face_analyser.py +++ b/facefusion/face_analyser.py @@ -1,586 +1,124 @@ -from typing import Any, Optional, List, Tuple -from time import sleep -import cv2 +from typing import List, Optional + import numpy -import onnxruntime -import facefusion.globals -from facefusion import process_manager +from facefusion import state_manager from facefusion.common_helper import get_first -from facefusion.face_helper import estimate_matrix_by_face_landmark_5, warp_face_by_face_landmark_5, warp_face_by_translation, create_static_anchors, distance_to_face_landmark_5, distance_to_bounding_box, convert_face_landmark_68_to_5, apply_nms, categorize_age, categorize_gender +from facefusion.face_classifier import classify_face +from facefusion.face_detector import detect_faces, detect_rotated_faces +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_landmarks, estimate_face_landmark_68_5 +from facefusion.face_recognizer import calc_embedding from facefusion.face_store import get_static_faces, set_static_faces -from facefusion.execution import apply_execution_provider_options -from facefusion.download import conditional_download -from facefusion.filesystem import resolve_relative_path, is_file -from facefusion.thread_helper import thread_lock, thread_semaphore, conditional_thread_semaphore -from facefusion.typing import VisionFrame, Face, FaceSet, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, ModelSet, BoundingBox, FaceLandmarkSet, FaceLandmark5, FaceLandmark68, Score, FaceScoreSet, Embedding -from facefusion.vision import resize_frame_resolution, unpack_resolution - -FACE_ANALYSER = None -MODELS : ModelSet =\ -{ - 'face_detector_retinaface': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/retinaface_10g.onnx', - 'path': resolve_relative_path('../.assets/models/retinaface_10g.onnx') - }, - 'face_detector_scrfd': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/scrfd_2.5g.onnx', - 'path': resolve_relative_path('../.assets/models/scrfd_2.5g.onnx') - }, - 'face_detector_yoloface': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/yoloface_8n.onnx', - 'path': resolve_relative_path('../.assets/models/yoloface_8n.onnx') - }, - 'face_detector_yunet': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/yunet_2023mar.onnx', - 'path': resolve_relative_path('../.assets/models/yunet_2023mar.onnx') - }, - 'face_recognizer_arcface_blendswap': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_w600k_r50.onnx', - 'path': resolve_relative_path('../.assets/models/arcface_w600k_r50.onnx') - }, - 'face_recognizer_arcface_inswapper': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_w600k_r50.onnx', - 'path': resolve_relative_path('../.assets/models/arcface_w600k_r50.onnx') - }, - 'face_recognizer_arcface_simswap': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_simswap.onnx', - 'path': resolve_relative_path('../.assets/models/arcface_simswap.onnx') - }, - 'face_recognizer_arcface_uniface': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_w600k_r50.onnx', - 'path': resolve_relative_path('../.assets/models/arcface_w600k_r50.onnx') - }, - 'face_landmarker_68': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/2dfan4.onnx', - 'path': resolve_relative_path('../.assets/models/2dfan4.onnx') - }, - 'face_landmarker_68_5': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_landmarker_68_5.onnx', - 'path': resolve_relative_path('../.assets/models/face_landmarker_68_5.onnx') - }, - 'gender_age': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gender_age.onnx', - 'path': resolve_relative_path('../.assets/models/gender_age.onnx') - } -} +from facefusion.typing import BoundingBox, Face, FaceLandmark5, FaceLandmarkSet, FaceScoreSet, Score, VisionFrame -def get_face_analyser() -> Any: - global FACE_ANALYSER - - face_detectors = {} - face_landmarkers = {} - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if FACE_ANALYSER is None: - if facefusion.globals.face_detector_model in [ 'many', 'retinaface' ]: - face_detectors['retinaface'] = onnxruntime.InferenceSession(MODELS.get('face_detector_retinaface').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - if facefusion.globals.face_detector_model in [ 'many', 'scrfd' ]: - face_detectors['scrfd'] = onnxruntime.InferenceSession(MODELS.get('face_detector_scrfd').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - if facefusion.globals.face_detector_model in [ 'many', 'yoloface' ]: - face_detectors['yoloface'] = onnxruntime.InferenceSession(MODELS.get('face_detector_yoloface').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - if facefusion.globals.face_detector_model in [ 'yunet' ]: - face_detectors['yunet'] = cv2.FaceDetectorYN.create(MODELS.get('face_detector_yunet').get('path'), '', (0, 0)) - if facefusion.globals.face_recognizer_model == 'arcface_blendswap': - face_recognizer = onnxruntime.InferenceSession(MODELS.get('face_recognizer_arcface_blendswap').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - if facefusion.globals.face_recognizer_model == 'arcface_inswapper': - face_recognizer = onnxruntime.InferenceSession(MODELS.get('face_recognizer_arcface_inswapper').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - if facefusion.globals.face_recognizer_model == 'arcface_simswap': - face_recognizer = onnxruntime.InferenceSession(MODELS.get('face_recognizer_arcface_simswap').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - if facefusion.globals.face_recognizer_model == 'arcface_uniface': - face_recognizer = onnxruntime.InferenceSession(MODELS.get('face_recognizer_arcface_uniface').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - face_landmarkers['68'] = onnxruntime.InferenceSession(MODELS.get('face_landmarker_68').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - face_landmarkers['68_5'] = onnxruntime.InferenceSession(MODELS.get('face_landmarker_68_5').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - gender_age = onnxruntime.InferenceSession(MODELS.get('gender_age').get('path'), providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - FACE_ANALYSER =\ - { - 'face_detectors': face_detectors, - 'face_recognizer': face_recognizer, - 'face_landmarkers': face_landmarkers, - 'gender_age': gender_age - } - return FACE_ANALYSER - - -def clear_face_analyser() -> Any: - global FACE_ANALYSER - - FACE_ANALYSER = None - - -def pre_check() -> bool: - download_directory_path = resolve_relative_path('../.assets/models') - model_urls =\ - [ - MODELS.get('face_landmarker_68').get('url'), - MODELS.get('face_landmarker_68_5').get('url'), - MODELS.get('gender_age').get('url') - ] - model_paths =\ - [ - MODELS.get('face_landmarker_68').get('path'), - MODELS.get('face_landmarker_68_5').get('path'), - MODELS.get('gender_age').get('path') - ] - - if facefusion.globals.face_detector_model in [ 'many', 'retinaface' ]: - model_urls.append(MODELS.get('face_detector_retinaface').get('url')) - model_paths.append(MODELS.get('face_detector_retinaface').get('path')) - if facefusion.globals.face_detector_model in [ 'many', 'scrfd' ]: - model_urls.append(MODELS.get('face_detector_scrfd').get('url')) - model_paths.append(MODELS.get('face_detector_scrfd').get('path')) - if facefusion.globals.face_detector_model in [ 'many', 'yoloface' ]: - model_urls.append(MODELS.get('face_detector_yoloface').get('url')) - model_paths.append(MODELS.get('face_detector_yoloface').get('path')) - if facefusion.globals.face_detector_model in [ 'yunet' ]: - model_urls.append(MODELS.get('face_detector_yunet').get('url')) - model_paths.append(MODELS.get('face_detector_yunet').get('path')) - if facefusion.globals.face_recognizer_model == 'arcface_blendswap': - model_urls.append(MODELS.get('face_recognizer_arcface_blendswap').get('url')) - model_paths.append(MODELS.get('face_recognizer_arcface_blendswap').get('path')) - if facefusion.globals.face_recognizer_model == 'arcface_inswapper': - model_urls.append(MODELS.get('face_recognizer_arcface_inswapper').get('url')) - model_paths.append(MODELS.get('face_recognizer_arcface_inswapper').get('path')) - if facefusion.globals.face_recognizer_model == 'arcface_simswap': - model_urls.append(MODELS.get('face_recognizer_arcface_simswap').get('url')) - model_paths.append(MODELS.get('face_recognizer_arcface_simswap').get('path')) - if facefusion.globals.face_recognizer_model == 'arcface_uniface': - model_urls.append(MODELS.get('face_recognizer_arcface_uniface').get('url')) - model_paths.append(MODELS.get('face_recognizer_arcface_uniface').get('path')) - - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, model_urls) - process_manager.end() - return all(is_file(model_path) for model_path in model_paths) - - -def detect_with_retinaface(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[FaceLandmark5], List[Score]]: - face_detector = get_face_analyser().get('face_detectors').get('retinaface') - face_detector_width, face_detector_height = unpack_resolution(face_detector_size) - temp_vision_frame = resize_frame_resolution(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] - feature_strides = [ 8, 16, 32 ] - feature_map_channel = 3 - anchor_total = 2 - bounding_box_list = [] - face_landmark_5_list = [] - score_list = [] - - detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) - with thread_semaphore(): - detections = face_detector.run(None, - { - face_detector.get_inputs()[0].name: detect_vision_frame - }) - for index, feature_stride in enumerate(feature_strides): - keep_indices = numpy.where(detections[index] >= facefusion.globals.face_detector_score)[0] - if keep_indices.any(): - 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_box_raw = detections[index + feature_map_channel] * feature_stride - face_landmark_5_raw = detections[index + feature_map_channel * 2] * feature_stride - for bounding_box in distance_to_bounding_box(anchors, bounding_box_raw)[keep_indices]: - bounding_box_list.append(numpy.array( - [ - bounding_box[0] * ratio_width, - bounding_box[1] * ratio_height, - bounding_box[2] * ratio_width, - bounding_box[3] * ratio_height - ])) - for face_landmark_5 in distance_to_face_landmark_5(anchors, face_landmark_5_raw)[keep_indices]: - face_landmark_5_list.append(face_landmark_5 * [ ratio_width, ratio_height ]) - for score in detections[index][keep_indices]: - score_list.append(score[0]) - return bounding_box_list, face_landmark_5_list, score_list - - -def detect_with_scrfd(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[FaceLandmark5], List[Score]]: - face_detector = get_face_analyser().get('face_detectors').get('scrfd') - face_detector_width, face_detector_height = unpack_resolution(face_detector_size) - temp_vision_frame = resize_frame_resolution(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] - feature_strides = [ 8, 16, 32 ] - feature_map_channel = 3 - anchor_total = 2 - bounding_box_list = [] - face_landmark_5_list = [] - score_list = [] - - detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) - with thread_semaphore(): - detections = face_detector.run(None, - { - face_detector.get_inputs()[0].name: detect_vision_frame - }) - for index, feature_stride in enumerate(feature_strides): - keep_indices = numpy.where(detections[index] >= facefusion.globals.face_detector_score)[0] - if keep_indices.any(): - 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_box_raw = detections[index + feature_map_channel] * feature_stride - face_landmark_5_raw = detections[index + feature_map_channel * 2] * feature_stride - for bounding_box in distance_to_bounding_box(anchors, bounding_box_raw)[keep_indices]: - bounding_box_list.append(numpy.array( - [ - bounding_box[0] * ratio_width, - bounding_box[1] * ratio_height, - bounding_box[2] * ratio_width, - bounding_box[3] * ratio_height - ])) - for face_landmark_5 in distance_to_face_landmark_5(anchors, face_landmark_5_raw)[keep_indices]: - face_landmark_5_list.append(face_landmark_5 * [ ratio_width, ratio_height ]) - for score in detections[index][keep_indices]: - score_list.append(score[0]) - return bounding_box_list, face_landmark_5_list, score_list - - -def detect_with_yoloface(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[FaceLandmark5], List[Score]]: - face_detector = get_face_analyser().get('face_detectors').get('yoloface') - face_detector_width, face_detector_height = unpack_resolution(face_detector_size) - temp_vision_frame = resize_frame_resolution(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] - bounding_box_list = [] - face_landmark_5_list = [] - score_list = [] - - detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) - with thread_semaphore(): - detections = face_detector.run(None, - { - face_detector.get_inputs()[0].name: detect_vision_frame - }) - detections = numpy.squeeze(detections).T - bounding_box_raw, score_raw, face_landmark_5_raw = numpy.split(detections, [ 4, 5 ], axis = 1) - keep_indices = numpy.where(score_raw > facefusion.globals.face_detector_score)[0] - if keep_indices.any(): - bounding_box_raw, face_landmark_5_raw, score_raw = bounding_box_raw[keep_indices], face_landmark_5_raw[keep_indices], score_raw[keep_indices] - for bounding_box in bounding_box_raw: - bounding_box_list.append(numpy.array( - [ - (bounding_box[0] - bounding_box[2] / 2) * ratio_width, - (bounding_box[1] - bounding_box[3] / 2) * ratio_height, - (bounding_box[0] + bounding_box[2] / 2) * ratio_width, - (bounding_box[1] + bounding_box[3] / 2) * ratio_height - ])) - face_landmark_5_raw[:, 0::3] = (face_landmark_5_raw[:, 0::3]) * ratio_width - face_landmark_5_raw[:, 1::3] = (face_landmark_5_raw[:, 1::3]) * ratio_height - for face_landmark_5 in face_landmark_5_raw: - face_landmark_5_list.append(numpy.array(face_landmark_5.reshape(-1, 3)[:, :2])) - score_list = score_raw.ravel().tolist() - return bounding_box_list, face_landmark_5_list, score_list - - -def detect_with_yunet(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[FaceLandmark5], List[Score]]: - face_detector = get_face_analyser().get('face_detectors').get('yunet') - face_detector_width, face_detector_height = unpack_resolution(face_detector_size) - temp_vision_frame = resize_frame_resolution(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] - bounding_box_list = [] - face_landmark_5_list = [] - score_list = [] - - face_detector.setInputSize((temp_vision_frame.shape[1], temp_vision_frame.shape[0])) - face_detector.setScoreThreshold(facefusion.globals.face_detector_score) - with thread_semaphore(): - _, detections = face_detector.detect(temp_vision_frame) - if numpy.any(detections): - for detection in detections: - bounding_box_list.append(numpy.array( - [ - detection[0] * ratio_width, - detection[1] * ratio_height, - (detection[0] + detection[2]) * ratio_width, - (detection[1] + detection[3]) * ratio_height - ])) - face_landmark_5_list.append(detection[4:14].reshape((5, 2)) * [ ratio_width, ratio_height ]) - score_list.append(detection[14]) - return bounding_box_list, face_landmark_5_list, score_list - - -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 = (detect_vision_frame - 127.5) / 128.0 - detect_vision_frame = numpy.expand_dims(detect_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) - return detect_vision_frame - - -def create_faces(vision_frame : VisionFrame, bounding_box_list : List[BoundingBox], face_landmark_5_list : List[FaceLandmark5], score_list : List[Score]) -> List[Face]: +def create_faces(vision_frame : VisionFrame, bounding_boxes : List[BoundingBox], face_scores : List[Score], face_landmarks_5 : List[FaceLandmark5]) -> List[Face]: faces = [] - if facefusion.globals.face_detector_score > 0: - sort_indices = numpy.argsort(-numpy.array(score_list)) - bounding_box_list = [ bounding_box_list[index] for index in sort_indices ] - face_landmark_5_list = [face_landmark_5_list[index] for index in sort_indices] - score_list = [ score_list[index] for index in sort_indices ] - iou_threshold = 0.1 if facefusion.globals.face_detector_model == 'many' else 0.4 - keep_indices = apply_nms(bounding_box_list, iou_threshold) - for index in keep_indices: - bounding_box = bounding_box_list[index] - face_landmark_5_68 = face_landmark_5_list[index] - face_landmark_68_5 = expand_face_landmark_68_from_5(face_landmark_5_68) - face_landmark_68 = face_landmark_68_5 - face_landmark_68_score = 0.0 - if facefusion.globals.face_landmarker_score > 0: - face_landmark_68, face_landmark_68_score = detect_face_landmark_68(vision_frame, bounding_box) - if face_landmark_68_score > facefusion.globals.face_landmarker_score: - face_landmark_5_68 = convert_face_landmark_68_to_5(face_landmark_68) - landmarks : FaceLandmarkSet =\ - { - '5': face_landmark_5_list[index], - '5/68': face_landmark_5_68, - '68': face_landmark_68, - '68/5': face_landmark_68_5 - } - scores : FaceScoreSet = \ - { - 'detector': score_list[index], - 'landmarker': face_landmark_68_score - } - embedding, normed_embedding = calc_embedding(vision_frame, landmarks.get('5/68')) - gender, age = detect_gender_age(vision_frame, bounding_box) - faces.append(Face( - bounding_box = bounding_box, - landmarks = landmarks, - scores = scores, - embedding = embedding, - normed_embedding = normed_embedding, - gender = gender, - age = age - )) + 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_landmarks(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 + } + embedding, normed_embedding = calc_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 = embedding, + normed_embedding = normed_embedding, + gender = gender, + age = age, + race = race + )) return faces -def calc_embedding(temp_vision_frame : VisionFrame, face_landmark_5 : FaceLandmark5) -> Tuple[Embedding, Embedding]: - face_recognizer = get_face_analyser().get('face_recognizer') - crop_vision_frame, matrix = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5, 'arcface_112_v2', (112, 112)) - 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) - with conditional_thread_semaphore(facefusion.globals.execution_providers): - embedding = face_recognizer.run(None, - { - face_recognizer.get_inputs()[0].name: crop_vision_frame - })[0] - embedding = embedding.ravel() - normed_embedding = embedding / numpy.linalg.norm(embedding) - return embedding, normed_embedding - - -def detect_face_landmark_68(temp_vision_frame : VisionFrame, bounding_box : BoundingBox) -> Tuple[FaceLandmark68, Score]: - face_landmarker = get_face_analyser().get('face_landmarkers').get('68') - scale = 195 / numpy.subtract(bounding_box[2:], bounding_box[:2]).max() - translation = (256 - numpy.add(bounding_box[2:], bounding_box[:2]) * scale) * 0.5 - crop_vision_frame, affine_matrix = warp_face_by_translation(temp_vision_frame, translation, scale, (256, 256)) - crop_vision_frame = cv2.cvtColor(crop_vision_frame, cv2.COLOR_RGB2Lab) - if numpy.mean(crop_vision_frame[:, :, 0]) < 30: - crop_vision_frame[:, :, 0] = cv2.createCLAHE(clipLimit = 2).apply(crop_vision_frame[:, :, 0]) - crop_vision_frame = cv2.cvtColor(crop_vision_frame, cv2.COLOR_Lab2RGB) - crop_vision_frame = crop_vision_frame.transpose(2, 0, 1).astype(numpy.float32) / 255.0 - with conditional_thread_semaphore(facefusion.globals.execution_providers): - face_landmark_68, face_heatmap = face_landmarker.run(None, - { - face_landmarker.get_inputs()[0].name: [ crop_vision_frame ] - }) - face_landmark_68 = face_landmark_68[:, :, :2][0] / 64 - face_landmark_68 = face_landmark_68.reshape(1, -1, 2) * 256 - face_landmark_68 = cv2.transform(face_landmark_68, cv2.invertAffineTransform(affine_matrix)) - face_landmark_68 = face_landmark_68.reshape(-1, 2) - face_landmark_68_score = numpy.amax(face_heatmap, axis = (2, 3)) - face_landmark_68_score = numpy.mean(face_landmark_68_score) - return face_landmark_68, face_landmark_68_score - - -def expand_face_landmark_68_from_5(face_landmark_5 : FaceLandmark5) -> FaceLandmark68: - face_landmarker = get_face_analyser().get('face_landmarkers').get('68_5') - 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) - with conditional_thread_semaphore(facefusion.globals.execution_providers): - face_landmark_68_5 = face_landmarker.run(None, - { - face_landmarker.get_inputs()[0].name: [ face_landmark_5 ] - })[0][0] - 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 detect_gender_age(temp_vision_frame : VisionFrame, bounding_box : BoundingBox) -> Tuple[int, int]: - gender_age = get_face_analyser().get('gender_age') - bounding_box = bounding_box.reshape(2, -1) - scale = 64 / numpy.subtract(*bounding_box[::-1]).max() - translation = 48 - bounding_box.sum(axis = 0) * scale * 0.5 - crop_vision_frame, affine_matrix = warp_face_by_translation(temp_vision_frame, translation, scale, (96, 96)) - 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) - with conditional_thread_semaphore(facefusion.globals.execution_providers): - prediction = gender_age.run(None, - { - gender_age.get_inputs()[0].name: crop_vision_frame - })[0][0] - gender = int(numpy.argmax(prediction[:2])) - age = int(numpy.round(prediction[2] * 100)) - return gender, age - - -def get_one_face(vision_frame : VisionFrame, position : int = 0) -> Optional[Face]: - many_faces = get_many_faces(vision_frame) - if many_faces: - try: - return many_faces[position] - except IndexError: - return many_faces[-1] +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(vision_frames : List[VisionFrame], position : int = 0) -> Optional[Face]: - average_face = None - faces = [] - embedding_list = [] - normed_embedding_list = [] +def get_average_face(faces : List[Face]) -> Optional[Face]: + embeddings = [] + normed_embeddings = [] - for vision_frame in vision_frames: - face = get_one_face(vision_frame, position) - if face: - faces.append(face) - embedding_list.append(face.embedding) - normed_embedding_list.append(face.normed_embedding) if faces: first_face = get_first(faces) - average_face = Face( + + for face in faces: + embeddings.append(face.embedding) + normed_embeddings.append(face.normed_embedding) + + return Face( bounding_box = first_face.bounding_box, - landmarks = first_face.landmarks, - scores = first_face.scores, - embedding = numpy.mean(embedding_list, axis = 0), - normed_embedding = numpy.mean(normed_embedding_list, axis = 0), + score_set = first_face.score_set, + landmark_set = first_face.landmark_set, + angle = first_face.angle, + embedding = numpy.mean(embeddings, axis = 0), + normed_embedding = numpy.mean(normed_embeddings, axis = 0), gender = first_face.gender, - age = first_face.age + age = first_face.age, + race = first_face.race ) - return average_face + return None -def get_many_faces(vision_frame : VisionFrame) -> List[Face]: - faces = [] - try: - faces_cache = get_static_faces(vision_frame) - if faces_cache: - faces = faces_cache - else: - bounding_box_list = [] - face_landmark_5_list = [] - score_list = [] +def get_many_faces(vision_frames : List[VisionFrame]) -> List[Face]: + many_faces : List[Face] = [] - if facefusion.globals.face_detector_model in [ 'many', 'retinaface']: - bounding_box_list_retinaface, face_landmark_5_list_retinaface, score_list_retinaface = detect_with_retinaface(vision_frame, facefusion.globals.face_detector_size) - bounding_box_list.extend(bounding_box_list_retinaface) - face_landmark_5_list.extend(face_landmark_5_list_retinaface) - score_list.extend(score_list_retinaface) - if facefusion.globals.face_detector_model in [ 'many', 'scrfd' ]: - bounding_box_list_scrfd, face_landmark_5_list_scrfd, score_list_scrfd = detect_with_scrfd(vision_frame, facefusion.globals.face_detector_size) - bounding_box_list.extend(bounding_box_list_scrfd) - face_landmark_5_list.extend(face_landmark_5_list_scrfd) - score_list.extend(score_list_scrfd) - if facefusion.globals.face_detector_model in [ 'many', 'yoloface' ]: - bounding_box_list_yoloface, face_landmark_5_list_yoloface, score_list_yoloface = detect_with_yoloface(vision_frame, facefusion.globals.face_detector_size) - bounding_box_list.extend(bounding_box_list_yoloface) - face_landmark_5_list.extend(face_landmark_5_list_yoloface) - score_list.extend(score_list_yoloface) - if facefusion.globals.face_detector_model in [ 'yunet' ]: - bounding_box_list_yunet, face_landmark_5_list_yunet, score_list_yunet = detect_with_yunet(vision_frame, facefusion.globals.face_detector_size) - bounding_box_list.extend(bounding_box_list_yunet) - face_landmark_5_list.extend(face_landmark_5_list_yunet) - score_list.extend(score_list_yunet) - if bounding_box_list and face_landmark_5_list and score_list: - faces = create_faces(vision_frame, bounding_box_list, face_landmark_5_list, score_list) - if faces: - set_static_faces(vision_frame, faces) - if facefusion.globals.face_analyser_order: - faces = sort_by_order(faces, facefusion.globals.face_analyser_order) - if facefusion.globals.face_analyser_age: - faces = filter_by_age(faces, facefusion.globals.face_analyser_age) - if facefusion.globals.face_analyser_gender: - faces = filter_by_gender(faces, facefusion.globals.face_analyser_gender) - except (AttributeError, ValueError): - pass - return faces + 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_rotated_faces(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) -def find_similar_faces(reference_faces : FaceSet, vision_frame : VisionFrame, face_distance : float) -> List[Face]: - similar_faces : List[Face] = [] - many_faces = get_many_faces(vision_frame) + 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 reference_faces: - for reference_set in reference_faces: - if not similar_faces: - for reference_face in reference_faces[reference_set]: - for face in many_faces: - if compare_faces(face, reference_face, face_distance): - similar_faces.append(face) - return similar_faces - - -def compare_faces(face : Face, reference_face : Face, face_distance : float) -> bool: - current_face_distance = calc_face_distance(face, reference_face) - return current_face_distance < face_distance - - -def calc_face_distance(face : Face, reference_face : Face) -> float: - if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'): - return 1 - numpy.dot(face.normed_embedding, reference_face.normed_embedding) - return 0 - - -def sort_by_order(faces : List[Face], order : FaceAnalyserOrder) -> List[Face]: - if order == 'left-right': - return sorted(faces, key = lambda face: face.bounding_box[0]) - if order == 'right-left': - return sorted(faces, key = lambda face: face.bounding_box[0], reverse = True) - if order == 'top-bottom': - return sorted(faces, key = lambda face: face.bounding_box[1]) - if order == 'bottom-top': - return sorted(faces, key = lambda face: face.bounding_box[1], reverse = True) - if order == 'small-large': - return sorted(faces, key = lambda face: (face.bounding_box[2] - face.bounding_box[0]) * (face.bounding_box[3] - face.bounding_box[1])) - if order == 'large-small': - return sorted(faces, key = lambda face: (face.bounding_box[2] - face.bounding_box[0]) * (face.bounding_box[3] - face.bounding_box[1]), reverse = True) - if order == 'best-worst': - return sorted(faces, key = lambda face: face.scores.get('detector'), reverse = True) - if order == 'worst-best': - return sorted(faces, key = lambda face: face.scores.get('detector')) - return faces - - -def filter_by_age(faces : List[Face], age : FaceAnalyserAge) -> List[Face]: - filter_faces = [] - for face in faces: - if categorize_age(face.age) == age: - filter_faces.append(face) - return filter_faces - - -def filter_by_gender(faces : List[Face], gender : FaceAnalyserGender) -> List[Face]: - filter_faces = [] - for face in faces: - if categorize_gender(face.gender) == gender: - filter_faces.append(face) - return filter_faces + if faces: + many_faces.extend(faces) + set_static_faces(vision_frame, faces) + return many_faces diff --git a/facefusion/face_classifier.py b/facefusion/face_classifier.py new file mode 100644 index 00000000..2addbe4b --- /dev/null +++ b/facefusion/face_classifier.py @@ -0,0 +1,128 @@ +from typing import List, Tuple + +import numpy + +from facefusion import inference_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources +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.typing import Age, FaceLandmark5, Gender, InferencePool, ModelOptions, ModelSet, Race, VisionFrame + +MODEL_SET : ModelSet =\ +{ + 'fairface': + { + 'hashes': + { + 'face_classifier': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/fairface.hash', + 'path': resolve_relative_path('../.assets/models/fairface.hash') + } + }, + 'sources': + { + 'face_classifier': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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_sources = get_model_options().get('sources') + return inference_manager.get_inference_pool(__name__, model_sources) + + +def clear_inference_pool() -> None: + inference_manager.clear_inference_pool(__name__) + + +def get_model_options() -> ModelOptions: + return MODEL_SET.get('fairface') + + +def pre_check() -> bool: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +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 + 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 00000000..09d104b2 --- /dev/null +++ b/facefusion/face_detector.py @@ -0,0 +1,309 @@ +from typing import List, Tuple + +import cv2 +import numpy + +from facefusion import inference_manager, state_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.face_helper import create_rotated_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.typing import Angle, BoundingBox, Detection, DownloadSet, FaceLandmark5, InferencePool, ModelSet, Score, VisionFrame +from facefusion.vision import resize_frame_resolution, unpack_resolution + +MODEL_SET : ModelSet =\ +{ + 'retinaface': + { + 'hashes': + { + 'retinaface': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/retinaface_10g.hash', + 'path': resolve_relative_path('../.assets/models/retinaface_10g.hash') + } + }, + 'sources': + { + 'retinaface': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/retinaface_10g.onnx', + 'path': resolve_relative_path('../.assets/models/retinaface_10g.onnx') + } + } + }, + 'scrfd': + { + 'hashes': + { + 'scrfd': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/scrfd_2.5g.hash', + 'path': resolve_relative_path('../.assets/models/scrfd_2.5g.hash') + } + }, + 'sources': + { + 'scrfd': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/scrfd_2.5g.onnx', + 'path': resolve_relative_path('../.assets/models/scrfd_2.5g.onnx') + } + } + }, + 'yoloface': + { + 'hashes': + { + 'yoloface': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/yoloface_8n.hash', + 'path': resolve_relative_path('../.assets/models/yoloface_8n.hash') + } + }, + 'sources': + { + 'yoloface': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/yoloface_8n.onnx', + 'path': resolve_relative_path('../.assets/models/yoloface_8n.onnx') + } + } + } +} + + +def get_inference_pool() -> InferencePool: + _, model_sources = collect_model_downloads() + model_context = __name__ + '.' + state_manager.get_item('face_detector_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('face_detector_model') + inference_manager.clear_inference_pool(model_context) + + +def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_hashes = {} + model_sources = {} + + if state_manager.get_item('face_detector_model') in [ 'many', 'retinaface' ]: + model_hashes['retinaface'] = MODEL_SET.get('retinaface').get('hashes').get('retinaface') + model_sources['retinaface'] = MODEL_SET.get('retinaface').get('sources').get('retinaface') + if state_manager.get_item('face_detector_model') in [ 'many', 'scrfd' ]: + model_hashes['scrfd'] = MODEL_SET.get('scrfd').get('hashes').get('scrfd') + model_sources['scrfd'] = MODEL_SET.get('scrfd').get('sources').get('scrfd') + if state_manager.get_item('face_detector_model') in [ 'many', 'yoloface' ]: + model_hashes['yoloface'] = MODEL_SET.get('yoloface').get('hashes').get('yoloface') + model_sources['yoloface'] = MODEL_SET.get('yoloface').get('sources').get('yoloface') + return model_hashes, model_sources + + +def pre_check() -> bool: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes, model_sources = collect_model_downloads() + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +def detect_faces(vision_frame : VisionFrame) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + 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(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(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', 'yoloface' ]: + bounding_boxes, face_scores, face_landmarks_5 = detect_with_yoloface(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) for all_bounding_box in all_bounding_boxes ] + return all_bounding_boxes, all_face_scores, all_face_landmarks_5 + + +def detect_rotated_faces(vision_frame : VisionFrame, angle : Angle) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + rotated_matrix, rotated_size = create_rotated_matrix_and_size(angle, vision_frame.shape[:2][::-1]) + rotated_vision_frame = cv2.warpAffine(vision_frame, rotated_matrix, rotated_size) + rotated_inverse_matrix = cv2.invertAffineTransform(rotated_matrix) + bounding_boxes, face_scores, face_landmarks_5 = detect_faces(rotated_vision_frame) + bounding_boxes = [ transform_bounding_box(bounding_box, rotated_inverse_matrix) for bounding_box in bounding_boxes ] + face_landmarks_5 = [ transform_points(face_landmark_5, rotated_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_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = resize_frame_resolution(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) + detection = forward_with_retinaface(detect_vision_frame) + + for index, feature_stride in enumerate(feature_strides): + keep_indices = numpy.where(detection[index] >= state_manager.get_item('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_box_raw = detection[index + feature_map_channel] * feature_stride + face_landmark_5_raw = detection[index + feature_map_channel * 2] * feature_stride + + for bounding_box in distance_to_bounding_box(anchors, bounding_box_raw)[keep_indices]: + bounding_boxes.append(numpy.array( + [ + bounding_box[0] * ratio_width, + bounding_box[1] * ratio_height, + bounding_box[2] * ratio_width, + bounding_box[3] * ratio_height, + ])) + + for score in detection[index][keep_indices]: + face_scores.append(score[0]) + + for face_landmark_5 in distance_to_face_landmark_5(anchors, face_landmark_5_raw)[keep_indices]: + face_landmarks_5.append(face_landmark_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_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = resize_frame_resolution(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) + detection = forward_with_scrfd(detect_vision_frame) + + for index, feature_stride in enumerate(feature_strides): + keep_indices = numpy.where(detection[index] >= state_manager.get_item('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_box_raw = detection[index + feature_map_channel] * feature_stride + face_landmark_5_raw = detection[index + feature_map_channel * 2] * feature_stride + + for bounding_box in distance_to_bounding_box(anchors, bounding_box_raw)[keep_indices]: + bounding_boxes.append(numpy.array( + [ + bounding_box[0] * ratio_width, + bounding_box[1] * ratio_height, + bounding_box[2] * ratio_width, + bounding_box[3] * ratio_height, + ])) + + for score in detection[index][keep_indices]: + face_scores.append(score[0]) + + for face_landmark_5 in distance_to_face_landmark_5(anchors, face_landmark_5_raw)[keep_indices]: + face_landmarks_5.append(face_landmark_5 * [ ratio_width, ratio_height ]) + + return bounding_boxes, face_scores, face_landmarks_5 + + +def detect_with_yoloface(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[Score], List[FaceLandmark5]]: + bounding_boxes = [] + face_scores = [] + face_landmarks_5 = [] + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = resize_frame_resolution(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) + detection = forward_with_yoloface(detect_vision_frame) + detection = numpy.squeeze(detection).T + bounding_box_raw, score_raw, face_landmark_5_raw = numpy.split(detection, [ 4, 5 ], axis = 1) + keep_indices = numpy.where(score_raw > state_manager.get_item('face_detector_score'))[0] + + if numpy.any(keep_indices): + bounding_box_raw, face_landmark_5_raw, score_raw = bounding_box_raw[keep_indices], face_landmark_5_raw[keep_indices], score_raw[keep_indices] + + for bounding_box in bounding_box_raw: + bounding_boxes.append(numpy.array( + [ + (bounding_box[0] - bounding_box[2] / 2) * ratio_width, + (bounding_box[1] - bounding_box[3] / 2) * ratio_height, + (bounding_box[0] + bounding_box[2] / 2) * ratio_width, + (bounding_box[1] + bounding_box[3] / 2) * ratio_height, + ])) + + face_scores = score_raw.ravel().tolist() + face_landmark_5_raw[:, 0::3] = (face_landmark_5_raw[:, 0::3]) * ratio_width + face_landmark_5_raw[:, 1::3] = (face_landmark_5_raw[:, 1::3]) * ratio_height + + for face_landmark_5 in face_landmark_5_raw: + face_landmarks_5.append(numpy.array(face_landmark_5.reshape(-1, 3)[:, :2])) + + 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_yoloface(detect_vision_frame : VisionFrame) -> Detection: + face_detector = get_inference_pool().get('yoloface') + + 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 = (detect_vision_frame - 127.5) / 128.0 + 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/face_helper.py b/facefusion/face_helper.py index 83eff56e..1265cac6 100644 --- a/facefusion/face_helper.py +++ b/facefusion/face_helper.py @@ -1,10 +1,11 @@ -from typing import Any, Tuple, List -from cv2.typing import Size from functools import lru_cache +from typing import List, Sequence, Tuple + import cv2 import numpy +from cv2.typing import Size -from facefusion.typing import BoundingBox, FaceLandmark5, FaceLandmark68, VisionFrame, Mask, Matrix, Translation, WarpTemplate, WarpTemplateSet, FaceAnalyserAge, FaceAnalyserGender +from facefusion.typing import Anchors, Angle, BoundingBox, Distance, FaceDetectorModel, FaceLandmark5, FaceLandmark68, Mask, Matrix, Points, Scale, Score, Translation, VisionFrame, WarpTemplate, WarpTemplateSet WARP_TEMPLATES : WarpTemplateSet =\ { @@ -86,7 +87,7 @@ def paste_back(temp_vision_frame : VisionFrame, crop_vision_frame : VisionFrame, @lru_cache(maxsize = None) -def create_static_anchors(feature_stride : int, anchor_total : int, stride_height : int, stride_width : int) -> numpy.ndarray[Any, Any]: +def create_static_anchors(feature_stride : int, anchor_total : int, stride_height : int, stride_width : int) -> Anchors: y, x = numpy.mgrid[:stride_height, :stride_width][::-1] anchors = numpy.stack((y, x), axis = -1) anchors = (anchors * feature_stride).reshape((-1, 2)) @@ -94,14 +95,50 @@ def create_static_anchors(feature_stride : int, anchor_total : int, stride_heigh return anchors -def create_bounding_box_from_face_landmark_68(face_landmark_68 : FaceLandmark68) -> BoundingBox: +def create_rotated_matrix_and_size(angle : Angle, size : Size) -> Tuple[Matrix, Size]: + rotated_matrix = cv2.getRotationMatrix2D((size[0] / 2, size[1] / 2), angle, 1) + rotated_size = numpy.dot(numpy.abs(rotated_matrix[:, :2]), size) + rotated_matrix[:, -1] += (rotated_size - size) * 0.5 #type:ignore[misc] + rotated_size = int(rotated_size[0]), int(rotated_size[1]) + return rotated_matrix, rotated_size + + +def create_bounding_box(face_landmark_68 : FaceLandmark68) -> BoundingBox: min_x, min_y = numpy.min(face_landmark_68, axis = 0) max_x, max_y = numpy.max(face_landmark_68, axis = 0) - bounding_box = numpy.array([ min_x, min_y, max_x, max_y ]).astype(numpy.int16) + bounding_box = normalize_bounding_box(numpy.array([ min_x, min_y, max_x, max_y ])) return bounding_box -def distance_to_bounding_box(points : numpy.ndarray[Any, Any], distance : numpy.ndarray[Any, Any]) -> BoundingBox: +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] @@ -110,14 +147,21 @@ def distance_to_bounding_box(points : numpy.ndarray[Any, Any], distance : numpy. return bounding_box -def distance_to_face_landmark_5(points : numpy.ndarray[Any, Any], distance : numpy.ndarray[Any, Any]) -> FaceLandmark5: +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 convert_face_landmark_68_to_5(face_landmark_68 : FaceLandmark68) -> FaceLandmark5: +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), @@ -129,41 +173,38 @@ def convert_face_landmark_68_to_5(face_landmark_68 : FaceLandmark68) -> FaceLand return face_landmark_5 -def apply_nms(bounding_box_list : List[BoundingBox], iou_threshold : float) -> List[int]: - keep_indices = [] - dimension_list = numpy.reshape(bounding_box_list, (-1, 4)) - x1 = dimension_list[:, 0] - y1 = dimension_list[:, 1] - x2 = dimension_list[:, 2] - y2 = dimension_list[:, 3] - areas = (x2 - x1 + 1) * (y2 - y1 + 1) - indices = numpy.arange(len(bounding_box_list)) - while indices.size > 0: - index = indices[0] - remain_indices = indices[1:] - keep_indices.append(index) - xx1 = numpy.maximum(x1[index], x1[remain_indices]) - yy1 = numpy.maximum(y1[index], y1[remain_indices]) - xx2 = numpy.minimum(x2[index], x2[remain_indices]) - yy2 = numpy.minimum(y2[index], y2[remain_indices]) - width = numpy.maximum(0, xx2 - xx1 + 1) - height = numpy.maximum(0, yy2 - yy1 + 1) - iou = width * height / (areas[index] + areas[remain_indices] - width * height) - indices = indices[numpy.where(iou <= iou_threshold)[0] + 1] +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], face_scores : List[Score], score_threshold : float, nms_threshold : float) -> Sequence[int]: + normed_bounding_boxes = [ (x1, y1, x2 - x1, y2 - y1) for (x1, y1, x2, y2) in bounding_boxes ] + keep_indices = cv2.dnn.NMSBoxes(normed_bounding_boxes, face_scores, score_threshold = score_threshold, nms_threshold = nms_threshold) return keep_indices -def categorize_age(age : int) -> FaceAnalyserAge: - if age < 13: - return 'child' - elif age < 19: - return 'teen' - elif age < 60: - return 'adult' - return 'senior' +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 categorize_gender(gender : int) -> FaceAnalyserGender: - if gender == 0: - return 'female' - return 'male' +def merge_matrix(matrices : List[Matrix]) -> Matrix: + merged_matrix = numpy.vstack([ matrices[0], [ 0, 0, 1 ] ]) + for matrix in matrices[1:]: + matrix = numpy.vstack([ matrix, [ 0, 0, 1 ] ]) + merged_matrix = numpy.dot(merged_matrix, matrix) + return merged_matrix[:2, :] diff --git a/facefusion/face_landmarker.py b/facefusion/face_landmarker.py new file mode 100644 index 00000000..6760c271 --- /dev/null +++ b/facefusion/face_landmarker.py @@ -0,0 +1,217 @@ +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 +from facefusion.face_helper import create_rotated_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.typing import Angle, BoundingBox, DownloadSet, FaceLandmark5, FaceLandmark68, InferencePool, ModelSet, Prediction, Score, VisionFrame + +MODEL_SET : ModelSet =\ +{ + '2dfan4': + { + 'hashes': + { + '2dfan4': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/2dfan4.hash', + 'path': resolve_relative_path('../.assets/models/2dfan4.hash') + } + }, + 'sources': + { + '2dfan4': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/2dfan4.onnx', + 'path': resolve_relative_path('../.assets/models/2dfan4.onnx') + } + }, + 'size': (256, 256) + }, + 'peppa_wutz': + { + 'hashes': + { + 'peppa_wutz': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/peppa_wutz.hash', + 'path': resolve_relative_path('../.assets/models/peppa_wutz.hash') + } + }, + 'sources': + { + 'peppa_wutz': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/peppa_wutz.onnx', + 'path': resolve_relative_path('../.assets/models/peppa_wutz.onnx') + } + }, + 'size': (256, 256) + }, + 'fan_68_5': + { + 'hashes': + { + 'fan_68_5': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/fan_68_5.hash', + 'path': resolve_relative_path('../.assets/models/fan_68_5.hash') + } + }, + 'sources': + { + 'fan_68_5': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/fan_68_5.onnx', + 'path': resolve_relative_path('../.assets/models/fan_68_5.onnx') + } + } + } +} + + +def get_inference_pool() -> InferencePool: + _, model_sources = collect_model_downloads() + model_context = __name__ + '.' + state_manager.get_item('face_landmarker_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('face_landmarker_model') + inference_manager.clear_inference_pool(model_context) + + +def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_hashes =\ + { + 'fan_68_5': MODEL_SET.get('fan_68_5').get('hashes').get('fan_68_5') + } + model_sources =\ + { + 'fan_68_5': MODEL_SET.get('fan_68_5').get('sources').get('fan_68_5') + } + + if state_manager.get_item('face_landmarker_model') in [ 'many', '2dfan4' ]: + model_hashes['2dfan4'] = MODEL_SET.get('2dfan4').get('hashes').get('2dfan4') + model_sources['2dfan4'] = MODEL_SET.get('2dfan4').get('sources').get('2dfan4') + if state_manager.get_item('face_landmarker_model') in [ 'many', 'peppa_wutz' ]: + model_hashes['peppa_wutz'] = MODEL_SET.get('peppa_wutz').get('hashes').get('peppa_wutz') + model_sources['peppa_wutz'] = MODEL_SET.get('peppa_wutz').get('sources').get('peppa_wutz') + return model_hashes, model_sources + + +def pre_check() -> bool: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes, model_sources = collect_model_downloads() + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +def detect_face_landmarks(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 = MODEL_SET.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 + rotated_matrix, rotated_size = create_rotated_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, rotated_matrix, rotated_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(rotated_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 = MODEL_SET.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 + rotated_matrix, rotated_size = create_rotated_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, rotated_matrix, rotated_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(rotated_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 index f606ab8c..deffdb72 100755 --- a/facefusion/face_masker.py +++ b/facefusion/face_masker.py @@ -1,32 +1,57 @@ -from typing import Any, Dict, List -from cv2.typing import Size from functools import lru_cache -from time import sleep +from typing import Dict, List, Tuple + import cv2 import numpy -import onnxruntime +from cv2.typing import Size -import facefusion.globals -from facefusion import process_manager -from facefusion.thread_helper import thread_lock, conditional_thread_semaphore -from facefusion.typing import FaceLandmark68, VisionFrame, Mask, Padding, FaceMaskRegion, ModelSet -from facefusion.execution import apply_execution_provider_options -from facefusion.filesystem import resolve_relative_path, is_file -from facefusion.download import conditional_download +from facefusion import inference_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.filesystem import resolve_relative_path +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.typing import DownloadSet, FaceLandmark68, FaceMaskRegion, InferencePool, Mask, ModelSet, Padding, VisionFrame -FACE_OCCLUDER = None -FACE_PARSER = None -MODELS : ModelSet =\ +MODEL_SET : ModelSet =\ { 'face_occluder': { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_occluder.onnx', - 'path': resolve_relative_path('../.assets/models/face_occluder.onnx') + 'hashes': + { + 'face_occluder': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/dfl_xseg.hash', + 'path': resolve_relative_path('../.assets/models/dfl_xseg.hash') + } + }, + 'sources': + { + 'face_occluder': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/dfl_xseg.onnx', + 'path': resolve_relative_path('../.assets/models/dfl_xseg.onnx') + } + }, + 'size': (256, 256) }, 'face_parser': { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_parser.onnx', - 'path': resolve_relative_path('../.assets/models/face_parser.onnx') + 'hashes': + { + 'face_parser': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/resnet_34.hash', + 'path': resolve_relative_path('../.assets/models/resnet_34.hash') + } + }, + 'sources': + { + 'face_parser': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/resnet_34.onnx', + 'path': resolve_relative_path('../.assets/models/resnet_34.onnx') + } + }, + 'size': (512, 512) } } FACE_MASK_REGIONS : Dict[FaceMaskRegion, int] =\ @@ -44,67 +69,41 @@ FACE_MASK_REGIONS : Dict[FaceMaskRegion, int] =\ } -def get_face_occluder() -> Any: - global FACE_OCCLUDER - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if FACE_OCCLUDER is None: - model_path = MODELS.get('face_occluder').get('path') - FACE_OCCLUDER = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return FACE_OCCLUDER +def get_inference_pool() -> InferencePool: + _, model_sources = collect_model_downloads() + return inference_manager.get_inference_pool(__name__, model_sources) -def get_face_parser() -> Any: - global FACE_PARSER - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if FACE_PARSER is None: - model_path = MODELS.get('face_parser').get('path') - FACE_PARSER = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return FACE_PARSER +def clear_inference_pool() -> None: + inference_manager.clear_inference_pool(__name__) -def clear_face_occluder() -> None: - global FACE_OCCLUDER - - FACE_OCCLUDER = None - - -def clear_face_parser() -> None: - global FACE_PARSER - - FACE_PARSER = None +def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_hashes =\ + { + 'face_occluder': MODEL_SET.get('face_occluder').get('hashes').get('face_occluder'), + 'face_parser': MODEL_SET.get('face_parser').get('hashes').get('face_parser') + } + model_sources =\ + { + 'face_occluder': MODEL_SET.get('face_occluder').get('sources').get('face_occluder'), + 'face_parser': MODEL_SET.get('face_parser').get('sources').get('face_parser') + } + return model_hashes, model_sources def pre_check() -> bool: download_directory_path = resolve_relative_path('../.assets/models') - model_urls =\ - [ - MODELS.get('face_occluder').get('url'), - MODELS.get('face_parser').get('url') - ] - model_paths =\ - [ - MODELS.get('face_occluder').get('path'), - MODELS.get('face_parser').get('path') - ] + model_hashes, model_sources = collect_model_downloads() - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, model_urls) - process_manager.end() - return all(is_file(model_path) for model_path in model_paths) + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) @lru_cache(maxsize = None) def create_static_box_mask(crop_size : Size, face_mask_blur : float, face_mask_padding : Padding) -> Mask: 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, numpy.float32) + 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 @@ -115,15 +114,11 @@ def create_static_box_mask(crop_size : Size, face_mask_blur : float, face_mask_p def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask: - face_occluder = get_face_occluder() - prepare_vision_frame = cv2.resize(crop_vision_frame, face_occluder.get_inputs()[0].shape[1:3][::-1]) + model_size = MODEL_SET.get('face_occluder').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 prepare_vision_frame = prepare_vision_frame.transpose(0, 1, 2, 3) - with conditional_thread_semaphore(facefusion.globals.execution_providers): - occlusion_mask : Mask = face_occluder.run(None, - { - face_occluder.get_inputs()[0].name: prepare_vision_frame - })[0][0] + occlusion_mask = forward_occlude_face(prepare_vision_frame) occlusion_mask = occlusion_mask.transpose(0, 1, 2).clip(0, 1).astype(numpy.float32) occlusion_mask = cv2.resize(occlusion_mask, crop_vision_frame.shape[:2][::-1]) occlusion_mask = (cv2.GaussianBlur(occlusion_mask.clip(0, 1), (0, 0), 5).clip(0.5, 1) - 0.5) * 2 @@ -131,15 +126,14 @@ def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask: def create_region_mask(crop_vision_frame : VisionFrame, face_mask_regions : List[FaceMaskRegion]) -> Mask: - face_parser = get_face_parser() - prepare_vision_frame = cv2.flip(cv2.resize(crop_vision_frame, (512, 512)), 1) - prepare_vision_frame = numpy.expand_dims(prepare_vision_frame, axis = 0).astype(numpy.float32)[:, :, ::-1] / 127.5 - 1 + model_size = MODEL_SET.get('face_parser').get('size') + prepare_vision_frame = cv2.resize(crop_vision_frame, model_size) + prepare_vision_frame = prepare_vision_frame[:, :, ::-1].astype(numpy.float32) / 255 + 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) - with conditional_thread_semaphore(facefusion.globals.execution_providers): - region_mask : Mask = face_parser.run(None, - { - face_parser.get_inputs()[0].name: prepare_vision_frame - })[0][0] + region_mask = forward_parse_face(prepare_vision_frame) region_mask = numpy.isin(region_mask.argmax(0), [ FACE_MASK_REGIONS[region] for 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 @@ -149,7 +143,31 @@ def create_region_mask(crop_vision_frame : VisionFrame, face_mask_regions : List def create_mouth_mask(face_landmark_68 : FaceLandmark68) -> Mask: convex_hull = cv2.convexHull(face_landmark_68[numpy.r_[3:14, 31:36]].astype(numpy.int32)) mouth_mask : Mask = numpy.zeros((512, 512)).astype(numpy.float32) - mouth_mask = cv2.fillConvexPoly(mouth_mask, convex_hull, 1.0) + mouth_mask = cv2.fillConvexPoly(mouth_mask, convex_hull, 1.0) #type:ignore[call-overload] mouth_mask = cv2.erode(mouth_mask.clip(0, 1), numpy.ones((21, 3))) mouth_mask = cv2.GaussianBlur(mouth_mask, (0, 0), sigmaX = 1, sigmaY = 15) return mouth_mask + + +def forward_occlude_face(prepare_vision_frame : VisionFrame) -> Mask: + face_occluder = get_inference_pool().get('face_occluder') + + 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: + face_parser = get_inference_pool().get('face_parser') + + 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 00000000..5ed85c42 --- /dev/null +++ b/facefusion/face_recognizer.py @@ -0,0 +1,81 @@ +from typing import Tuple + +import numpy + +from facefusion import inference_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources +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.typing import Embedding, FaceLandmark5, InferencePool, ModelOptions, ModelSet, VisionFrame + +MODEL_SET : ModelSet =\ +{ + 'arcface': + { + 'hashes': + { + 'face_recognizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_w600k_r50.hash', + 'path': resolve_relative_path('../.assets/models/arcface_w600k_r50.hash') + } + }, + 'sources': + { + 'face_recognizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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_sources = get_model_options().get('sources') + return inference_manager.get_inference_pool(__name__, model_sources) + + +def clear_inference_pool() -> None: + inference_manager.clear_inference_pool(__name__) + + +def get_model_options() -> ModelOptions: + return MODEL_SET.get('arcface') + + +def pre_check() -> bool: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +def calc_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) + embedding = forward(crop_vision_frame) + embedding = embedding.ravel() + normed_embedding = embedding / numpy.linalg.norm(embedding) + return embedding, normed_embedding + + +def forward(crop_vision_frame : VisionFrame) -> Embedding: + face_recognizer = get_inference_pool().get('face_recognizer') + + with conditional_thread_semaphore(): + embedding = face_recognizer.run(None, + { + 'input': crop_vision_frame + })[0] + + return embedding diff --git a/facefusion/face_selector.py b/facefusion/face_selector.py new file mode 100644 index 00000000..a3385e96 --- /dev/null +++ b/facefusion/face_selector.py @@ -0,0 +1,91 @@ +from typing import List + +import numpy + +from facefusion import state_manager +from facefusion.typing import Face, FaceSelectorOrder, FaceSet, Gender, Race + + +def find_similar_faces(faces : List[Face], reference_faces : FaceSet, face_distance : float) -> List[Face]: + similar_faces : List[Face] = [] + + if faces and reference_faces: + for reference_set in reference_faces: + if not similar_faces: + for reference_face in reference_faces[reference_set]: + for face in faces: + if compare_faces(face, reference_face, face_distance): + similar_faces.append(face) + return similar_faces + + +def compare_faces(face : Face, reference_face : Face, face_distance : float) -> bool: + current_face_distance = calc_face_distance(face, reference_face) + return current_face_distance < face_distance + + +def calc_face_distance(face : Face, reference_face : Face) -> float: + if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'): + return 1 - numpy.dot(face.normed_embedding, reference_face.normed_embedding) + return 0 + + +def sort_and_filter_faces(faces : List[Face]) -> List[Face]: + if faces: + if state_manager.get_item('face_selector_order'): + faces = sort_by_order(faces, state_manager.get_item('face_selector_order')) + if state_manager.get_item('face_selector_gender'): + faces = filter_by_gender(faces, state_manager.get_item('face_selector_gender')) + if state_manager.get_item('face_selector_race'): + faces = filter_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_by_age(faces, state_manager.get_item('face_selector_age_start'), state_manager.get_item('face_selector_age_end')) + return faces + + +def sort_by_order(faces : List[Face], order : FaceSelectorOrder) -> List[Face]: + if order == 'left-right': + return sorted(faces, key = lambda face: face.bounding_box[0]) + if order == 'right-left': + return sorted(faces, key = lambda face: face.bounding_box[0], reverse = True) + if order == 'top-bottom': + return sorted(faces, key = lambda face: face.bounding_box[1]) + if order == 'bottom-top': + return sorted(faces, key = lambda face: face.bounding_box[1], reverse = True) + if order == 'small-large': + return sorted(faces, key = lambda face: (face.bounding_box[2] - face.bounding_box[0]) * (face.bounding_box[3] - face.bounding_box[1])) + if order == 'large-small': + return sorted(faces, key = lambda face: (face.bounding_box[2] - face.bounding_box[0]) * (face.bounding_box[3] - face.bounding_box[1]), reverse = True) + if order == 'best-worst': + return sorted(faces, key = lambda face: face.score_set.get('detector'), reverse = True) + if order == 'worst-best': + return sorted(faces, key = lambda face: face.score_set.get('detector')) + return faces + + +def filter_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_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_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 index 7540bc9d..7957c50f 100644 --- a/facefusion/face_store.py +++ b/facefusion/face_store.py @@ -1,16 +1,21 @@ -from typing import Optional, List import hashlib +from typing import List, Optional + import numpy -from facefusion.typing import VisionFrame, Face, FaceStore, FaceSet +from facefusion.typing import Face, FaceSet, FaceStore, VisionFrame -FACE_STORE: FaceStore =\ +FACE_STORE : FaceStore =\ { 'static_faces': {}, 'reference_faces': {} } +def get_face_store() -> FaceStore: + return FACE_STORE + + def get_static_faces(vision_frame : VisionFrame) -> Optional[List[Face]]: frame_hash = create_frame_hash(vision_frame) if frame_hash in FACE_STORE['static_faces']: diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py index f1e6724f..68107c51 100644 --- a/facefusion/ffmpeg.py +++ b/facefusion/ffmpeg.py @@ -1,32 +1,39 @@ -from typing import List, Optional import os +import shutil import subprocess +import tempfile +from typing import List, Optional + import filetype -import facefusion.globals -from facefusion import logger, process_manager -from facefusion.typing import OutputVideoPreset, Fps, AudioBuffer -from facefusion.filesystem import get_temp_frames_pattern, get_temp_file_path +from facefusion import logger, process_manager, state_manager +from facefusion.filesystem import remove_file +from facefusion.temp_helper import get_temp_file_path, get_temp_frames_pattern +from facefusion.typing import AudioBuffer, Fps, OutputVideoPreset from facefusion.vision import restrict_video_fps -def run_ffmpeg(args : List[str]) -> bool: - commands = [ 'ffmpeg', '-hide_banner', '-loglevel', 'error' ] +def run_ffmpeg(args : List[str]) -> subprocess.Popen[bytes]: + commands = [ shutil.which('ffmpeg'), '-hide_banner', '-loglevel', 'error' ] commands.extend(args) process = subprocess.Popen(commands, stderr = subprocess.PIPE, stdout = subprocess.PIPE) while process_manager.is_processing(): try: - if facefusion.globals.log_level == 'debug': + if state_manager.get_item('log_level') == 'debug': log_debug(process) - return process.wait(timeout = 0.5) == 0 + process.wait(timeout = 0.5) except subprocess.TimeoutExpired: continue - return process.returncode == 0 + return process + + if process_manager.is_stopping(): + process.terminate() + return process def open_ffmpeg(args : List[str]) -> subprocess.Popen[bytes]: - commands = [ 'ffmpeg', '-hide_banner', '-loglevel', 'quiet' ] + commands = [ shutil.which('ffmpeg'), '-hide_banner', '-loglevel', 'quiet' ] commands.extend(args) return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.PIPE) @@ -37,66 +44,89 @@ def log_debug(process : subprocess.Popen[bytes]) -> None: for error in errors: if error.strip(): - logger.debug(error.strip(), __name__.upper()) + logger.debug(error.strip(), __name__) def extract_frames(target_path : str, temp_video_resolution : str, temp_video_fps : Fps) -> bool: - trim_frame_start = facefusion.globals.trim_frame_start - trim_frame_end = facefusion.globals.trim_frame_end - temp_frames_pattern = get_temp_frames_pattern(target_path, '%04d') + trim_frame_start = state_manager.get_item('trim_frame_start') + trim_frame_end = state_manager.get_item('trim_frame_end') + temp_frames_pattern = get_temp_frames_pattern(target_path, '%08d') commands = [ '-i', target_path, '-s', str(temp_video_resolution), '-q:v', '0' ] - if trim_frame_start is not None and trim_frame_end is not None: + if isinstance(trim_frame_start, int) and isinstance(trim_frame_end, int): commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ':end_frame=' + str(trim_frame_end) + ',fps=' + str(temp_video_fps) ]) - elif trim_frame_start is not None: + elif isinstance(trim_frame_start, int): commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ',fps=' + str(temp_video_fps) ]) - elif trim_frame_end is not None: + elif isinstance(trim_frame_end, int): commands.extend([ '-vf', 'trim=end_frame=' + str(trim_frame_end) + ',fps=' + str(temp_video_fps) ]) else: commands.extend([ '-vf', 'fps=' + str(temp_video_fps) ]) commands.extend([ '-vsync', '0', temp_frames_pattern ]) - return run_ffmpeg(commands) + return run_ffmpeg(commands).returncode == 0 def merge_video(target_path : str, output_video_resolution : str, output_video_fps : Fps) -> bool: temp_video_fps = restrict_video_fps(target_path, output_video_fps) temp_file_path = get_temp_file_path(target_path) - temp_frames_pattern = get_temp_frames_pattern(target_path, '%04d') - commands = [ '-r', str(temp_video_fps), '-i', temp_frames_pattern, '-s', str(output_video_resolution), '-c:v', facefusion.globals.output_video_encoder ] + temp_frames_pattern = get_temp_frames_pattern(target_path, '%08d') + commands = [ '-r', str(temp_video_fps), '-i', temp_frames_pattern, '-s', str(output_video_resolution), '-c:v', state_manager.get_item('output_video_encoder') ] - if facefusion.globals.output_video_encoder in [ 'libx264', 'libx265' ]: - output_video_compression = round(51 - (facefusion.globals.output_video_quality * 0.51)) - commands.extend([ '-crf', str(output_video_compression), '-preset', facefusion.globals.output_video_preset ]) - if facefusion.globals.output_video_encoder in [ 'libvpx-vp9' ]: - output_video_compression = round(63 - (facefusion.globals.output_video_quality * 0.63)) + if state_manager.get_item('output_video_encoder') in [ 'libx264', 'libx265' ]: + output_video_compression = round(51 - (state_manager.get_item('output_video_quality') * 0.51)) + commands.extend([ '-crf', str(output_video_compression), '-preset', state_manager.get_item('output_video_preset') ]) + if state_manager.get_item('output_video_encoder') in [ 'libvpx-vp9' ]: + output_video_compression = round(63 - (state_manager.get_item('output_video_quality') * 0.63)) commands.extend([ '-crf', str(output_video_compression) ]) - if facefusion.globals.output_video_encoder in [ 'h264_nvenc', 'hevc_nvenc' ]: - output_video_compression = round(51 - (facefusion.globals.output_video_quality * 0.51)) - commands.extend([ '-cq', str(output_video_compression), '-preset', map_nvenc_preset(facefusion.globals.output_video_preset) ]) - if facefusion.globals.output_video_encoder in [ 'h264_amf', 'hevc_amf' ]: - output_video_compression = round(51 - (facefusion.globals.output_video_quality * 0.51)) - commands.extend([ '-qp_i', str(output_video_compression), '-qp_p', str(output_video_compression), '-quality', map_amf_preset(facefusion.globals.output_video_preset) ]) + if state_manager.get_item('output_video_encoder') in [ 'h264_nvenc', 'hevc_nvenc' ]: + output_video_compression = round(51 - (state_manager.get_item('output_video_quality') * 0.51)) + commands.extend([ '-cq', str(output_video_compression), '-preset', map_nvenc_preset(state_manager.get_item('output_video_preset')) ]) + if state_manager.get_item('output_video_encoder') in [ 'h264_amf', 'hevc_amf' ]: + output_video_compression = round(51 - (state_manager.get_item('output_video_quality') * 0.51)) + commands.extend([ '-qp_i', str(output_video_compression), '-qp_p', str(output_video_compression), '-quality', map_amf_preset(state_manager.get_item('output_video_preset')) ]) + if state_manager.get_item('output_video_encoder') in [ 'h264_videotoolbox', 'hevc_videotoolbox' ]: + commands.extend([ '-q:v', str(state_manager.get_item('output_video_quality')) ]) commands.extend([ '-vf', 'framerate=fps=' + str(output_video_fps), '-pix_fmt', 'yuv420p', '-colorspace', 'bt709', '-y', temp_file_path ]) - return run_ffmpeg(commands) + return run_ffmpeg(commands).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() + commands = [ '-f', 'concat', '-safe', '0', '-i', concat_video_file.name, '-c:v', 'copy', '-c:a', state_manager.get_item('output_audio_encoder'), '-y', os.path.abspath(output_path) ] + process = run_ffmpeg(commands) + process.communicate() + remove_file(concat_video_path) + return process.returncode == 0 def copy_image(target_path : str, temp_image_resolution : str) -> bool: temp_file_path = get_temp_file_path(target_path) - is_webp = filetype.guess_mime(target_path) == 'image/webp' - temp_image_compression = 100 if is_webp else 0 + temp_image_compression = calc_image_compression(target_path, 100) commands = [ '-i', target_path, '-s', str(temp_image_resolution), '-q:v', str(temp_image_compression), '-y', temp_file_path ] - return run_ffmpeg(commands) + return run_ffmpeg(commands).returncode == 0 def finalize_image(target_path : str, output_path : str, output_image_resolution : str) -> bool: temp_file_path = get_temp_file_path(target_path) - output_image_compression = round(31 - (facefusion.globals.output_image_quality * 0.31)) + output_image_compression = calc_image_compression(target_path, state_manager.get_item('output_image_quality')) commands = [ '-i', temp_file_path, '-s', str(output_image_resolution), '-q:v', str(output_image_compression), '-y', output_path ] - return run_ffmpeg(commands) + return run_ffmpeg(commands).returncode == 0 + + +def calc_image_compression(image_path : str, image_quality : int) -> int: + is_webp = filetype.guess_mime(image_path) == 'image/webp' + if is_webp: + image_quality = 100 - image_quality + return round(31 - (image_quality * 0.31)) def read_audio_buffer(target_path : str, sample_rate : int, channel_total : int) -> Optional[AudioBuffer]: - commands = [ '-i', target_path, '-vn', '-f', 's16le', '-acodec', 'pcm_s16le', '-ar', str(sample_rate), '-ac', str(channel_total), '-'] + commands = [ '-i', target_path, '-vn', '-f', 's16le', '-acodec', 'pcm_s16le', '-ar', str(sample_rate), '-ac', str(channel_total), '-' ] process = open_ffmpeg(commands) audio_buffer, _ = process.communicate() if process.returncode == 0: @@ -105,25 +135,25 @@ def read_audio_buffer(target_path : str, sample_rate : int, channel_total : int) def restore_audio(target_path : str, output_path : str, output_video_fps : Fps) -> bool: - trim_frame_start = facefusion.globals.trim_frame_start - trim_frame_end = facefusion.globals.trim_frame_end + trim_frame_start = state_manager.get_item('trim_frame_start') + trim_frame_end = state_manager.get_item('trim_frame_end') temp_file_path = get_temp_file_path(target_path) commands = [ '-i', temp_file_path ] - if trim_frame_start is not None: + if isinstance(trim_frame_start, int): start_time = trim_frame_start / output_video_fps commands.extend([ '-ss', str(start_time) ]) - if trim_frame_end is not None: + if isinstance(trim_frame_end, int): end_time = trim_frame_end / output_video_fps commands.extend([ '-to', str(end_time) ]) - commands.extend([ '-i', target_path, '-c', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-shortest', '-y', output_path ]) - return run_ffmpeg(commands) + commands.extend([ '-i', target_path, '-c:v', 'copy', '-c:a', state_manager.get_item('output_audio_encoder'), '-map', '0:v:0', '-map', '1:a:0', '-shortest', '-y', output_path ]) + return run_ffmpeg(commands).returncode == 0 def replace_audio(target_path : str, audio_path : str, output_path : str) -> bool: temp_file_path = get_temp_file_path(target_path) - commands = [ '-i', temp_file_path, '-i', audio_path, '-af', 'apad', '-shortest', '-y', output_path ] - return run_ffmpeg(commands) + commands = [ '-i', temp_file_path, '-i', audio_path, '-c:a', state_manager.get_item('output_audio_encoder'), '-af', 'apad', '-shortest', '-y', output_path ] + return run_ffmpeg(commands).returncode == 0 def map_nvenc_preset(output_video_preset : OutputVideoPreset) -> Optional[str]: diff --git a/facefusion/filesystem.py b/facefusion/filesystem.py index cacb7538..ac01c944 100644 --- a/facefusion/filesystem.py +++ b/facefusion/filesystem.py @@ -1,70 +1,34 @@ -from typing import List, Optional -import glob import os import shutil -import tempfile -import filetype from pathlib import Path +from typing import List, Optional + +import filetype -import facefusion.globals from facefusion.common_helper import is_windows if is_windows(): import ctypes -def get_temp_frame_paths(target_path : str) -> List[str]: - temp_frames_pattern = get_temp_frames_pattern(target_path, '*') - return sorted(glob.glob(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 + '.' + facefusion.globals.temp_frame_format) - - -def get_temp_file_path(target_path : str) -> str: - _, target_extension = os.path.splitext(os.path.basename(target_path)) - temp_directory_path = get_temp_directory_path(target_path) - return os.path.join(temp_directory_path, 'temp' + target_extension) - - -def get_temp_directory_path(target_path : str) -> str: - target_name, _ = os.path.splitext(os.path.basename(target_path)) - temp_directory_path = os.path.join(tempfile.gettempdir(), 'facefusion') - return os.path.join(temp_directory_path, target_name) - - -def create_temp(target_path : str) -> None: - temp_directory_path = get_temp_directory_path(target_path) - Path(temp_directory_path).mkdir(parents = True, exist_ok = True) - - -def move_temp(target_path : str, output_path : str) -> None: - temp_file_path = get_temp_file_path(target_path) - - if is_file(temp_file_path): - if is_file(output_path): - os.remove(output_path) - shutil.move(temp_file_path, output_path) - - -def clear_temp(target_path : str) -> None: - temp_directory_path = get_temp_directory_path(target_path) - parent_directory_path = os.path.dirname(temp_directory_path) - - if not facefusion.globals.keep_temp and is_directory(temp_directory_path): - shutil.rmtree(temp_directory_path, ignore_errors = True) - if os.path.exists(parent_directory_path) and not os.listdir(parent_directory_path): - os.rmdir(parent_directory_path) - - def get_file_size(file_path : str) -> int: if is_file(file_path): return os.path.getsize(file_path) return 0 +def same_file_extension(file_paths : List[str]) -> bool: + file_extensions : List[str] = [] + + for file_path in file_paths: + _, file_extension = os.path.splitext(file_path.lower()) + + if file_extensions and file_extension not in file_extensions: + return False + file_extensions.append(file_extension) + return True + + def is_file(file_path : str) -> bool: return bool(file_path and os.path.isfile(file_path)) @@ -73,6 +37,12 @@ def is_directory(directory_path : str) -> bool: return bool(directory_path and os.path.isdir(directory_path)) +def in_directory(file_path : str) -> bool: + if file_path and not is_directory(file_path): + return is_directory(os.path.dirname(file_path)) + return False + + def is_audio(audio_path : str) -> bool: return is_file(audio_path) and filetype.helpers.is_audio(audio_path) @@ -113,6 +83,48 @@ def resolve_relative_path(path : str) -> str: return os.path.abspath(os.path.join(os.path.dirname(__file__), path)) +def sanitize_path_for_windows(full_path : str) -> Optional[str]: + buffer_size = 0 + + while True: + unicode_buffer = ctypes.create_unicode_buffer(buffer_size) + buffer_limit = ctypes.windll.kernel32.GetShortPathNameW(full_path, unicode_buffer, buffer_size) #type:ignore[attr-defined] + + if buffer_size > buffer_limit: + return unicode_buffer.value + if buffer_limit == 0: + return None + buffer_size = buffer_limit + + +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 create_directory(directory_path : str) -> bool: + if directory_path and not is_file(directory_path): + Path(directory_path).mkdir(parents = True, exist_ok = True) + return is_directory(directory_path) + return False + + def list_directory(directory_path : str) -> Optional[List[str]]: if is_directory(directory_path): files = os.listdir(directory_path) @@ -121,15 +133,8 @@ def list_directory(directory_path : str) -> Optional[List[str]]: return None -def sanitize_path_for_windows(full_path : str) -> Optional[str]: - buffer_size = 0 - - while True: - unicode_buffer = ctypes.create_unicode_buffer(buffer_size) - buffer_threshold = ctypes.windll.kernel32.GetShortPathNameW(full_path, unicode_buffer, buffer_size) #type:ignore[attr-defined] - - if buffer_size > buffer_threshold: - return unicode_buffer.value - if buffer_threshold == 0: - return None - buffer_size = buffer_threshold +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 diff --git a/facefusion/globals.py b/facefusion/globals.py deleted file mode 100755 index 062cbf23..00000000 --- a/facefusion/globals.py +++ /dev/null @@ -1,60 +0,0 @@ -from typing import List, Optional - -from facefusion.typing import LogLevel, VideoMemoryStrategy, FaceSelectorMode, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, FaceMaskType, FaceMaskRegion, OutputVideoEncoder, OutputVideoPreset, FaceDetectorModel, FaceRecognizerModel, TempFrameFormat, Padding - -# general -config_path : Optional[str] = None -source_paths : Optional[List[str]] = None -target_path : Optional[str] = None -output_path : Optional[str] = None -# misc -force_download : Optional[bool] = None -skip_download : Optional[bool] = None -headless : Optional[bool] = None -log_level : Optional[LogLevel] = None -# execution -execution_device_id : Optional[str] = None -execution_providers : List[str] = [] -execution_thread_count : Optional[int] = None -execution_queue_count : Optional[int] = None -# memory -video_memory_strategy : Optional[VideoMemoryStrategy] = None -system_memory_limit : Optional[int] = None -# face analyser -face_analyser_order : Optional[FaceAnalyserOrder] = None -face_analyser_age : Optional[FaceAnalyserAge] = None -face_analyser_gender : Optional[FaceAnalyserGender] = None -face_detector_model : Optional[FaceDetectorModel] = None -face_detector_size : Optional[str] = None -face_detector_score : Optional[float] = None -face_landmarker_score : Optional[float] = None -face_recognizer_model : Optional[FaceRecognizerModel] = None -# face selector -face_selector_mode : Optional[FaceSelectorMode] = None -reference_face_position : Optional[int] = None -reference_face_distance : Optional[float] = None -reference_frame_number : Optional[int] = None -# face mask -face_mask_types : Optional[List[FaceMaskType]] = None -face_mask_blur : Optional[float] = None -face_mask_padding : Optional[Padding] = None -face_mask_regions : Optional[List[FaceMaskRegion]] = None -# frame extraction -trim_frame_start : Optional[int] = None -trim_frame_end : Optional[int] = None -temp_frame_format : Optional[TempFrameFormat] = None -keep_temp : Optional[bool] = None -# output creation -output_image_quality : Optional[int] = None -output_image_resolution : Optional[str] = None -output_video_encoder : Optional[OutputVideoEncoder] = None -output_video_preset : Optional[OutputVideoPreset] = None -output_video_quality : Optional[int] = None -output_video_resolution : Optional[str] = None -output_video_fps : Optional[float] = None -skip_audio : Optional[bool] = None -# frame processors -frame_processors : List[str] = [] -# uis -open_browser : Optional[bool] = None -ui_layouts : List[str] = [] diff --git a/facefusion/hash_helper.py b/facefusion/hash_helper.py new file mode 100644 index 00000000..9d334b97 --- /dev/null +++ b/facefusion/hash_helper.py @@ -0,0 +1,32 @@ +import os +import zlib +from typing import Optional + +from facefusion.filesystem import 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, 'r') as hash_file: + hash_content = hash_file.read().strip() + + 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, _ = os.path.split(validate_path) + validate_file_name, _ = os.path.splitext(_) + + 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 00000000..76bb4556 --- /dev/null +++ b/facefusion/inference_manager.py @@ -0,0 +1,75 @@ +from functools import lru_cache +from time import sleep +from typing import List + +import onnx +from onnxruntime import InferenceSession + +from facefusion import process_manager, state_manager +from facefusion.app_context import detect_app_context +from facefusion.execution import create_execution_providers, has_execution_provider +from facefusion.thread_helper import thread_lock +from facefusion.typing import DownloadSet, ExecutionProviderKey, InferencePool, InferencePoolSet, ModelInitializer + +INFERENCE_POOLS : InferencePoolSet =\ +{ + 'cli': {}, # type:ignore[typeddict-item] + 'ui': {} # type:ignore[typeddict-item] +} + + +def get_inference_pool(model_context : str, model_sources : DownloadSet) -> InferencePool: + global INFERENCE_POOLS + + with thread_lock(): + while process_manager.is_checking(): + sleep(0.5) + app_context = detect_app_context() + inference_context = get_inference_context(model_context) + + if not INFERENCE_POOLS.get(app_context).get(inference_context): + execution_provider_keys = resolve_execution_provider_keys(model_context) + INFERENCE_POOLS[app_context][inference_context] = create_inference_pool(model_sources, state_manager.get_item('execution_device_id'), execution_provider_keys) + + return INFERENCE_POOLS.get(app_context).get(inference_context) + + +def create_inference_pool(model_sources : DownloadSet, execution_device_id : str, execution_provider_keys : List[ExecutionProviderKey]) -> InferencePool: + inference_pool : InferencePool = {} + + for model_name in model_sources.keys(): + inference_pool[model_name] = create_inference_session(model_sources.get(model_name).get('path'), execution_device_id, execution_provider_keys) + return inference_pool + + +def clear_inference_pool(model_context : str) -> None: + global INFERENCE_POOLS + + app_context = detect_app_context() + inference_context = get_inference_context(model_context) + + if INFERENCE_POOLS.get(app_context).get(inference_context): + del INFERENCE_POOLS[app_context][inference_context] + + +def create_inference_session(model_path : str, execution_device_id : str, execution_provider_keys : List[ExecutionProviderKey]) -> InferenceSession: + execution_providers = create_execution_providers(execution_device_id, execution_provider_keys) + return InferenceSession(model_path, providers = execution_providers) + + +@lru_cache(maxsize = None) +def get_static_model_initializer(model_path : str) -> ModelInitializer: + model = onnx.load(model_path) + return onnx.numpy_helper.to_array(model.graph.initializer[-1]) + + +def resolve_execution_provider_keys(model_context : str) -> List[ExecutionProviderKey]: + if has_execution_provider('coreml') and (model_context.startswith('facefusion.processors.modules.age_modifier') or model_context.startswith('facefusion.processors.modules.frame_colorizer')): + return [ 'cpu' ] + return state_manager.get_item('execution_providers') + + +def get_inference_context(model_context : str) -> str: + execution_provider_keys = resolve_execution_provider_keys(model_context) + inference_context = model_context + '.' + '_'.join(execution_provider_keys) + return inference_context diff --git a/facefusion/installer.py b/facefusion/installer.py index 4735dd4c..66a51feb 100644 --- a/facefusion/installer.py +++ b/facefusion/installer.py @@ -1,36 +1,33 @@ -from typing import Dict, Tuple -import sys import os -import tempfile +import shutil +import signal import subprocess -import inquirer +import sys +import tempfile from argparse import ArgumentParser, HelpFormatter +from typing import Dict, Tuple from facefusion import metadata, wording from facefusion.common_helper import is_linux, is_macos, is_windows -if is_macos(): - os.environ['SYSTEM_VERSION_COMPAT'] = '0' - ONNXRUNTIMES : Dict[str, Tuple[str, str]] = {} if is_macos(): - ONNXRUNTIMES['default'] = ('onnxruntime', '1.17.3') + ONNXRUNTIMES['default'] = ('onnxruntime', '1.19.2') else: - ONNXRUNTIMES['default'] = ('onnxruntime', '1.17.3') - ONNXRUNTIMES['cuda-12.2'] = ('onnxruntime-gpu', '1.17.1') - ONNXRUNTIMES['cuda-11.8'] = ('onnxruntime-gpu', '1.17.1') - ONNXRUNTIMES['openvino'] = ('onnxruntime-openvino', '1.15.0') + ONNXRUNTIMES['default'] = ('onnxruntime', '1.19.2') + ONNXRUNTIMES['cuda'] = ('onnxruntime-gpu', '1.19.2') + ONNXRUNTIMES['openvino'] = ('onnxruntime-openvino', '1.18.0') if is_linux(): - ONNXRUNTIMES['rocm-5.4.2'] = ('onnxruntime-rocm', '1.16.3') - ONNXRUNTIMES['rocm-5.6'] = ('onnxruntime-rocm', '1.16.3') + ONNXRUNTIMES['rocm'] = ('onnxruntime-rocm', '1.18.0') if is_windows(): - ONNXRUNTIMES['directml'] = ('onnxruntime-directml', '1.17.3') + ONNXRUNTIMES['directml'] = ('onnxruntime-directml', '1.19.2') def cli() -> None: - program = ArgumentParser(formatter_class = lambda prog: HelpFormatter(prog, max_help_position = 200)) - program.add_argument('--onnxruntime', help = wording.get('help.install_dependency').format(dependency = 'onnxruntime'), choices = ONNXRUNTIMES.keys()) + signal.signal(signal.SIGINT, lambda signal_number, frame: sys.exit(0)) + program = ArgumentParser(formatter_class = lambda prog: HelpFormatter(prog, max_help_position = 50)) + program.add_argument('--onnxruntime', help = wording.get('help.install_dependency').format(dependency = 'onnxruntime'), choices = ONNXRUNTIMES.keys(), required = True) program.add_argument('--skip-conda', help = wording.get('help.skip_conda'), action = 'store_true') program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version') run(program) @@ -38,41 +35,59 @@ def cli() -> None: def run(program : ArgumentParser) -> None: args = program.parse_args() - python_id = 'cp' + str(sys.version_info.major) + str(sys.version_info.minor) + has_conda = 'CONDA_PREFIX' in os.environ + onnxruntime_name, onnxruntime_version = ONNXRUNTIMES.get(args.onnxruntime) - if not args.skip_conda and 'CONDA_PREFIX' not in os.environ: + if not args.skip_conda and not has_conda: sys.stdout.write(wording.get('conda_not_activated') + os.linesep) sys.exit(1) - if args.onnxruntime: - answers =\ - { - 'onnxruntime': args.onnxruntime - } - else: - answers = inquirer.prompt( - [ - inquirer.List('onnxruntime', message = wording.get('help.install_dependency').format(dependency = 'onnxruntime'), choices = list(ONNXRUNTIMES.keys())) - ]) - if answers: - onnxruntime = answers['onnxruntime'] - onnxruntime_name, onnxruntime_version = ONNXRUNTIMES[onnxruntime] - subprocess.call([ 'pip', 'install', '-r', 'requirements.txt', '--force-reinstall' ]) - if onnxruntime == 'rocm-5.4.2' or onnxruntime == 'rocm-5.6': - if python_id in [ 'cp39', 'cp310', 'cp311' ]: - rocm_version = onnxruntime.replace('-', '') - rocm_version = rocm_version.replace('.', '') - wheel_name = 'onnxruntime_training-' + onnxruntime_version + '+' + rocm_version + '-' + python_id + '-' + python_id + '-manylinux_2_17_x86_64.manylinux2014_x86_64.whl' - wheel_path = os.path.join(tempfile.gettempdir(), wheel_name) - wheel_url = 'https://download.onnxruntime.ai/' + wheel_name - subprocess.call([ 'curl', '--silent', '--location', '--continue-at', '-', '--output', wheel_path, wheel_url ]) - subprocess.call([ 'pip', 'uninstall', wheel_path, '-y', '-q' ]) - subprocess.call([ 'pip', 'install', wheel_path, '--force-reinstall' ]) - os.remove(wheel_path) - else: - subprocess.call([ 'pip', 'uninstall', 'onnxruntime', onnxruntime_name, '-y', '-q' ]) - if onnxruntime == 'cuda-12.2': - subprocess.call([ 'pip', 'install', onnxruntime_name + '==' + onnxruntime_version, '--extra-index-url', 'https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple', '--force-reinstall' ]) - else: - subprocess.call([ 'pip', 'install', onnxruntime_name + '==' + onnxruntime_version, '--force-reinstall' ]) - subprocess.call([ 'pip', 'install', 'numpy==1.26.4', '--force-reinstall' ]) + subprocess.call([ shutil.which('pip'), 'install', '-r', 'requirements.txt', '--force-reinstall' ]) + + if args.onnxruntime == 'rocm': + python_id = 'cp' + str(sys.version_info.major) + str(sys.version_info.minor) + + if python_id == 'cp310': + wheel_name = 'onnxruntime_rocm-' + onnxruntime_version +'-' + python_id + '-' + python_id + '-linux_x86_64.whl' + wheel_path = os.path.join(tempfile.gettempdir(), wheel_name) + wheel_url = 'https://repo.radeon.com/rocm/manylinux/rocm-rel-6.2/' + wheel_name + subprocess.call([ shutil.which('curl'), '--silent', '--location', '--continue-at', '-', '--output', wheel_path, wheel_url ]) + subprocess.call([ shutil.which('pip'), 'uninstall', 'onnxruntime', wheel_path, '-y', '-q' ]) + subprocess.call([ shutil.which('pip'), 'install', wheel_path, '--force-reinstall' ]) + os.remove(wheel_path) + else: + subprocess.call([ shutil.which('pip'), 'uninstall', 'onnxruntime', onnxruntime_name, '-y', '-q' ]) + subprocess.call([ shutil.which('pip'), 'install', onnxruntime_name + '==' + onnxruntime_version, '--force-reinstall' ]) + + 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 = [ 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 = [ 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) ]) + + if onnxruntime_version == '1.18.0': + subprocess.call([ shutil.which('pip'), 'install', 'numpy==1.26.4', '--force-reinstall' ]) diff --git a/facefusion/processors/frame/__init__.py b/facefusion/jobs/__init__.py similarity index 100% rename from facefusion/processors/frame/__init__.py rename to facefusion/jobs/__init__.py diff --git a/facefusion/jobs/job_helper.py b/facefusion/jobs/job_helper.py new file mode 100644 index 00000000..26f468ef --- /dev/null +++ b/facefusion/jobs/job_helper.py @@ -0,0 +1,15 @@ +import os +from datetime import datetime +from typing import Optional + + +def get_step_output_path(job_id : str, step_index : int, output_path : str) -> Optional[str]: + if output_path: + output_directory_path, _ = os.path.split(output_path) + output_file_name, output_file_extension = os.path.splitext(_) + 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 00000000..a7b6e841 --- /dev/null +++ b/facefusion/jobs/job_list.py @@ -0,0 +1,34 @@ +from datetime import datetime +from typing import Optional, Tuple + +from facefusion.date_helper import describe_time_ago +from facefusion.jobs import job_manager +from facefusion.typing import JobStatus, TableContents, TableHeaders + + +def compose_job_list(job_status : JobStatus) -> Tuple[TableHeaders, TableContents]: + jobs = job_manager.find_jobs(job_status) + job_headers : TableHeaders = [ 'job id', 'steps', 'date created', 'date updated', 'job status' ] + job_contents : TableContents = [] + + 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 00000000..2e396b32 --- /dev/null +++ b/facefusion/jobs/job_manager.py @@ -0,0 +1,263 @@ +import glob +import os +from copy import copy +from typing import List, Optional + +from facefusion.choices import job_statuses +from facefusion.date_helper import get_current_date_time +from facefusion.filesystem import create_directory, is_directory, is_file, move_file, remove_directory, remove_file +from facefusion.jobs.job_helper import get_step_output_path +from facefusion.json import read_json, write_json +from facefusion.temp_helper import create_base_directory +from facefusion.typing 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 job_statuses ] + + create_base_directory() + 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() -> bool: + drafted_job_ids = find_job_ids('drafted') + + if drafted_job_ids: + for job_id in drafted_job_ids: + if not submit_job(job_id): + return False + return True + return False + + +def delete_job(job_id : str) -> bool: + return delete_job_file(job_id) + + +def delete_jobs() -> bool: + job_ids = find_job_ids('drafted') + find_job_ids('queued') + find_job_ids('failed') + find_job_ids('completed') + + if job_ids: + for job_id in job_ids: + if not delete_job(job_id): + return False + return True + return False + + +def find_jobs(job_status : JobStatus) -> JobSet: + job_ids = find_job_ids(job_status) + jobs : JobSet = {} + + for job_id in job_ids: + jobs[job_id] = read_job_file(job_id) + return jobs + + +def find_job_ids(job_status : JobStatus) -> List[str]: + job_pattern = os.path.join(JOBS_PATH, job_status, '*.json') + job_files = glob.glob(job_pattern) + job_files.sort(key = os.path.getmtime) + job_ids = [] + + for job_file in job_files: + job_id, _ = os.path.splitext(os.path.basename(job_file)) + 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 job_statuses: + job_pattern = os.path.join(JOBS_PATH, job_status, job_file_name) + job_paths = glob.glob(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 00000000..e1adddb2 --- /dev/null +++ b/facefusion/jobs/job_runner.py @@ -0,0 +1,106 @@ +from facefusion.ffmpeg import concat_video +from facefusion.filesystem import is_image, is_video, move_file, remove_file +from facefusion.jobs import job_helper, job_manager +from facefusion.typing 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) -> bool: + queued_job_ids = job_manager.find_job_ids('queued') + + if queued_job_ids: + for job_id in queued_job_ids: + if not run_job(job_id, process_step): + return False + return True + 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) -> bool: + failed_job_ids = job_manager.find_job_ids('failed') + + if failed_job_ids: + for job_id in failed_job_ids: + if not retry_job(job_id, process_step): + return False + return True + 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 all(map(is_video, temp_output_paths)): + if not concat_video(output_path, temp_output_paths): + return False + if any(map(is_image, 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) + 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) + output_set.setdefault(output_path, []).append(step_output_path) + return output_set diff --git a/facefusion/jobs/job_store.py b/facefusion/jobs/job_store.py new file mode 100644 index 00000000..9d330d09 --- /dev/null +++ b/facefusion/jobs/job_store.py @@ -0,0 +1,27 @@ +from typing import List + +from facefusion.typing 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(step_keys : List[str]) -> None: + for step_key in step_keys: + JOB_STORE['job_keys'].append(step_key) + + +def register_step_keys(job_keys : List[str]) -> None: + for job_key in job_keys: + JOB_STORE['step_keys'].append(job_key) diff --git a/facefusion/json.py b/facefusion/json.py new file mode 100644 index 00000000..dcb182c0 --- /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.typing import Content + + +def read_json(json_path : str) -> Optional[Content]: + if is_file(json_path): + try: + with open(json_path, 'r') 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/logger.py b/facefusion/logger.py index e6e3ba08..1ea484e0 100644 --- a/facefusion/logger.py +++ b/facefusion/logger.py @@ -1,32 +1,74 @@ -from typing import Dict -from logging import basicConfig, getLogger, Logger, DEBUG, INFO, WARNING, ERROR +from logging import Logger, basicConfig, getLogger +from typing import Tuple -from facefusion.typing import LogLevel +from facefusion.choices import log_level_set +from facefusion.common_helper import get_first, get_last +from facefusion.typing import LogLevel, TableContents, TableHeaders def init(log_level : LogLevel) -> None: - basicConfig(format = None) - get_package_logger().setLevel(get_log_levels()[log_level]) + basicConfig(format = '%(message)s') + get_package_logger().setLevel(log_level_set.get(log_level)) def get_package_logger() -> Logger: return getLogger('facefusion') -def debug(message : str, scope : str) -> None: - get_package_logger().debug('[' + scope + '] ' + message) +def debug(message : str, module_name : str) -> None: + get_package_logger().debug(create_message(message, module_name)) -def info(message : str, scope : str) -> None: - get_package_logger().info('[' + scope + '] ' + message) +def info(message : str, module_name : str) -> None: + get_package_logger().info(create_message(message, module_name)) -def warn(message : str, scope : str) -> None: - get_package_logger().warning('[' + scope + '] ' + message) +def warn(message : str, module_name : str) -> None: + get_package_logger().warning(create_message(message, module_name)) -def error(message : str, scope : str) -> None: - get_package_logger().error('[' + scope + '] ' + message) +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: + scopes = module_name.split('.') + first_scope = get_first(scopes) + last_scope = get_last(scopes) + + if first_scope and last_scope: + return '[' + first_scope.upper() + '.' + last_scope.upper() + '] ' + message + return message + + +def table(headers : TableHeaders, contents : TableContents) -> None: + package_logger = get_package_logger() + table_column, table_separator = create_table_parts(headers, contents) + + package_logger.info(table_separator) + package_logger.info(table_column.format(*headers)) + package_logger.info(table_separator) + + for content in contents: + package_logger.info(table_column.format(*content)) + + package_logger.info(table_separator) + + +def create_table_parts(headers : TableHeaders, contents : TableContents) -> 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) + '-+' def enable() -> None: @@ -35,13 +77,3 @@ def enable() -> None: def disable() -> None: get_package_logger().disabled = True - - -def get_log_levels() -> Dict[LogLevel, int]: - return\ - { - 'error': ERROR, - 'warn': WARNING, - 'info': INFO, - 'debug': DEBUG - } diff --git a/facefusion/metadata.py b/facefusion/metadata.py index 99e20c39..78b8aa74 100644 --- a/facefusion/metadata.py +++ b/facefusion/metadata.py @@ -1,13 +1,17 @@ +from typing import Optional + METADATA =\ { 'name': 'FaceFusion', - 'description': 'Next generation face swapper and enhancer', - 'version': '2.6.1', + 'description': 'Industry leading face manipulation platform', + 'version': '3.0.0', 'license': 'MIT', 'author': 'Henry Ruhs', 'url': 'https://facefusion.io' } -def get(key : str) -> str: - return METADATA[key] +def get(key : str) -> Optional[str]: + if key in METADATA: + return METADATA.get(key) + return None diff --git a/facefusion/normalizer.py b/facefusion/normalizer.py index 07b8071b..560dc5ff 100644 --- a/facefusion/normalizer.py +++ b/facefusion/normalizer.py @@ -1,24 +1,6 @@ from typing import List, Optional -import hashlib -import os -import facefusion.globals -from facefusion.filesystem import is_directory -from facefusion.typing import Padding, Fps - - -def normalize_output_path(target_path : Optional[str], output_path : Optional[str]) -> Optional[str]: - if target_path and output_path: - target_name, target_extension = os.path.splitext(os.path.basename(target_path)) - if is_directory(output_path): - output_hash = hashlib.sha1(str(facefusion.globals.__dict__).encode('utf-8')).hexdigest()[:8] - output_name = target_name + '-' + output_hash - return os.path.join(output_path, output_name + target_extension) - output_name, output_extension = os.path.splitext(os.path.basename(output_path)) - output_directory_path = os.path.dirname(output_path) - if is_directory(output_directory_path) and output_extension: - return os.path.join(output_directory_path, output_name + target_extension) - return None +from facefusion.typing import Fps, Padding def normalize_padding(padding : Optional[List[int]]) -> Optional[Padding]: @@ -34,6 +16,6 @@ def normalize_padding(padding : Optional[List[int]]) -> Optional[Padding]: def normalize_fps(fps : Optional[float]) -> Optional[Fps]: - if fps is not None: + 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 index 3d5cce0f..6ba526ad 100644 --- a/facefusion/process_manager.py +++ b/facefusion/process_manager.py @@ -1,6 +1,6 @@ from typing import Generator, List -from facefusion.typing import QueuePayload, ProcessState +from facefusion.typing import ProcessState, QueuePayload PROCESS_STATE : ProcessState = 'pending' diff --git a/facefusion/processors/choices.py b/facefusion/processors/choices.py new file mode 100755 index 00000000..9ab399eb --- /dev/null +++ b/facefusion/processors/choices.py @@ -0,0 +1,46 @@ +from typing import List, Sequence + +from facefusion.common_helper import create_float_range, create_int_range +from facefusion.processors.typing import AgeModifierModel, ExpressionRestorerModel, FaceDebuggerItem, FaceEditorModel, FaceEnhancerModel, FaceSwapperSet, FrameColorizerModel, FrameEnhancerModel, LipSyncerModel + +age_modifier_models : List[AgeModifierModel] = [ 'styleganex_age' ] +expression_restorer_models : List[ExpressionRestorerModel] = [ 'live_portrait' ] +face_debugger_items : List[FaceDebuggerItem] = [ 'bounding-box', 'face-landmark-5', 'face-landmark-5/68', 'face-landmark-68', 'face-landmark-68/5', 'face-mask', 'face-detector-score', 'face-landmarker-score', 'age', 'gender', 'race' ] +face_editor_models : List[FaceEditorModel] = [ 'live_portrait' ] +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_swapper_set : FaceSwapperSet =\ +{ + 'blendswap_256': [ '256x256', '384x384', '512x512', '768x768', '1024x1024' ], + 'ghost_256_unet_1': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'ghost_256_unet_2': [ '256x256', '512x512', '768x768', '1024x1024' ], + 'ghost_256_unet_3': [ '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_512_unofficial': [ '512x512', '768x768', '1024x1024' ], + 'uniface_256': [ '256x256', '512x512', '768x768', '1024x1024' ] +} +frame_colorizer_models : List[FrameColorizerModel] = [ 'ddcolor', 'ddcolor_artistic', 'deoldify', 'deoldify_artistic', 'deoldify_stable' ] +frame_colorizer_sizes : List[str] = [ '192x192', '256x256', '384x384', '512x512' ] +frame_enhancer_models : List[FrameEnhancerModel] = [ 'clear_reality_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', 'span_kendata_x4', 'ultra_sharp_x4' ] +lip_syncer_models : List[LipSyncerModel] = [ 'wav2lip', 'wav2lip_gan' ] + +age_modifier_direction_range : Sequence[int] = create_int_range(-100, 100, 1) +expression_restorer_factor_range : Sequence[int] = create_int_range(0, 100, 1) +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) +face_enhancer_blend_range : Sequence[int] = create_int_range(0, 100, 1) +frame_colorizer_blend_range : Sequence[int] = create_int_range(0, 100, 1) +frame_enhancer_blend_range : Sequence[int] = create_int_range(0, 100, 1) diff --git a/facefusion/processors/core.py b/facefusion/processors/core.py new file mode 100644 index 00000000..7cff5ebf --- /dev/null +++ b/facefusion/processors/core.py @@ -0,0 +1,110 @@ +import importlib +import os +from concurrent.futures import ThreadPoolExecutor, as_completed +from queue import Queue +from types import ModuleType +from typing import Any, List + +from tqdm import tqdm + +from facefusion import logger, state_manager, wording +from facefusion.exit_helper import hard_exit +from facefusion.typing import ProcessFrames, QueuePayload + +PROCESSORS_METHODS =\ +[ + 'get_inference_pool', + 'clear_inference_pool', + 'register_args', + 'apply_args', + 'pre_check', + 'pre_process', + 'post_process', + 'get_reference_frame', + 'process_frame', + 'process_frames', + 'process_image', + 'process_video' +] + + +def load_processor_module(processor : str) -> Any: + try: + processor_module = importlib.import_module('facefusion.processors.modules.' + processor) + for method_name in PROCESSORS_METHODS: + if not hasattr(processor_module, method_name): + raise NotImplementedError + except ModuleNotFoundError as exception: + logger.error(wording.get('processor_not_loaded').format(processor = processor), __name__) + logger.debug(exception.msg, __name__) + hard_exit(1) + except NotImplementedError: + logger.error(wording.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 + + +def clear_processors_modules(processors : List[str]) -> None: + for processor in processors: + processor_module = load_processor_module(processor) + processor_module.clear_inference_pool() + + +def multi_process_frames(source_paths : List[str], temp_frame_paths : List[str], process_frames : ProcessFrames) -> None: + queue_payloads = create_queue_payloads(temp_frame_paths) + with tqdm(total = len(queue_payloads), desc = wording.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'), + 'execution_thread_count': state_manager.get_item('execution_thread_count'), + 'execution_queue_count': state_manager.get_item('execution_queue_count') + }) + with ThreadPoolExecutor(max_workers = state_manager.get_item('execution_thread_count')) as executor: + futures = [] + queue : Queue[QueuePayload] = create_queue(queue_payloads) + queue_per_future = max(len(queue_payloads) // state_manager.get_item('execution_thread_count') * state_manager.get_item('execution_queue_count'), 1) + + while not queue.empty(): + future = executor.submit(process_frames, source_paths, pick_queue(queue, queue_per_future), progress.update) + futures.append(future) + + for future_done in as_completed(futures): + future_done.result() + + +def create_queue(queue_payloads : List[QueuePayload]) -> Queue[QueuePayload]: + queue : Queue[QueuePayload] = Queue() + for queue_payload in queue_payloads: + queue.put(queue_payload) + return queue + + +def pick_queue(queue : Queue[QueuePayload], queue_per_future : int) -> List[QueuePayload]: + queues = [] + for _ in range(queue_per_future): + if not queue.empty(): + queues.append(queue.get()) + return queues + + +def create_queue_payloads(temp_frame_paths : List[str]) -> List[QueuePayload]: + queue_payloads = [] + temp_frame_paths = sorted(temp_frame_paths, key = os.path.basename) + + for frame_number, frame_path in enumerate(temp_frame_paths): + frame_payload : QueuePayload =\ + { + 'frame_number': frame_number, + 'frame_path': frame_path + } + queue_payloads.append(frame_payload) + return queue_payloads diff --git a/facefusion/processors/frame/choices.py b/facefusion/processors/frame/choices.py deleted file mode 100755 index 4764f3e8..00000000 --- a/facefusion/processors/frame/choices.py +++ /dev/null @@ -1,16 +0,0 @@ -from typing import List - -from facefusion.common_helper import create_int_range -from facefusion.processors.frame.typings import FaceDebuggerItem, FaceEnhancerModel, FaceSwapperModel, FrameColorizerModel, FrameEnhancerModel, LipSyncerModel - -face_debugger_items : List[FaceDebuggerItem] = [ 'bounding-box', 'face-landmark-5', 'face-landmark-5/68', 'face-landmark-68', 'face-landmark-68/5', 'face-mask', 'face-detector-score', 'face-landmarker-score', 'age', 'gender' ] -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_swapper_models : List[FaceSwapperModel] = [ 'blendswap_256', 'inswapper_128', 'inswapper_128_fp16', 'simswap_256', 'simswap_512_unofficial', 'uniface_256' ] -frame_colorizer_models : List[FrameColorizerModel] = [ 'ddcolor', 'ddcolor_artistic', 'deoldify', 'deoldify_artistic', 'deoldify_stable' ] -frame_colorizer_sizes : List[str] = [ '192x192', '256x256', '384x384', '512x512' ] -frame_enhancer_models : List[FrameEnhancerModel] = [ 'clear_reality_x4', 'lsdir_x4', 'nomos8k_sc_x4', 'real_esrgan_x2', 'real_esrgan_x2_fp16', 'real_esrgan_x4', 'real_esrgan_x4_fp16', 'real_hatgan_x4', 'span_kendata_x4', 'ultra_sharp_x4' ] -lip_syncer_models : List[LipSyncerModel] = [ 'wav2lip_gan' ] - -face_enhancer_blend_range : List[int] = create_int_range(0, 100, 1) -frame_colorizer_blend_range : List[int] = create_int_range(0, 100, 1) -frame_enhancer_blend_range : List[int] = create_int_range(0, 100, 1) diff --git a/facefusion/processors/frame/core.py b/facefusion/processors/frame/core.py deleted file mode 100644 index 8b3f12b6..00000000 --- a/facefusion/processors/frame/core.py +++ /dev/null @@ -1,116 +0,0 @@ -import os -import sys -import importlib -from concurrent.futures import ThreadPoolExecutor, as_completed -from queue import Queue -from types import ModuleType -from typing import Any, List -from tqdm import tqdm - -import facefusion.globals -from facefusion.typing import ProcessFrames, QueuePayload -from facefusion.execution import encode_execution_providers -from facefusion import logger, wording - -FRAME_PROCESSORS_MODULES : List[ModuleType] = [] -FRAME_PROCESSORS_METHODS =\ -[ - 'get_frame_processor', - 'clear_frame_processor', - 'get_options', - 'set_options', - 'register_args', - 'apply_args', - 'pre_check', - 'post_check', - 'pre_process', - 'post_process', - 'get_reference_frame', - 'process_frame', - 'process_frames', - 'process_image', - 'process_video' -] - - -def load_frame_processor_module(frame_processor : str) -> Any: - try: - frame_processor_module = importlib.import_module('facefusion.processors.frame.modules.' + frame_processor) - for method_name in FRAME_PROCESSORS_METHODS: - if not hasattr(frame_processor_module, method_name): - raise NotImplementedError - except ModuleNotFoundError as exception: - logger.error(wording.get('frame_processor_not_loaded').format(frame_processor = frame_processor), __name__.upper()) - logger.debug(exception.msg, __name__.upper()) - sys.exit(1) - except NotImplementedError: - logger.error(wording.get('frame_processor_not_implemented').format(frame_processor = frame_processor), __name__.upper()) - sys.exit(1) - return frame_processor_module - - -def get_frame_processors_modules(frame_processors : List[str]) -> List[ModuleType]: - global FRAME_PROCESSORS_MODULES - - if not FRAME_PROCESSORS_MODULES: - for frame_processor in frame_processors: - frame_processor_module = load_frame_processor_module(frame_processor) - FRAME_PROCESSORS_MODULES.append(frame_processor_module) - return FRAME_PROCESSORS_MODULES - - -def clear_frame_processors_modules() -> None: - global FRAME_PROCESSORS_MODULES - - for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - frame_processor_module.clear_frame_processor() - FRAME_PROCESSORS_MODULES = [] - - -def multi_process_frames(source_paths : List[str], temp_frame_paths : List[str], process_frames : ProcessFrames) -> None: - queue_payloads = create_queue_payloads(temp_frame_paths) - with tqdm(total = len(queue_payloads), desc = wording.get('processing'), unit = 'frame', ascii = ' =', disable = facefusion.globals.log_level in [ 'warn', 'error' ]) as progress: - progress.set_postfix( - { - 'execution_providers': encode_execution_providers(facefusion.globals.execution_providers), - 'execution_thread_count': facefusion.globals.execution_thread_count, - 'execution_queue_count': facefusion.globals.execution_queue_count - }) - with ThreadPoolExecutor(max_workers = facefusion.globals.execution_thread_count) as executor: - futures = [] - queue : Queue[QueuePayload] = create_queue(queue_payloads) - queue_per_future = max(len(queue_payloads) // facefusion.globals.execution_thread_count * facefusion.globals.execution_queue_count, 1) - while not queue.empty(): - future = executor.submit(process_frames, source_paths, pick_queue(queue, queue_per_future), progress.update) - futures.append(future) - for future_done in as_completed(futures): - future_done.result() - - -def create_queue(queue_payloads : List[QueuePayload]) -> Queue[QueuePayload]: - queue : Queue[QueuePayload] = Queue() - for queue_payload in queue_payloads: - queue.put(queue_payload) - return queue - - -def pick_queue(queue : Queue[QueuePayload], queue_per_future : int) -> List[QueuePayload]: - queues = [] - for _ in range(queue_per_future): - if not queue.empty(): - queues.append(queue.get()) - return queues - - -def create_queue_payloads(temp_frame_paths : List[str]) -> List[QueuePayload]: - queue_payloads = [] - temp_frame_paths = sorted(temp_frame_paths, key = os.path.basename) - - for frame_number, frame_path in enumerate(temp_frame_paths): - frame_payload : QueuePayload =\ - { - 'frame_number': frame_number, - 'frame_path': frame_path - } - queue_payloads.append(frame_payload) - return queue_payloads diff --git a/facefusion/processors/frame/globals.py b/facefusion/processors/frame/globals.py deleted file mode 100755 index 76ab2b2a..00000000 --- a/facefusion/processors/frame/globals.py +++ /dev/null @@ -1,14 +0,0 @@ -from typing import List, Optional - -from facefusion.processors.frame.typings import FaceDebuggerItem, FaceEnhancerModel, FaceSwapperModel, FrameColorizerModel, FrameEnhancerModel, LipSyncerModel - -face_debugger_items : Optional[List[FaceDebuggerItem]] = None -face_enhancer_model : Optional[FaceEnhancerModel] = None -face_enhancer_blend : Optional[int] = None -face_swapper_model : Optional[FaceSwapperModel] = None -frame_colorizer_model : Optional[FrameColorizerModel] = None -frame_colorizer_blend : Optional[int] = None -frame_colorizer_size : Optional[str] = None -frame_enhancer_model : Optional[FrameEnhancerModel] = None -frame_enhancer_blend : Optional[int] = None -lip_syncer_model : Optional[LipSyncerModel] = None diff --git a/facefusion/processors/frame/modules/face_debugger.py b/facefusion/processors/frame/modules/face_debugger.py deleted file mode 100755 index ded5c645..00000000 --- a/facefusion/processors/frame/modules/face_debugger.py +++ /dev/null @@ -1,192 +0,0 @@ -from typing import Any, List, Literal -from argparse import ArgumentParser -import cv2 -import numpy - -import facefusion.globals -import facefusion.processors.frame.core as frame_processors -from facefusion import config, process_manager, wording -from facefusion.face_analyser import get_one_face, get_many_faces, find_similar_faces, clear_face_analyser -from facefusion.face_masker import create_static_box_mask, create_occlusion_mask, create_region_mask, clear_face_occluder, clear_face_parser -from facefusion.face_helper import warp_face_by_face_landmark_5, categorize_age, categorize_gender -from facefusion.face_store import get_reference_faces -from facefusion.content_analyser import clear_content_analyser -from facefusion.typing import Face, VisionFrame, UpdateProgress, ProcessMode, QueuePayload -from facefusion.vision import read_image, read_static_image, write_image -from facefusion.processors.frame.typings import FaceDebuggerInputs -from facefusion.processors.frame import globals as frame_processors_globals, choices as frame_processors_choices - -NAME = __name__.upper() - - -def get_frame_processor() -> None: - pass - - -def clear_frame_processor() -> None: - pass - - -def get_options(key : Literal['model']) -> None: - pass - - -def set_options(key : Literal['model'], value : Any) -> None: - pass - - -def register_args(program : ArgumentParser) -> None: - program.add_argument('--face-debugger-items', help = wording.get('help.face_debugger_items').format(choices = ', '.join(frame_processors_choices.face_debugger_items)), default = config.get_str_list('frame_processors.face_debugger_items', 'face-landmark-5/68 face-mask'), choices = frame_processors_choices.face_debugger_items, nargs = '+', metavar = 'FACE_DEBUGGER_ITEMS') - - -def apply_args(program : ArgumentParser) -> None: - args = program.parse_args() - frame_processors_globals.face_debugger_items = args.face_debugger_items - - -def pre_check() -> bool: - return True - - -def post_check() -> bool: - return True - - -def pre_process(mode : ProcessMode) -> bool: - return True - - -def post_process() -> None: - read_static_image.cache_clear() - if facefusion.globals.video_memory_strategy == 'strict' or facefusion.globals.video_memory_strategy == 'moderate': - clear_frame_processor() - if facefusion.globals.video_memory_strategy == 'strict': - clear_face_analyser() - clear_content_analyser() - clear_face_occluder() - clear_face_parser() - - -def debug_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: - primary_color = (0, 0, 255) - secondary_color = (0, 255, 0) - tertiary_color = (255, 255, 0) - bounding_box = target_face.bounding_box.astype(numpy.int32) - temp_vision_frame = temp_vision_frame.copy() - has_face_landmark_5_fallback = numpy.array_equal(target_face.landmarks.get('5'), target_face.landmarks.get('5/68')) - has_face_landmark_68_fallback = numpy.array_equal(target_face.landmarks.get('68'), target_face.landmarks.get('68/5')) - - if 'bounding-box' in frame_processors_globals.face_debugger_items: - cv2.rectangle(temp_vision_frame, (bounding_box[0], bounding_box[1]), (bounding_box[2], bounding_box[3]), primary_color, 2) - if 'face-mask' in frame_processors_globals.face_debugger_items: - crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmarks.get('5/68'), 'arcface_128_v2', (512, 512)) - inverse_matrix = cv2.invertAffineTransform(affine_matrix) - temp_size = temp_vision_frame.shape[:2][::-1] - crop_mask_list = [] - if 'box' in facefusion.globals.face_mask_types: - box_mask = create_static_box_mask(crop_vision_frame.shape[:2][::-1], 0, facefusion.globals.face_mask_padding) - crop_mask_list.append(box_mask) - if 'occlusion' in facefusion.globals.face_mask_types: - occlusion_mask = create_occlusion_mask(crop_vision_frame) - crop_mask_list.append(occlusion_mask) - if 'region' in facefusion.globals.face_mask_types: - region_mask = create_region_mask(crop_vision_frame, facefusion.globals.face_mask_regions) - crop_mask_list.append(region_mask) - crop_mask = numpy.minimum.reduce(crop_mask_list).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_vision_frame[inverse_vision_frame > 0] = 255 - inverse_contours = cv2.findContours(inverse_vision_frame, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)[0] - cv2.drawContours(temp_vision_frame, inverse_contours, -1, tertiary_color if has_face_landmark_5_fallback else secondary_color, 2) - if 'face-landmark-5' in frame_processors_globals.face_debugger_items and numpy.any(target_face.landmarks.get('5')): - face_landmark_5 = target_face.landmarks.get('5').astype(numpy.int32) - for index in range(face_landmark_5.shape[0]): - cv2.circle(temp_vision_frame, (face_landmark_5[index][0], face_landmark_5[index][1]), 3, primary_color, -1) - if 'face-landmark-5/68' in frame_processors_globals.face_debugger_items and numpy.any(target_face.landmarks.get('5/68')): - face_landmark_5_68 = target_face.landmarks.get('5/68').astype(numpy.int32) - for index in range(face_landmark_5_68.shape[0]): - cv2.circle(temp_vision_frame, (face_landmark_5_68[index][0], face_landmark_5_68[index][1]), 3, tertiary_color if has_face_landmark_5_fallback else secondary_color, -1) - if 'face-landmark-68' in frame_processors_globals.face_debugger_items and numpy.any(target_face.landmarks.get('68')): - face_landmark_68 = target_face.landmarks.get('68').astype(numpy.int32) - for index in range(face_landmark_68.shape[0]): - cv2.circle(temp_vision_frame, (face_landmark_68[index][0], face_landmark_68[index][1]), 3, tertiary_color if has_face_landmark_68_fallback else secondary_color, -1) - if 'face-landmark-68/5' in frame_processors_globals.face_debugger_items and numpy.any(target_face.landmarks.get('68')): - face_landmark_68 = target_face.landmarks.get('68/5').astype(numpy.int32) - for index in range(face_landmark_68.shape[0]): - cv2.circle(temp_vision_frame, (face_landmark_68[index][0], face_landmark_68[index][1]), 3, primary_color, -1) - if bounding_box[3] - bounding_box[1] > 50 and bounding_box[2] - bounding_box[0] > 50: - top = bounding_box[1] - left = bounding_box[0] - 20 - if 'face-detector-score' in frame_processors_globals.face_debugger_items: - face_score_text = str(round(target_face.scores.get('detector'), 2)) - top = top + 20 - cv2.putText(temp_vision_frame, face_score_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, primary_color, 2) - if 'face-landmarker-score' in frame_processors_globals.face_debugger_items: - face_score_text = str(round(target_face.scores.get('landmarker'), 2)) - top = top + 20 - cv2.putText(temp_vision_frame, face_score_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, tertiary_color if has_face_landmark_5_fallback else secondary_color, 2) - if 'age' in frame_processors_globals.face_debugger_items: - face_age_text = categorize_age(target_face.age) - top = top + 20 - cv2.putText(temp_vision_frame, face_age_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, primary_color, 2) - if 'gender' in frame_processors_globals.face_debugger_items: - face_gender_text = categorize_gender(target_face.gender) - top = top + 20 - cv2.putText(temp_vision_frame, face_gender_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, primary_color, 2) - return temp_vision_frame - - -def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: - pass - - -def process_frame(inputs : FaceDebuggerInputs) -> VisionFrame: - reference_faces = inputs.get('reference_faces') - target_vision_frame = inputs.get('target_vision_frame') - - if facefusion.globals.face_selector_mode == 'many': - many_faces = get_many_faces(target_vision_frame) - if many_faces: - for target_face in many_faces: - target_vision_frame = debug_face(target_face, target_vision_frame) - if facefusion.globals.face_selector_mode == 'one': - target_face = get_one_face(target_vision_frame) - if target_face: - target_vision_frame = debug_face(target_face, target_vision_frame) - if facefusion.globals.face_selector_mode == 'reference': - similar_faces = find_similar_faces(reference_faces, target_vision_frame, facefusion.globals.reference_face_distance) - if similar_faces: - for similar_face in similar_faces: - target_vision_frame = debug_face(similar_face, target_vision_frame) - return target_vision_frame - - -def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - - for queue_payload in process_manager.manage(queue_payloads): - target_vision_path = queue_payload['frame_path'] - target_vision_frame = read_image(target_vision_path) - output_vision_frame = process_frame( - { - 'reference_faces': reference_faces, - 'target_vision_frame': target_vision_frame - }) - write_image(target_vision_path, output_vision_frame) - update_progress(1) - - -def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - target_vision_frame = read_static_image(target_path) - output_vision_frame = process_frame( - { - 'reference_faces': reference_faces, - 'target_vision_frame': target_vision_frame - }) - write_image(output_path, output_vision_frame) - - -def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: - frame_processors.multi_process_frames(source_paths, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/modules/face_enhancer.py b/facefusion/processors/frame/modules/face_enhancer.py deleted file mode 100755 index 2be9f971..00000000 --- a/facefusion/processors/frame/modules/face_enhancer.py +++ /dev/null @@ -1,301 +0,0 @@ -from typing import Any, List, Literal, Optional -from argparse import ArgumentParser -from time import sleep -import cv2 -import numpy -import onnxruntime - -import facefusion.globals -import facefusion.processors.frame.core as frame_processors -from facefusion import config, process_manager, logger, wording -from facefusion.face_analyser import get_many_faces, clear_face_analyser, find_similar_faces, get_one_face -from facefusion.face_masker import create_static_box_mask, create_occlusion_mask, clear_face_occluder -from facefusion.face_helper import warp_face_by_face_landmark_5, paste_back -from facefusion.execution import apply_execution_provider_options -from facefusion.content_analyser import clear_content_analyser -from facefusion.face_store import get_reference_faces -from facefusion.normalizer import normalize_output_path -from facefusion.thread_helper import thread_lock, thread_semaphore -from facefusion.typing import Face, VisionFrame, UpdateProgress, ProcessMode, ModelSet, OptionsWithModel, QueuePayload -from facefusion.common_helper import create_metavar -from facefusion.filesystem import is_file, is_image, is_video, resolve_relative_path -from facefusion.download import conditional_download, is_download_done -from facefusion.vision import read_image, read_static_image, write_image -from facefusion.processors.frame.typings import FaceEnhancerInputs -from facefusion.processors.frame import globals as frame_processors_globals -from facefusion.processors.frame import choices as frame_processors_choices - -FRAME_PROCESSOR = None -NAME = __name__.upper() -MODELS : ModelSet =\ -{ - 'codeformer': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/codeformer.onnx', - 'path': resolve_relative_path('../.assets/models/codeformer.onnx'), - 'template': 'ffhq_512', - 'size': (512, 512) - }, - 'gfpgan_1.2': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.2.onnx', - 'path': resolve_relative_path('../.assets/models/gfpgan_1.2.onnx'), - 'template': 'ffhq_512', - 'size': (512, 512) - }, - 'gfpgan_1.3': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.3.onnx', - 'path': resolve_relative_path('../.assets/models/gfpgan_1.3.onnx'), - 'template': 'ffhq_512', - 'size': (512, 512) - }, - 'gfpgan_1.4': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.4.onnx', - 'path': resolve_relative_path('../.assets/models/gfpgan_1.4.onnx'), - 'template': 'ffhq_512', - 'size': (512, 512) - }, - 'gpen_bfr_256': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_256.onnx', - 'path': resolve_relative_path('../.assets/models/gpen_bfr_256.onnx'), - 'template': 'arcface_128_v2', - 'size': (256, 256) - }, - 'gpen_bfr_512': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_512.onnx', - 'path': resolve_relative_path('../.assets/models/gpen_bfr_512.onnx'), - 'template': 'ffhq_512', - 'size': (512, 512) - }, - 'gpen_bfr_1024': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_1024.onnx', - 'path': resolve_relative_path('../.assets/models/gpen_bfr_1024.onnx'), - 'template': 'ffhq_512', - 'size': (1024, 1024) - }, - 'gpen_bfr_2048': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_2048.onnx', - 'path': resolve_relative_path('../.assets/models/gpen_bfr_2048.onnx'), - 'template': 'ffhq_512', - 'size': (2048, 2048) - }, - 'restoreformer_plus_plus': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/restoreformer_plus_plus.onnx', - 'path': resolve_relative_path('../.assets/models/restoreformer_plus_plus.onnx'), - 'template': 'ffhq_512', - 'size': (512, 512) - } -} -OPTIONS : Optional[OptionsWithModel] = None - - -def get_frame_processor() -> Any: - global FRAME_PROCESSOR - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if FRAME_PROCESSOR is None: - model_path = get_options('model').get('path') - FRAME_PROCESSOR = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return FRAME_PROCESSOR - - -def clear_frame_processor() -> None: - global FRAME_PROCESSOR - - FRAME_PROCESSOR = None - - -def get_options(key : Literal['model']) -> Any: - global OPTIONS - - if OPTIONS is None: - OPTIONS =\ - { - 'model': MODELS[frame_processors_globals.face_enhancer_model] - } - return OPTIONS.get(key) - - -def set_options(key : Literal['model'], value : Any) -> None: - global OPTIONS - - OPTIONS[key] = value - - -def register_args(program : ArgumentParser) -> None: - program.add_argument('--face-enhancer-model', help = wording.get('help.face_enhancer_model'), default = config.get_str_value('frame_processors.face_enhancer_model', 'gfpgan_1.4'), choices = frame_processors_choices.face_enhancer_models) - program.add_argument('--face-enhancer-blend', help = wording.get('help.face_enhancer_blend'), type = int, default = config.get_int_value('frame_processors.face_enhancer_blend', '80'), choices = frame_processors_choices.face_enhancer_blend_range, metavar = create_metavar(frame_processors_choices.face_enhancer_blend_range)) - - -def apply_args(program : ArgumentParser) -> None: - args = program.parse_args() - frame_processors_globals.face_enhancer_model = args.face_enhancer_model - frame_processors_globals.face_enhancer_blend = args.face_enhancer_blend - - -def pre_check() -> bool: - download_directory_path = resolve_relative_path('../.assets/models') - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, [ model_url ]) - process_manager.end() - return is_file(model_path) - - -def post_check() -> bool: - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download and not is_download_done(model_url, model_path): - logger.error(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) - return False - if not is_file(model_path): - logger.error(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def pre_process(mode : ProcessMode) -> bool: - if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path): - logger.error(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) - return False - if mode == 'output' and not normalize_output_path(facefusion.globals.target_path, facefusion.globals.output_path): - logger.error(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def post_process() -> None: - read_static_image.cache_clear() - if facefusion.globals.video_memory_strategy == 'strict' or facefusion.globals.video_memory_strategy == 'moderate': - clear_frame_processor() - if facefusion.globals.video_memory_strategy == 'strict': - clear_face_analyser() - clear_content_analyser() - clear_face_occluder() - - -def enhance_face(target_face: Face, temp_vision_frame : VisionFrame) -> VisionFrame: - model_template = get_options('model').get('template') - model_size = get_options('model').get('size') - crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmarks.get('5/68'), model_template, model_size) - box_mask = create_static_box_mask(crop_vision_frame.shape[:2][::-1], facefusion.globals.face_mask_blur, (0, 0, 0, 0)) - crop_mask_list =\ - [ - box_mask - ] - - if 'occlusion' in facefusion.globals.face_mask_types: - occlusion_mask = create_occlusion_mask(crop_vision_frame) - crop_mask_list.append(occlusion_mask) - crop_vision_frame = prepare_crop_frame(crop_vision_frame) - crop_vision_frame = apply_enhance(crop_vision_frame) - crop_vision_frame = normalize_crop_frame(crop_vision_frame) - crop_mask = numpy.minimum.reduce(crop_mask_list).clip(0, 1) - paste_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, crop_mask, affine_matrix) - temp_vision_frame = blend_frame(temp_vision_frame, paste_vision_frame) - return temp_vision_frame - - -def apply_enhance(crop_vision_frame : VisionFrame) -> VisionFrame: - frame_processor = get_frame_processor() - frame_processor_inputs = {} - - for frame_processor_input in frame_processor.get_inputs(): - if frame_processor_input.name == 'input': - frame_processor_inputs[frame_processor_input.name] = crop_vision_frame - if frame_processor_input.name == 'weight': - weight = numpy.array([ 1 ]).astype(numpy.double) - frame_processor_inputs[frame_processor_input.name] = weight - with thread_semaphore(): - crop_vision_frame = frame_processor.run(None, frame_processor_inputs)[0][0] - return crop_vision_frame - - -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_frame(temp_vision_frame : VisionFrame, paste_vision_frame : VisionFrame) -> VisionFrame: - face_enhancer_blend = 1 - (frame_processors_globals.face_enhancer_blend / 100) - temp_vision_frame = cv2.addWeighted(temp_vision_frame, face_enhancer_blend, paste_vision_frame, 1 - face_enhancer_blend, 0) - return temp_vision_frame - - -def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: - return enhance_face(target_face, temp_vision_frame) - - -def process_frame(inputs : FaceEnhancerInputs) -> VisionFrame: - reference_faces = inputs.get('reference_faces') - target_vision_frame = inputs.get('target_vision_frame') - - if facefusion.globals.face_selector_mode == 'many': - many_faces = get_many_faces(target_vision_frame) - if many_faces: - for target_face in many_faces: - target_vision_frame = enhance_face(target_face, target_vision_frame) - if facefusion.globals.face_selector_mode == 'one': - target_face = get_one_face(target_vision_frame) - if target_face: - target_vision_frame = enhance_face(target_face, target_vision_frame) - if facefusion.globals.face_selector_mode == 'reference': - similar_faces = find_similar_faces(reference_faces, target_vision_frame, facefusion.globals.reference_face_distance) - if similar_faces: - for similar_face in similar_faces: - target_vision_frame = enhance_face(similar_face, target_vision_frame) - return target_vision_frame - - -def process_frames(source_path : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - - for queue_payload in process_manager.manage(queue_payloads): - target_vision_path = queue_payload['frame_path'] - target_vision_frame = read_image(target_vision_path) - output_vision_frame = process_frame( - { - 'reference_faces': reference_faces, - 'target_vision_frame': target_vision_frame - }) - write_image(target_vision_path, output_vision_frame) - update_progress(1) - - -def process_image(source_path : str, target_path : str, output_path : str) -> None: - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - target_vision_frame = read_static_image(target_path) - output_vision_frame = process_frame( - { - 'reference_faces': reference_faces, - 'target_vision_frame': target_vision_frame - }) - write_image(output_path, output_vision_frame) - - -def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: - frame_processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/modules/face_swapper.py b/facefusion/processors/frame/modules/face_swapper.py deleted file mode 100755 index 047fa1f7..00000000 --- a/facefusion/processors/frame/modules/face_swapper.py +++ /dev/null @@ -1,369 +0,0 @@ -from typing import Any, List, Literal, Optional -from argparse import ArgumentParser -from time import sleep -import numpy -import onnx -import onnxruntime -from onnx import numpy_helper - -import facefusion.globals -import facefusion.processors.frame.core as frame_processors -from facefusion import config, process_manager, logger, wording -from facefusion.execution import has_execution_provider, apply_execution_provider_options -from facefusion.face_analyser import get_one_face, get_average_face, get_many_faces, find_similar_faces, clear_face_analyser -from facefusion.face_masker import create_static_box_mask, create_occlusion_mask, create_region_mask, clear_face_occluder, clear_face_parser -from facefusion.face_helper import warp_face_by_face_landmark_5, paste_back -from facefusion.face_store import get_reference_faces -from facefusion.content_analyser import clear_content_analyser -from facefusion.normalizer import normalize_output_path -from facefusion.thread_helper import thread_lock, conditional_thread_semaphore -from facefusion.typing import Face, Embedding, VisionFrame, UpdateProgress, ProcessMode, ModelSet, OptionsWithModel, QueuePayload -from facefusion.filesystem import is_file, is_image, has_image, is_video, filter_image_paths, resolve_relative_path -from facefusion.download import conditional_download, is_download_done -from facefusion.vision import read_image, read_static_image, read_static_images, write_image -from facefusion.processors.frame.typings import FaceSwapperInputs -from facefusion.processors.frame import globals as frame_processors_globals -from facefusion.processors.frame import choices as frame_processors_choices - -FRAME_PROCESSOR = None -MODEL_INITIALIZER = None -NAME = __name__.upper() -MODELS : ModelSet =\ -{ - 'blendswap_256': - { - 'type': 'blendswap', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/blendswap_256.onnx', - 'path': resolve_relative_path('../.assets/models/blendswap_256.onnx'), - 'template': 'ffhq_512', - 'size': (256, 256), - 'mean': [ 0.0, 0.0, 0.0 ], - 'standard_deviation': [ 1.0, 1.0, 1.0 ] - }, - 'inswapper_128': - { - 'type': 'inswapper', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx', - 'path': resolve_relative_path('../.assets/models/inswapper_128.onnx'), - 'template': 'arcface_128_v2', - 'size': (128, 128), - 'mean': [ 0.0, 0.0, 0.0 ], - 'standard_deviation': [ 1.0, 1.0, 1.0 ] - }, - 'inswapper_128_fp16': - { - 'type': 'inswapper', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128_fp16.onnx', - 'path': resolve_relative_path('../.assets/models/inswapper_128_fp16.onnx'), - 'template': 'arcface_128_v2', - 'size': (128, 128), - 'mean': [ 0.0, 0.0, 0.0 ], - 'standard_deviation': [ 1.0, 1.0, 1.0 ] - }, - 'simswap_256': - { - 'type': 'simswap', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/simswap_256.onnx', - 'path': resolve_relative_path('../.assets/models/simswap_256.onnx'), - 'template': 'arcface_112_v1', - 'size': (256, 256), - 'mean': [ 0.485, 0.456, 0.406 ], - 'standard_deviation': [ 0.229, 0.224, 0.225 ] - }, - 'simswap_512_unofficial': - { - 'type': 'simswap', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/simswap_512_unofficial.onnx', - 'path': resolve_relative_path('../.assets/models/simswap_512_unofficial.onnx'), - 'template': 'arcface_112_v1', - 'size': (512, 512), - 'mean': [ 0.0, 0.0, 0.0 ], - 'standard_deviation': [ 1.0, 1.0, 1.0 ] - }, - 'uniface_256': - { - 'type': 'uniface', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/uniface_256.onnx', - 'path': resolve_relative_path('../.assets/models/uniface_256.onnx'), - 'template': 'ffhq_512', - 'size': (256, 256), - 'mean': [ 0.0, 0.0, 0.0 ], - 'standard_deviation': [ 1.0, 1.0, 1.0 ] - } -} -OPTIONS : Optional[OptionsWithModel] = None - - -def get_frame_processor() -> Any: - global FRAME_PROCESSOR - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if FRAME_PROCESSOR is None: - model_path = get_options('model').get('path') - FRAME_PROCESSOR = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return FRAME_PROCESSOR - - -def clear_frame_processor() -> None: - global FRAME_PROCESSOR - - FRAME_PROCESSOR = None - - -def get_model_initializer() -> Any: - global MODEL_INITIALIZER - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if MODEL_INITIALIZER is None: - model_path = get_options('model').get('path') - model = onnx.load(model_path) - MODEL_INITIALIZER = numpy_helper.to_array(model.graph.initializer[-1]) - return MODEL_INITIALIZER - - -def clear_model_initializer() -> None: - global MODEL_INITIALIZER - - MODEL_INITIALIZER = None - - -def get_options(key : Literal['model']) -> Any: - global OPTIONS - - if OPTIONS is None: - OPTIONS =\ - { - 'model': MODELS[frame_processors_globals.face_swapper_model] - } - return OPTIONS.get(key) - - -def set_options(key : Literal['model'], value : Any) -> None: - global OPTIONS - - OPTIONS[key] = value - - -def register_args(program : ArgumentParser) -> None: - if has_execution_provider('CoreMLExecutionProvider') or has_execution_provider('OpenVINOExecutionProvider'): - face_swapper_model_fallback = 'inswapper_128' - else: - face_swapper_model_fallback = 'inswapper_128_fp16' - program.add_argument('--face-swapper-model', help = wording.get('help.face_swapper_model'), default = config.get_str_value('frame_processors.face_swapper_model', face_swapper_model_fallback), choices = frame_processors_choices.face_swapper_models) - - -def apply_args(program : ArgumentParser) -> None: - args = program.parse_args() - frame_processors_globals.face_swapper_model = args.face_swapper_model - if args.face_swapper_model == 'blendswap_256': - facefusion.globals.face_recognizer_model = 'arcface_blendswap' - if args.face_swapper_model == 'inswapper_128' or args.face_swapper_model == 'inswapper_128_fp16': - facefusion.globals.face_recognizer_model = 'arcface_inswapper' - if args.face_swapper_model == 'simswap_256' or args.face_swapper_model == 'simswap_512_unofficial': - facefusion.globals.face_recognizer_model = 'arcface_simswap' - if args.face_swapper_model == 'uniface_256': - facefusion.globals.face_recognizer_model = 'arcface_uniface' - - -def pre_check() -> bool: - download_directory_path = resolve_relative_path('../.assets/models') - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, [ model_url ]) - process_manager.end() - return is_file(model_path) - - -def post_check() -> bool: - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download and not is_download_done(model_url, model_path): - logger.error(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) - return False - if not is_file(model_path): - logger.error(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def pre_process(mode : ProcessMode) -> bool: - if not has_image(facefusion.globals.source_paths): - logger.error(wording.get('select_image_source') + wording.get('exclamation_mark'), NAME) - return False - source_image_paths = filter_image_paths(facefusion.globals.source_paths) - source_frames = read_static_images(source_image_paths) - for source_frame in source_frames: - if not get_one_face(source_frame): - logger.error(wording.get('no_source_face_detected') + wording.get('exclamation_mark'), NAME) - return False - if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path): - logger.error(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) - return False - if mode == 'output' and not normalize_output_path(facefusion.globals.target_path, facefusion.globals.output_path): - logger.error(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def post_process() -> None: - read_static_image.cache_clear() - if facefusion.globals.video_memory_strategy == 'strict' or facefusion.globals.video_memory_strategy == 'moderate': - clear_model_initializer() - clear_frame_processor() - if facefusion.globals.video_memory_strategy == 'strict': - clear_face_analyser() - clear_content_analyser() - clear_face_occluder() - clear_face_parser() - - -def swap_face(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: - model_template = get_options('model').get('template') - model_size = get_options('model').get('size') - crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmarks.get('5/68'), model_template, model_size) - crop_mask_list = [] - - if 'box' in facefusion.globals.face_mask_types: - box_mask = create_static_box_mask(crop_vision_frame.shape[:2][::-1], facefusion.globals.face_mask_blur, facefusion.globals.face_mask_padding) - crop_mask_list.append(box_mask) - if 'occlusion' in facefusion.globals.face_mask_types: - occlusion_mask = create_occlusion_mask(crop_vision_frame) - crop_mask_list.append(occlusion_mask) - crop_vision_frame = prepare_crop_frame(crop_vision_frame) - crop_vision_frame = apply_swap(source_face, crop_vision_frame) - crop_vision_frame = normalize_crop_frame(crop_vision_frame) - if 'region' in facefusion.globals.face_mask_types: - region_mask = create_region_mask(crop_vision_frame, facefusion.globals.face_mask_regions) - crop_mask_list.append(region_mask) - crop_mask = numpy.minimum.reduce(crop_mask_list).clip(0, 1) - temp_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, crop_mask, affine_matrix) - return temp_vision_frame - - -def apply_swap(source_face : Face, crop_vision_frame : VisionFrame) -> VisionFrame: - frame_processor = get_frame_processor() - model_type = get_options('model').get('type') - frame_processor_inputs = {} - - for frame_processor_input in frame_processor.get_inputs(): - if frame_processor_input.name == 'source': - if model_type == 'blendswap' or model_type == 'uniface': - frame_processor_inputs[frame_processor_input.name] = prepare_source_frame(source_face) - else: - frame_processor_inputs[frame_processor_input.name] = prepare_source_embedding(source_face) - if frame_processor_input.name == 'target': - frame_processor_inputs[frame_processor_input.name] = crop_vision_frame - with conditional_thread_semaphore(facefusion.globals.execution_providers): - crop_vision_frame = frame_processor.run(None, frame_processor_inputs)[0][0] - return crop_vision_frame - - -def prepare_source_frame(source_face : Face) -> VisionFrame: - model_type = get_options('model').get('type') - source_vision_frame = read_static_image(facefusion.globals.source_paths[0]) - if model_type == 'blendswap': - source_vision_frame, _ = warp_face_by_face_landmark_5(source_vision_frame, source_face.landmarks.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.landmarks.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_options('model').get('type') - if model_type == 'inswapper': - model_initializer = get_model_initializer() - source_embedding = source_face.embedding.reshape((1, -1)) - source_embedding = numpy.dot(source_embedding, model_initializer) / numpy.linalg.norm(source_embedding) - else: - source_embedding = source_face.normed_embedding.reshape(1, -1) - return source_embedding - - -def prepare_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: - model_mean = get_options('model').get('mean') - model_standard_deviation = get_options('model').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: - 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[:, :, ::-1] - return crop_vision_frame - - -def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: - return swap_face(source_face, target_face, temp_vision_frame) - - -def process_frame(inputs : FaceSwapperInputs) -> VisionFrame: - reference_faces = inputs.get('reference_faces') - source_face = inputs.get('source_face') - target_vision_frame = inputs.get('target_vision_frame') - - if facefusion.globals.face_selector_mode == 'many': - many_faces = get_many_faces(target_vision_frame) - if many_faces: - for target_face in many_faces: - target_vision_frame = swap_face(source_face, target_face, target_vision_frame) - if facefusion.globals.face_selector_mode == 'one': - target_face = get_one_face(target_vision_frame) - if target_face: - target_vision_frame = swap_face(source_face, target_face, target_vision_frame) - if facefusion.globals.face_selector_mode == 'reference': - similar_faces = find_similar_faces(reference_faces, target_vision_frame, facefusion.globals.reference_face_distance) - if similar_faces: - for similar_face in similar_faces: - target_vision_frame = swap_face(source_face, similar_face, target_vision_frame) - return target_vision_frame - - -def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - source_frames = read_static_images(source_paths) - source_face = get_average_face(source_frames) - - for queue_payload in process_manager.manage(queue_payloads): - target_vision_path = queue_payload['frame_path'] - target_vision_frame = read_image(target_vision_path) - output_vision_frame = process_frame( - { - 'reference_faces': reference_faces, - 'source_face': source_face, - 'target_vision_frame': target_vision_frame - }) - write_image(target_vision_path, output_vision_frame) - update_progress(1) - - -def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - source_frames = read_static_images(source_paths) - source_face = get_average_face(source_frames) - target_vision_frame = read_static_image(target_path) - output_vision_frame = process_frame( - { - 'reference_faces': reference_faces, - 'source_face': source_face, - 'target_vision_frame': target_vision_frame - }) - write_image(output_path, output_vision_frame) - - -def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: - frame_processors.multi_process_frames(source_paths, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/modules/frame_colorizer.py b/facefusion/processors/frame/modules/frame_colorizer.py deleted file mode 100644 index 6125815a..00000000 --- a/facefusion/processors/frame/modules/frame_colorizer.py +++ /dev/null @@ -1,241 +0,0 @@ -from typing import Any, List, Literal, Optional -from argparse import ArgumentParser -from time import sleep -import cv2 -import numpy -import onnxruntime - -import facefusion.globals -import facefusion.processors.frame.core as frame_processors -from facefusion import config, process_manager, logger, wording -from facefusion.face_analyser import clear_face_analyser -from facefusion.content_analyser import clear_content_analyser -from facefusion.execution import apply_execution_provider_options -from facefusion.normalizer import normalize_output_path -from facefusion.thread_helper import thread_lock, thread_semaphore -from facefusion.typing import Face, VisionFrame, UpdateProgress, ProcessMode, ModelSet, OptionsWithModel, QueuePayload -from facefusion.common_helper import create_metavar -from facefusion.filesystem import is_file, resolve_relative_path, is_image, is_video -from facefusion.download import conditional_download, is_download_done -from facefusion.vision import read_image, read_static_image, write_image, unpack_resolution -from facefusion.processors.frame.typings import FrameColorizerInputs -from facefusion.processors.frame import globals as frame_processors_globals -from facefusion.processors.frame import choices as frame_processors_choices - -FRAME_PROCESSOR = None -NAME = __name__.upper() -MODELS : ModelSet =\ -{ - 'ddcolor': - { - 'type': 'ddcolor', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/ddcolor.onnx', - 'path': resolve_relative_path('../.assets/models/ddcolor.onnx') - }, - 'ddcolor_artistic': - { - 'type': 'ddcolor', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/ddcolor_artistic.onnx', - 'path': resolve_relative_path('../.assets/models/ddcolor_artistic.onnx') - }, - 'deoldify': - { - 'type': 'deoldify', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/deoldify.onnx', - 'path': resolve_relative_path('../.assets/models/deoldify.onnx') - }, - 'deoldify_artistic': - { - 'type': 'deoldify', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/deoldify_artistic.onnx', - 'path': resolve_relative_path('../.assets/models/deoldify_artistic.onnx') - }, - 'deoldify_stable': - { - 'type': 'deoldify', - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/deoldify_stable.onnx', - 'path': resolve_relative_path('../.assets/models/deoldify_stable.onnx') - } -} -OPTIONS : Optional[OptionsWithModel] = None - - -def get_frame_processor() -> Any: - global FRAME_PROCESSOR - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if FRAME_PROCESSOR is None: - model_path = get_options('model').get('path') - FRAME_PROCESSOR = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return FRAME_PROCESSOR - - -def clear_frame_processor() -> None: - global FRAME_PROCESSOR - - FRAME_PROCESSOR = None - - -def get_options(key : Literal['model']) -> Any: - global OPTIONS - - if OPTIONS is None: - OPTIONS =\ - { - 'model': MODELS[frame_processors_globals.frame_colorizer_model] - } - return OPTIONS.get(key) - - -def set_options(key : Literal['model'], value : Any) -> None: - global OPTIONS - - OPTIONS[key] = value - - -def register_args(program : ArgumentParser) -> None: - program.add_argument('--frame-colorizer-model', help = wording.get('help.frame_colorizer_model'), default = config.get_str_value('frame_processors.frame_colorizer_model', 'ddcolor'), choices = frame_processors_choices.frame_colorizer_models) - program.add_argument('--frame-colorizer-blend', help = wording.get('help.frame_colorizer_blend'), type = int, default = config.get_int_value('frame_processors.frame_colorizer_blend', '100'), choices = frame_processors_choices.frame_colorizer_blend_range, metavar = create_metavar(frame_processors_choices.frame_colorizer_blend_range)) - program.add_argument('--frame-colorizer-size', help = wording.get('help.frame_colorizer_size'), type = str, default = config.get_str_value('frame_processors.frame_colorizer_size', '256x256'), choices = frame_processors_choices.frame_colorizer_sizes) - - -def apply_args(program : ArgumentParser) -> None: - args = program.parse_args() - frame_processors_globals.frame_colorizer_model = args.frame_colorizer_model - frame_processors_globals.frame_colorizer_blend = args.frame_colorizer_blend - frame_processors_globals.frame_colorizer_size = args.frame_colorizer_size - - -def pre_check() -> bool: - download_directory_path = resolve_relative_path('../.assets/models') - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, [ model_url ]) - process_manager.end() - return is_file(model_path) - - -def post_check() -> bool: - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download and not is_download_done(model_url, model_path): - logger.error(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) - return False - if not is_file(model_path): - logger.error(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def pre_process(mode : ProcessMode) -> bool: - if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path): - logger.error(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) - return False - if mode == 'output' and not normalize_output_path(facefusion.globals.target_path, facefusion.globals.output_path): - logger.error(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def post_process() -> None: - read_static_image.cache_clear() - if facefusion.globals.video_memory_strategy == 'strict' or facefusion.globals.video_memory_strategy == 'moderate': - clear_frame_processor() - if facefusion.globals.video_memory_strategy == 'strict': - clear_face_analyser() - clear_content_analyser() - - -def colorize_frame(temp_vision_frame : VisionFrame) -> VisionFrame: - frame_processor = get_frame_processor() - prepare_vision_frame = prepare_temp_frame(temp_vision_frame) - with thread_semaphore(): - color_vision_frame = frame_processor.run(None, - { - frame_processor.get_inputs()[0].name: prepare_vision_frame - })[0][0] - color_vision_frame = merge_color_frame(temp_vision_frame, color_vision_frame) - color_vision_frame = blend_frame(temp_vision_frame, color_vision_frame) - return color_vision_frame - - -def prepare_temp_frame(temp_vision_frame : VisionFrame) -> VisionFrame: - model_size = unpack_resolution(frame_processors_globals.frame_colorizer_size) - model_type = get_options('model').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) - 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_options('model').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) - 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_frame(temp_vision_frame : VisionFrame, paste_vision_frame : VisionFrame) -> VisionFrame: - frame_colorizer_blend = 1 - (frame_processors_globals.frame_colorizer_blend / 100) - temp_vision_frame = cv2.addWeighted(temp_vision_frame, frame_colorizer_blend, paste_vision_frame, 1 - frame_colorizer_blend, 0) - return temp_vision_frame - - -def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: - pass - - -def process_frame(inputs : FrameColorizerInputs) -> VisionFrame: - target_vision_frame = inputs.get('target_vision_frame') - return colorize_frame(target_vision_frame) - - -def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: - for queue_payload in process_manager.manage(queue_payloads): - target_vision_path = queue_payload['frame_path'] - target_vision_frame = read_image(target_vision_path) - output_vision_frame = process_frame( - { - 'target_vision_frame': target_vision_frame - }) - write_image(target_vision_path, output_vision_frame) - update_progress(1) - - -def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: - target_vision_frame = read_static_image(target_path) - output_vision_frame = process_frame( - { - 'target_vision_frame': target_vision_frame - }) - write_image(output_path, output_vision_frame) - - -def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: - frame_processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/modules/frame_enhancer.py b/facefusion/processors/frame/modules/frame_enhancer.py deleted file mode 100644 index dd7234a2..00000000 --- a/facefusion/processors/frame/modules/frame_enhancer.py +++ /dev/null @@ -1,263 +0,0 @@ -from typing import Any, List, Literal, Optional -from argparse import ArgumentParser -from time import sleep -import cv2 -import numpy -import onnxruntime - -import facefusion.globals -import facefusion.processors.frame.core as frame_processors -from facefusion import config, process_manager, logger, wording -from facefusion.face_analyser import clear_face_analyser -from facefusion.content_analyser import clear_content_analyser -from facefusion.execution import apply_execution_provider_options -from facefusion.normalizer import normalize_output_path -from facefusion.thread_helper import thread_lock, conditional_thread_semaphore -from facefusion.typing import Face, VisionFrame, UpdateProgress, ProcessMode, ModelSet, OptionsWithModel, QueuePayload -from facefusion.common_helper import create_metavar -from facefusion.filesystem import is_file, resolve_relative_path, is_image, is_video -from facefusion.download import conditional_download, is_download_done -from facefusion.vision import read_image, read_static_image, write_image, merge_tile_frames, create_tile_frames -from facefusion.processors.frame.typings import FrameEnhancerInputs -from facefusion.processors.frame import globals as frame_processors_globals -from facefusion.processors.frame import choices as frame_processors_choices - -FRAME_PROCESSOR = None -NAME = __name__.upper() -MODELS : ModelSet =\ -{ - 'clear_reality_x4': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/clear_reality_x4.onnx', - 'path': resolve_relative_path('../.assets/models/clear_reality_x4.onnx'), - 'size': (128, 8, 4), - 'scale': 4 - }, - 'lsdir_x4': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/lsdir_x4.onnx', - 'path': resolve_relative_path('../.assets/models/lsdir_x4.onnx'), - 'size': (128, 8, 4), - 'scale': 4 - }, - 'nomos8k_sc_x4': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/nomos8k_sc_x4.onnx', - 'path': resolve_relative_path('../.assets/models/nomos8k_sc_x4.onnx'), - 'size': (128, 8, 4), - 'scale': 4 - }, - 'real_esrgan_x2': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/real_esrgan_x2.onnx', - 'path': resolve_relative_path('../.assets/models/real_esrgan_x2.onnx'), - 'size': (256, 16, 8), - 'scale': 2 - }, - 'real_esrgan_x2_fp16': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/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': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/real_esrgan_x4.onnx', - 'path': resolve_relative_path('../.assets/models/real_esrgan_x4.onnx'), - 'size': (256, 16, 8), - 'scale': 4 - }, - 'real_esrgan_x4_fp16': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/real_esrgan_x4_fp16.onnx', - 'path': resolve_relative_path('../.assets/models/real_esrgan_x4_fp16.onnx'), - 'size': (256, 16, 8), - 'scale': 4 - }, - 'real_hatgan_x4': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/real_hatgan_x4.onnx', - 'path': resolve_relative_path('../.assets/models/real_hatgan_x4.onnx'), - 'size': (256, 16, 8), - 'scale': 4 - }, - 'span_kendata_x4': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/span_kendata_x4.onnx', - 'path': resolve_relative_path('../.assets/models/span_kendata_x4.onnx'), - 'size': (128, 8, 4), - 'scale': 4 - }, - 'ultra_sharp_x4': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/ultra_sharp_x4.onnx', - 'path': resolve_relative_path('../.assets/models/ultra_sharp_x4.onnx'), - 'size': (128, 8, 4), - 'scale': 4 - } -} -OPTIONS : Optional[OptionsWithModel] = None - - -def get_frame_processor() -> Any: - global FRAME_PROCESSOR - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if FRAME_PROCESSOR is None: - model_path = get_options('model').get('path') - FRAME_PROCESSOR = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return FRAME_PROCESSOR - - -def clear_frame_processor() -> None: - global FRAME_PROCESSOR - - FRAME_PROCESSOR = None - - -def get_options(key : Literal['model']) -> Any: - global OPTIONS - - if OPTIONS is None: - OPTIONS =\ - { - 'model': MODELS[frame_processors_globals.frame_enhancer_model] - } - return OPTIONS.get(key) - - -def set_options(key : Literal['model'], value : Any) -> None: - global OPTIONS - - OPTIONS[key] = value - - -def register_args(program : ArgumentParser) -> None: - program.add_argument('--frame-enhancer-model', help = wording.get('help.frame_enhancer_model'), default = config.get_str_value('frame_processors.frame_enhancer_model', 'span_kendata_x4'), choices = frame_processors_choices.frame_enhancer_models) - program.add_argument('--frame-enhancer-blend', help = wording.get('help.frame_enhancer_blend'), type = int, default = config.get_int_value('frame_processors.frame_enhancer_blend', '80'), choices = frame_processors_choices.frame_enhancer_blend_range, metavar = create_metavar(frame_processors_choices.frame_enhancer_blend_range)) - - -def apply_args(program : ArgumentParser) -> None: - args = program.parse_args() - frame_processors_globals.frame_enhancer_model = args.frame_enhancer_model - frame_processors_globals.frame_enhancer_blend = args.frame_enhancer_blend - - -def pre_check() -> bool: - download_directory_path = resolve_relative_path('../.assets/models') - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, [ model_url ]) - process_manager.end() - return is_file(model_path) - - -def post_check() -> bool: - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download and not is_download_done(model_url, model_path): - logger.error(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) - return False - if not is_file(model_path): - logger.error(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def pre_process(mode : ProcessMode) -> bool: - if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path): - logger.error(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) - return False - if mode == 'output' and not normalize_output_path(facefusion.globals.target_path, facefusion.globals.output_path): - logger.error(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def post_process() -> None: - read_static_image.cache_clear() - if facefusion.globals.video_memory_strategy == 'strict' or facefusion.globals.video_memory_strategy == 'moderate': - clear_frame_processor() - if facefusion.globals.video_memory_strategy == 'strict': - clear_face_analyser() - clear_content_analyser() - - -def enhance_frame(temp_vision_frame : VisionFrame) -> VisionFrame: - frame_processor = get_frame_processor() - size = get_options('model').get('size') - scale = get_options('model').get('scale') - temp_height, temp_width = temp_vision_frame.shape[:2] - tile_vision_frames, pad_width, pad_height = create_tile_frames(temp_vision_frame, size) - - for index, tile_vision_frame in enumerate(tile_vision_frames): - with conditional_thread_semaphore(facefusion.globals.execution_providers): - tile_vision_frame = frame_processor.run(None, - { - frame_processor.get_inputs()[0].name : prepare_tile_frame(tile_vision_frame) - })[0] - tile_vision_frames[index] = normalize_tile_frame(tile_vision_frame) - merge_vision_frame = merge_tile_frames(tile_vision_frames, temp_width * scale, temp_height * scale, pad_width * scale, pad_height * scale, (size[0] * scale, size[1] * scale, size[2] * scale)) - temp_vision_frame = blend_frame(temp_vision_frame, merge_vision_frame) - return temp_vision_frame - - -def prepare_tile_frame(vision_tile_frame : VisionFrame) -> VisionFrame: - vision_tile_frame = numpy.expand_dims(vision_tile_frame[:, :, ::-1], axis = 0) - vision_tile_frame = vision_tile_frame.transpose(0, 3, 1, 2) - vision_tile_frame = vision_tile_frame.astype(numpy.float32) / 255 - return vision_tile_frame - - -def normalize_tile_frame(vision_tile_frame : VisionFrame) -> VisionFrame: - vision_tile_frame = vision_tile_frame.transpose(0, 2, 3, 1).squeeze(0) * 255 - vision_tile_frame = vision_tile_frame.clip(0, 255).astype(numpy.uint8)[:, :, ::-1] - return vision_tile_frame - - -def blend_frame(temp_vision_frame : VisionFrame, merge_vision_frame : VisionFrame) -> VisionFrame: - frame_enhancer_blend = 1 - (frame_processors_globals.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 = cv2.addWeighted(temp_vision_frame, frame_enhancer_blend, merge_vision_frame, 1 - frame_enhancer_blend, 0) - return temp_vision_frame - - -def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: - pass - - -def process_frame(inputs : FrameEnhancerInputs) -> VisionFrame: - target_vision_frame = inputs.get('target_vision_frame') - return enhance_frame(target_vision_frame) - - -def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: - for queue_payload in process_manager.manage(queue_payloads): - target_vision_path = queue_payload['frame_path'] - target_vision_frame = read_image(target_vision_path) - output_vision_frame = process_frame( - { - 'target_vision_frame': target_vision_frame - }) - write_image(target_vision_path, output_vision_frame) - update_progress(1) - - -def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: - target_vision_frame = read_static_image(target_path) - output_vision_frame = process_frame( - { - 'target_vision_frame': target_vision_frame - }) - write_image(output_path, output_vision_frame) - - -def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: - frame_processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/modules/lip_syncer.py b/facefusion/processors/frame/modules/lip_syncer.py deleted file mode 100755 index 05e76782..00000000 --- a/facefusion/processors/frame/modules/lip_syncer.py +++ /dev/null @@ -1,260 +0,0 @@ -from typing import Any, List, Literal, Optional -from argparse import ArgumentParser -from time import sleep -import cv2 -import numpy -import onnxruntime - -import facefusion.globals -import facefusion.processors.frame.core as frame_processors -from facefusion import config, process_manager, logger, wording -from facefusion.execution import apply_execution_provider_options -from facefusion.face_analyser import get_one_face, get_many_faces, find_similar_faces, clear_face_analyser -from facefusion.face_masker import create_static_box_mask, create_occlusion_mask, create_mouth_mask, clear_face_occluder, clear_face_parser -from facefusion.face_helper import warp_face_by_face_landmark_5, warp_face_by_bounding_box, paste_back, create_bounding_box_from_face_landmark_68 -from facefusion.face_store import get_reference_faces -from facefusion.content_analyser import clear_content_analyser -from facefusion.normalizer import normalize_output_path -from facefusion.thread_helper import thread_lock, conditional_thread_semaphore -from facefusion.typing import Face, VisionFrame, UpdateProgress, ProcessMode, ModelSet, OptionsWithModel, AudioFrame, QueuePayload -from facefusion.filesystem import is_file, has_audio, resolve_relative_path -from facefusion.download import conditional_download, is_download_done -from facefusion.audio import read_static_voice, get_voice_frame, create_empty_audio_frame -from facefusion.filesystem import is_image, is_video, filter_audio_paths -from facefusion.common_helper import get_first -from facefusion.vision import read_image, read_static_image, write_image, restrict_video_fps -from facefusion.processors.frame.typings import LipSyncerInputs -from facefusion.voice_extractor import clear_voice_extractor -from facefusion.processors.frame import globals as frame_processors_globals -from facefusion.processors.frame import choices as frame_processors_choices - -FRAME_PROCESSOR = None -NAME = __name__.upper() -MODELS : ModelSet =\ -{ - 'wav2lip_gan': - { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/wav2lip_gan.onnx', - 'path': resolve_relative_path('../.assets/models/wav2lip_gan.onnx') - } -} -OPTIONS : Optional[OptionsWithModel] = None - - -def get_frame_processor() -> Any: - global FRAME_PROCESSOR - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if FRAME_PROCESSOR is None: - model_path = get_options('model').get('path') - FRAME_PROCESSOR = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return FRAME_PROCESSOR - - -def clear_frame_processor() -> None: - global FRAME_PROCESSOR - - FRAME_PROCESSOR = None - - -def get_options(key : Literal['model']) -> Any: - global OPTIONS - - if OPTIONS is None: - OPTIONS =\ - { - 'model': MODELS[frame_processors_globals.lip_syncer_model] - } - return OPTIONS.get(key) - - -def set_options(key : Literal['model'], value : Any) -> None: - global OPTIONS - - OPTIONS[key] = value - - -def register_args(program : ArgumentParser) -> None: - program.add_argument('--lip-syncer-model', help = wording.get('help.lip_syncer_model'), default = config.get_str_value('frame_processors.lip_syncer_model', 'wav2lip_gan'), choices = frame_processors_choices.lip_syncer_models) - - -def apply_args(program : ArgumentParser) -> None: - args = program.parse_args() - frame_processors_globals.lip_syncer_model = args.lip_syncer_model - - -def pre_check() -> bool: - download_directory_path = resolve_relative_path('../.assets/models') - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, [ model_url ]) - process_manager.end() - return is_file(model_path) - - -def post_check() -> bool: - model_url = get_options('model').get('url') - model_path = get_options('model').get('path') - - if not facefusion.globals.skip_download and not is_download_done(model_url, model_path): - logger.error(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) - return False - if not is_file(model_path): - logger.error(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def pre_process(mode : ProcessMode) -> bool: - if not has_audio(facefusion.globals.source_paths): - logger.error(wording.get('select_audio_source') + wording.get('exclamation_mark'), NAME) - return False - if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path): - logger.error(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) - return False - if mode == 'output' and not normalize_output_path(facefusion.globals.target_path, facefusion.globals.output_path): - logger.error(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) - return False - return True - - -def post_process() -> None: - read_static_image.cache_clear() - read_static_voice.cache_clear() - if facefusion.globals.video_memory_strategy == 'strict' or facefusion.globals.video_memory_strategy == 'moderate': - clear_frame_processor() - if facefusion.globals.video_memory_strategy == 'strict': - clear_face_analyser() - clear_content_analyser() - clear_face_occluder() - clear_face_parser() - clear_voice_extractor() - - -def sync_lip(target_face : Face, temp_audio_frame : AudioFrame, temp_vision_frame : VisionFrame) -> VisionFrame: - frame_processor = get_frame_processor() - crop_mask_list = [] - temp_audio_frame = prepare_audio_frame(temp_audio_frame) - crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmarks.get('5/68'), 'ffhq_512', (512, 512)) - face_landmark_68 = cv2.transform(target_face.landmarks.get('68').reshape(1, -1, 2), affine_matrix).reshape(-1, 2) - bounding_box = create_bounding_box_from_face_landmark_68(face_landmark_68) - bounding_box[1] -= numpy.abs(bounding_box[3] - bounding_box[1]) * 0.125 - mouth_mask = create_mouth_mask(face_landmark_68) - crop_mask_list.append(mouth_mask) - box_mask = create_static_box_mask(crop_vision_frame.shape[:2][::-1], facefusion.globals.face_mask_blur, facefusion.globals.face_mask_padding) - crop_mask_list.append(box_mask) - - if 'occlusion' in facefusion.globals.face_mask_types: - occlusion_mask = create_occlusion_mask(crop_vision_frame) - crop_mask_list.append(occlusion_mask) - close_vision_frame, close_matrix = warp_face_by_bounding_box(crop_vision_frame, bounding_box, (96, 96)) - close_vision_frame = prepare_crop_frame(close_vision_frame) - with conditional_thread_semaphore(facefusion.globals.execution_providers): - close_vision_frame = frame_processor.run(None, - { - 'source': temp_audio_frame, - 'target': close_vision_frame - })[0] - crop_vision_frame = normalize_crop_frame(close_vision_frame) - crop_vision_frame = cv2.warpAffine(crop_vision_frame, cv2.invertAffineTransform(close_matrix), (512, 512), borderMode = cv2.BORDER_REPLICATE) - crop_mask = numpy.minimum.reduce(crop_mask_list) - paste_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, crop_mask, affine_matrix) - return paste_vision_frame - - -def prepare_audio_frame(temp_audio_frame : AudioFrame) -> AudioFrame: - 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) - 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: - crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0) - prepare_vision_frame = crop_vision_frame.copy() - prepare_vision_frame[:, 48:] = 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('float32') / 255.0 - return crop_vision_frame - - -def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: - 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) - return crop_vision_frame - - -def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: - pass - - -def process_frame(inputs : LipSyncerInputs) -> VisionFrame: - reference_faces = inputs.get('reference_faces') - source_audio_frame = inputs.get('source_audio_frame') - target_vision_frame = inputs.get('target_vision_frame') - - if facefusion.globals.face_selector_mode == 'many': - many_faces = get_many_faces(target_vision_frame) - if many_faces: - for target_face in many_faces: - target_vision_frame = sync_lip(target_face, source_audio_frame, target_vision_frame) - if facefusion.globals.face_selector_mode == 'one': - target_face = get_one_face(target_vision_frame) - if target_face: - target_vision_frame = sync_lip(target_face, source_audio_frame, target_vision_frame) - if facefusion.globals.face_selector_mode == 'reference': - similar_faces = find_similar_faces(reference_faces, target_vision_frame, facefusion.globals.reference_face_distance) - if similar_faces: - for similar_face in similar_faces: - target_vision_frame = sync_lip(similar_face, source_audio_frame, target_vision_frame) - return target_vision_frame - - -def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - source_audio_path = get_first(filter_audio_paths(source_paths)) - temp_video_fps = restrict_video_fps(facefusion.globals.target_path, facefusion.globals.output_video_fps) - - for queue_payload in process_manager.manage(queue_payloads): - frame_number = queue_payload['frame_number'] - target_vision_path = queue_payload['frame_path'] - source_audio_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() - target_vision_frame = read_image(target_vision_path) - output_vision_frame = process_frame( - { - 'reference_faces': reference_faces, - 'source_audio_frame': source_audio_frame, - 'target_vision_frame': target_vision_frame - }) - write_image(target_vision_path, output_vision_frame) - update_progress(1) - - -def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - source_audio_frame = create_empty_audio_frame() - target_vision_frame = read_static_image(target_path) - output_vision_frame = process_frame( - { - 'reference_faces': reference_faces, - 'source_audio_frame': source_audio_frame, - 'target_vision_frame': target_vision_frame - }) - write_image(output_path, output_vision_frame) - - -def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: - source_audio_paths = filter_audio_paths(facefusion.globals.source_paths) - temp_video_fps = restrict_video_fps(facefusion.globals.target_path, facefusion.globals.output_video_fps) - for source_audio_path in source_audio_paths: - read_static_voice(source_audio_path, temp_video_fps) - frame_processors.multi_process_frames(source_paths, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/typings.py b/facefusion/processors/frame/typings.py deleted file mode 100644 index aa95e9d5..00000000 --- a/facefusion/processors/frame/typings.py +++ /dev/null @@ -1,41 +0,0 @@ -from typing import Literal, TypedDict - -from facefusion.typing import Face, FaceSet, AudioFrame, VisionFrame - -FaceDebuggerItem = Literal['bounding-box', 'face-landmark-5', 'face-landmark-5/68', 'face-landmark-68', 'face-landmark-68/5', 'face-mask', 'face-detector-score', 'face-landmarker-score', 'age', 'gender'] -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'] -FaceSwapperModel = Literal['blendswap_256', 'inswapper_128', 'inswapper_128_fp16', 'simswap_256', 'simswap_512_unofficial', 'uniface_256'] -FrameColorizerModel = Literal['ddcolor', 'ddcolor_artistic', 'deoldify', 'deoldify_artistic', 'deoldify_stable'] -FrameEnhancerModel = Literal['clear_reality_x4', 'lsdir_x4', 'nomos8k_sc_x4', 'real_esrgan_x2', 'real_esrgan_x2_fp16', 'real_esrgan_x4', 'real_esrgan_x4_fp16', 'real_hatgan_x4', 'span_kendata_x4', 'ultra_sharp_x4'] -LipSyncerModel = Literal['wav2lip_gan'] - -FaceDebuggerInputs = TypedDict('FaceDebuggerInputs', -{ - 'reference_faces' : FaceSet, - 'target_vision_frame' : VisionFrame -}) -FaceEnhancerInputs = TypedDict('FaceEnhancerInputs', -{ - 'reference_faces' : FaceSet, - 'target_vision_frame' : VisionFrame -}) -FaceSwapperInputs = TypedDict('FaceSwapperInputs', -{ - 'reference_faces' : FaceSet, - 'source_face' : Face, - 'target_vision_frame' : VisionFrame -}) -FrameColorizerInputs = TypedDict('FrameColorizerInputs', -{ - 'target_vision_frame' : VisionFrame -}) -FrameEnhancerInputs = TypedDict('FrameEnhancerInputs', -{ - 'target_vision_frame' : VisionFrame -}) -LipSyncerInputs = TypedDict('LipSyncerInputs', -{ - 'reference_faces' : FaceSet, - 'source_audio_frame' : AudioFrame, - 'target_vision_frame' : VisionFrame -}) diff --git a/facefusion/processors/live_portrait.py b/facefusion/processors/live_portrait.py new file mode 100644 index 00000000..06ac82ef --- /dev/null +++ b/facefusion/processors/live_portrait.py @@ -0,0 +1,101 @@ +from typing import Tuple + +import numpy +import scipy + +from facefusion.processors.typing 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_euler_angles(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 = calc_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 calc_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/frame/modules/__init__.py b/facefusion/processors/modules/__init__.py similarity index 100% rename from facefusion/processors/frame/modules/__init__.py rename to facefusion/processors/modules/__init__.py diff --git a/facefusion/processors/modules/age_modifier.py b/facefusion/processors/modules/age_modifier.py new file mode 100755 index 00000000..c1e39498 --- /dev/null +++ b/facefusion/processors/modules/age_modifier.py @@ -0,0 +1,268 @@ +from argparse import ArgumentParser +from typing import Any, List + +import cv2 +import numpy +from cv2.typing import Size +from numpy.typing import NDArray + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, process_manager, state_manager, wording +from facefusion.common_helper import create_int_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.face_analyser import get_many_faces, get_one_face +from facefusion.face_helper import merge_matrix, paste_back, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_occlusion_mask, create_static_box_mask +from facefusion.face_selector import find_similar_faces, sort_and_filter_faces +from facefusion.face_store import get_reference_faces +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors import choices as processors_choices +from facefusion.processors.typing import AgeModifierInputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import thread_semaphore +from facefusion.typing import ApplyStateItem, Args, Face, InferencePool, Mask, ModelOptions, ModelSet, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import read_image, read_static_image, write_image + +MODEL_SET : ModelSet =\ +{ + 'styleganex_age': + { + 'hashes': + { + 'age_modifier': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/styleganex_age.hash', + 'path': resolve_relative_path('../.assets/models/styleganex_age.hash') + } + }, + 'sources': + { + 'age_modifier': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/styleganex_age.onnx', + 'path': resolve_relative_path('../.assets/models/styleganex_age.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + } +} + + +def get_inference_pool() -> InferencePool: + model_sources = get_model_options().get('sources') + model_context = __name__ + '.' + state_manager.get_item('age_modifier_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('age_modifier_model') + inference_manager.clear_inference_pool(model_context) + + +def get_model_options() -> ModelOptions: + age_modifier_model = state_manager.get_item('age_modifier_model') + return MODEL_SET.get(age_modifier_model) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--age-modifier-model', help = wording.get('help.age_modifier_model'), default = config.get_str_value('processors.age_modifier_model', 'styleganex_age'), choices = processors_choices.age_modifier_models) + group_processors.add_argument('--age-modifier-direction', help = wording.get('help.age_modifier_direction'), type = int, default = config.get_int_value('processors.age_modifier_direction', '0'), choices = processors_choices.age_modifier_direction_range, metavar = create_int_metavar(processors_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: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +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(wording.get('choose_image_or_video_target') + wording.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + 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_template = get_model_options().get('template') + model_size = get_model_options().get('size') + crop_size = (model_size[0] // 2, model_size[1] // 2) + face_landmark_5 = target_face.landmark_set.get('5/68').copy() + extend_face_landmark_5 = (face_landmark_5 - face_landmark_5[2]) * 2 + face_landmark_5[2] + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5, model_template, crop_size) + extend_vision_frame, extend_affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, extend_face_landmark_5, model_template, model_size) + extend_vision_frame_raw = extend_vision_frame.copy() + box_mask = create_static_box_mask(model_size, 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) + combined_matrix = merge_matrix([ extend_affine_matrix, cv2.invertAffineTransform(affine_matrix) ]) + occlusion_mask = cv2.warpAffine(occlusion_mask, combined_matrix, model_size) + crop_masks.append(occlusion_mask) + + crop_vision_frame = prepare_vision_frame(crop_vision_frame) + extend_vision_frame = prepare_vision_frame(extend_vision_frame) + extend_vision_frame = forward(crop_vision_frame, extend_vision_frame) + extend_vision_frame = normalize_extend_frame(extend_vision_frame) + extend_vision_frame = fix_color(extend_vision_frame_raw, extend_vision_frame) + extend_crop_mask = cv2.pyrUp(numpy.minimum.reduce(crop_masks).clip(0, 1)) + extend_affine_matrix *= extend_vision_frame.shape[0] / 512 + paste_vision_frame = paste_back(temp_vision_frame, extend_vision_frame, extend_crop_mask, extend_affine_matrix) + return paste_vision_frame + + +def forward(crop_vision_frame : VisionFrame, extend_vision_frame : VisionFrame) -> VisionFrame: + age_modifier = get_inference_pool().get('age_modifier') + age_modifier_inputs = {} + + 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] = prepare_direction(state_manager.get_item('age_modifier_direction')) + + with thread_semaphore(): + crop_vision_frame = age_modifier.run(None, age_modifier_inputs)[0][0] + + return crop_vision_frame + + +def fix_color(extend_vision_frame_raw : VisionFrame, extend_vision_frame : VisionFrame) -> VisionFrame: + color_difference = compute_color_difference(extend_vision_frame_raw, extend_vision_frame, (48, 48)) + color_difference_mask = create_static_box_mask(extend_vision_frame.shape[:2][::-1], 1.0, (0, 0, 0, 0)) + color_difference_mask = numpy.stack((color_difference_mask, ) * 3, axis = -1) + extend_vision_frame = normalize_color_difference(color_difference, color_difference_mask, extend_vision_frame) + return extend_vision_frame + + +def compute_color_difference(extend_vision_frame_raw : VisionFrame, extend_vision_frame : VisionFrame, size : Size) -> VisionFrame: + extend_vision_frame_raw = extend_vision_frame_raw.astype(numpy.float32) / 255 + extend_vision_frame_raw = cv2.resize(extend_vision_frame_raw, size, interpolation = cv2.INTER_AREA) + extend_vision_frame = extend_vision_frame.astype(numpy.float32) / 255 + extend_vision_frame = cv2.resize(extend_vision_frame, size, interpolation = cv2.INTER_AREA) + color_difference = extend_vision_frame_raw - extend_vision_frame + return color_difference + + +def normalize_color_difference(color_difference : VisionFrame, color_difference_mask : Mask, extend_vision_frame : VisionFrame) -> VisionFrame: + color_difference = cv2.resize(color_difference, extend_vision_frame.shape[:2][::-1], interpolation = cv2.INTER_CUBIC) + color_difference_mask = 1 - color_difference_mask.clip(0, 0.75) + extend_vision_frame = extend_vision_frame.astype(numpy.float32) / 255 + extend_vision_frame += color_difference * color_difference_mask + extend_vision_frame = extend_vision_frame.clip(0, 1) + extend_vision_frame = numpy.multiply(extend_vision_frame, 255).astype(numpy.uint8) + return extend_vision_frame + + +def prepare_direction(direction : int) -> NDArray[Any]: + direction = numpy.interp(float(direction), [ -100, 100 ], [ 2.5, -2.5 ]) #type:ignore[assignment] + return numpy.array(direction).astype(numpy.float32) + + +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: + 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.pyrDown(extend_vision_frame) + return extend_vision_frame + + +def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + return modify_age(target_face, temp_vision_frame) + + +def process_frame(inputs : AgeModifierInputs) -> VisionFrame: + reference_faces = inputs.get('reference_faces') + target_vision_frame = inputs.get('target_vision_frame') + many_faces = sort_and_filter_faces(get_many_faces([ target_vision_frame ])) + + if state_manager.get_item('face_selector_mode') == 'many': + if many_faces: + for target_face in many_faces: + target_vision_frame = modify_age(target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'one': + target_face = get_one_face(many_faces) + if target_face: + target_vision_frame = modify_age(target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'reference': + similar_faces = find_similar_faces(many_faces, reference_faces, state_manager.get_item('reference_face_distance')) + if similar_faces: + for similar_face in similar_faces: + target_vision_frame = modify_age(similar_face, target_vision_frame) + return target_vision_frame + + +def process_frames(source_path : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + + for queue_payload in process_manager.manage(queue_payloads): + target_vision_path = queue_payload['frame_path'] + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_path : str, target_path : str, output_path : str) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/modules/expression_restorer.py b/facefusion/processors/modules/expression_restorer.py new file mode 100755 index 00000000..1068bccc --- /dev/null +++ b/facefusion/processors/modules/expression_restorer.py @@ -0,0 +1,290 @@ +from argparse import ArgumentParser +from typing import List, Tuple + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, process_manager, state_manager, wording +from facefusion.common_helper import create_int_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.face_analyser import get_many_faces, get_one_face +from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_occlusion_mask, create_static_box_mask +from facefusion.face_selector import find_similar_faces, sort_and_filter_faces +from facefusion.face_store import get_reference_faces +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors import choices as processors_choices +from facefusion.processors.live_portrait import create_rotation, limit_expression +from facefusion.processors.typing import ExpressionRestorerInputs +from facefusion.processors.typing import LivePortraitExpression, LivePortraitFeatureVolume, LivePortraitMotionPoints, LivePortraitPitch, LivePortraitRoll, LivePortraitScale, LivePortraitTranslation, LivePortraitYaw +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore, thread_semaphore +from facefusion.typing import ApplyStateItem, Args, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import get_video_frame, read_image, read_static_image, write_image + +MODEL_SET : ModelSet =\ +{ + 'live_portrait': + { + 'hashes': + { + 'feature_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_feature_extractor.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_feature_extractor.hash') + }, + 'motion_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_motion_extractor.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_motion_extractor.hash') + }, + 'generator': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_generator.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_generator.hash') + } + }, + 'sources': + { + 'feature_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_feature_extractor.onnx', + 'path': resolve_relative_path('../.assets/models/live_portrait_feature_extractor.onnx') + }, + 'motion_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_motion_extractor.onnx', + 'path': resolve_relative_path('../.assets/models/live_portrait_motion_extractor.onnx') + }, + 'generator': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_generator.onnx', + 'path': resolve_relative_path('../.assets/models/live_portrait_generator.onnx') + } + }, + 'template': 'arcface_128_v2', + 'size': (512, 512) + } +} + + +def get_inference_pool() -> InferencePool: + model_sources = get_model_options().get('sources') + model_context = __name__ + '.' + state_manager.get_item('expression_restorer_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + inference_manager.clear_inference_pool(__name__) + + +def get_model_options() -> ModelOptions: + expression_restorer_model = state_manager.get_item('expression_restorer_model') + return MODEL_SET.get(expression_restorer_model) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--expression-restorer-model', help = wording.get('help.expression_restorer_model'), default = config.get_str_value('processors.expression_restorer_model', 'live_portrait'), choices = processors_choices.expression_restorer_models) + group_processors.add_argument('--expression-restorer-factor', help = wording.get('help.expression_restorer_factor'), type = int, default = config.get_int_value('processors.expression_restorer_factor', '80'), choices = processors_choices.expression_restorer_factor_range, metavar = create_int_metavar(processors_choices.expression_restorer_factor_range)) + facefusion.jobs.job_store.register_step_keys([ 'expression_restorer_model','expression_restorer_factor' ]) + + +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')) + + +def pre_check() -> bool: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +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(wording.get('choose_image_or_video_target') + wording.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + 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(source_vision_frame : VisionFrame, target_face : Face, 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 ])) + source_vision_frame = cv2.resize(source_vision_frame, temp_vision_frame.shape[:2][::-1]) + source_crop_vision_frame, _ = warp_face_by_face_landmark_5(source_vision_frame, target_face.landmark_set.get('5/68'), model_template, model_size) + target_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_static_box_mask(target_crop_vision_frame.shape[:2][::-1], 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(target_crop_vision_frame) + crop_masks.append(occlusion_mask) + + source_crop_vision_frame = prepare_crop_frame(source_crop_vision_frame) + target_crop_vision_frame = prepare_crop_frame(target_crop_vision_frame) + target_crop_vision_frame = apply_restore(source_crop_vision_frame, target_crop_vision_frame, expression_restorer_factor) + target_crop_vision_frame = normalize_crop_frame(target_crop_vision_frame) + crop_mask = numpy.minimum.reduce(crop_masks).clip(0, 1) + temp_vision_frame = paste_back(temp_vision_frame, target_crop_vision_frame, crop_mask, affine_matrix) + return temp_vision_frame + + +def apply_restore(source_crop_vision_frame : VisionFrame, target_crop_vision_frame : VisionFrame, expression_restorer_factor : float) -> VisionFrame: + feature_volume = forward_extract_feature(target_crop_vision_frame) + source_expression = forward_extract_motion(source_crop_vision_frame)[5] + pitch, yaw, roll, scale, translation, target_expression, motion_points = forward_extract_motion(target_crop_vision_frame) + rotation = create_rotation(pitch, yaw, roll) + source_expression[:, [ 0, 4, 5, 8, 9 ]] = target_expression[:, [ 0, 4, 5, 8, 9 ]] + source_expression = source_expression * expression_restorer_factor + target_expression * (1 - expression_restorer_factor) + source_expression = limit_expression(source_expression) + source_motion_points = scale * (motion_points @ rotation.T + source_expression) + translation + target_motion_points = scale * (motion_points @ rotation.T + target_expression) + translation + crop_vision_frame = forward_generate_frame(feature_volume, source_motion_points, target_motion_points) + 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_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 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 get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + pass + + +def process_frame(inputs : ExpressionRestorerInputs) -> VisionFrame: + reference_faces = inputs.get('reference_faces') + source_vision_frame = inputs.get('source_vision_frame') + target_vision_frame = inputs.get('target_vision_frame') + many_faces = sort_and_filter_faces(get_many_faces([ target_vision_frame ])) + + if state_manager.get_item('face_selector_mode') == 'many': + if many_faces: + for target_face in many_faces: + target_vision_frame = restore_expression(source_vision_frame, target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'one': + target_face = get_one_face(many_faces) + if target_face: + target_vision_frame = restore_expression(source_vision_frame, target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'reference': + similar_faces = find_similar_faces(many_faces, reference_faces, state_manager.get_item('reference_face_distance')) + if similar_faces: + for similar_face in similar_faces: + target_vision_frame = restore_expression(source_vision_frame, similar_face, target_vision_frame) + return target_vision_frame + + +def process_frames(source_path : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + + for queue_payload in process_manager.manage(queue_payloads): + frame_number = queue_payload.get('frame_number') + if state_manager.get_item('trim_frame_start'): + frame_number += state_manager.get_item('trim_frame_start') + source_vision_frame = get_video_frame(state_manager.get_item('target_path'), frame_number) + target_vision_path = queue_payload.get('frame_path') + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'source_vision_frame': source_vision_frame, + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_path : str, target_path : str, output_path : str) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + source_vision_frame = read_static_image(state_manager.get_item('target_path')) + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'source_vision_frame': source_vision_frame, + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/modules/face_debugger.py b/facefusion/processors/modules/face_debugger.py new file mode 100755 index 00000000..2dd11e46 --- /dev/null +++ b/facefusion/processors/modules/face_debugger.py @@ -0,0 +1,222 @@ +from argparse import ArgumentParser +from typing import List + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, logger, process_manager, state_manager, wording +from facefusion.face_analyser import get_many_faces, get_one_face +from facefusion.face_helper import warp_face_by_face_landmark_5 +from facefusion.face_masker import create_occlusion_mask, create_region_mask, create_static_box_mask +from facefusion.face_selector import find_similar_faces, sort_and_filter_faces +from facefusion.face_store import get_reference_faces +from facefusion.filesystem import in_directory, same_file_extension +from facefusion.processors import choices as processors_choices +from facefusion.processors.typing import FaceDebuggerInputs +from facefusion.program_helper import find_argument_group +from facefusion.typing import ApplyStateItem, Args, Face, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import read_image, read_static_image, write_image + + +def get_inference_pool() -> None: + 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 = wording.get('help.face_debugger_items').format(choices = ', '.join(processors_choices.face_debugger_items)), default = config.get_str_list('processors.face_debugger_items', 'face-landmark-5/68 face-mask'), choices = processors_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 == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + 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: + primary_color = (0, 0, 255) + primary_light_color = (100, 100, 255) + secondary_color = (0, 255, 0) + tertiary_color = (255, 255, 0) + bounding_box = target_face.bounding_box.astype(numpy.int32) + temp_vision_frame = temp_vision_frame.copy() + has_face_landmark_5_fallback = numpy.array_equal(target_face.landmark_set.get('5'), target_face.landmark_set.get('5/68')) + has_face_landmark_68_fallback = numpy.array_equal(target_face.landmark_set.get('68'), target_face.landmark_set.get('68/5')) + face_debugger_items = state_manager.get_item('face_debugger_items') + + if 'bounding-box' in face_debugger_items: + x1, y1, x2, y2 = bounding_box + cv2.rectangle(temp_vision_frame, (x1, y1), (x2, y2), primary_color, 2) + + if target_face.angle == 0: + cv2.line(temp_vision_frame, (x1, y1), (x2, y1), primary_light_color, 3) + elif target_face.angle == 180: + cv2.line(temp_vision_frame, (x1, y2), (x2, y2), primary_light_color, 3) + elif target_face.angle == 90: + cv2.line(temp_vision_frame, (x2, y1), (x2, y2), primary_light_color, 3) + elif target_face.angle == 270: + cv2.line(temp_vision_frame, (x1, y1), (x1, y2), primary_light_color, 3) + + if 'face-mask' in face_debugger_items: + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmark_set.get('5/68'), 'arcface_128_v2', (512, 512)) + inverse_matrix = cv2.invertAffineTransform(affine_matrix) + temp_size = temp_vision_frame.shape[:2][::-1] + crop_masks = [] + + if 'box' in state_manager.get_item('face_mask_types'): + box_mask = create_static_box_mask(crop_vision_frame.shape[:2][::-1], 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 '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_vision_frame[inverse_vision_frame > 0] = 255 #type:ignore[operator] + inverse_contours = cv2.findContours(inverse_vision_frame, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)[0] + cv2.drawContours(temp_vision_frame, inverse_contours, -1, tertiary_color if has_face_landmark_5_fallback else secondary_color, 2) + + if 'face-landmark-5' in face_debugger_items and numpy.any(target_face.landmark_set.get('5')): + face_landmark_5 = target_face.landmark_set.get('5').astype(numpy.int32) + for index in range(face_landmark_5.shape[0]): + cv2.circle(temp_vision_frame, (face_landmark_5[index][0], face_landmark_5[index][1]), 3, primary_color, -1) + + if 'face-landmark-5/68' in face_debugger_items and numpy.any(target_face.landmark_set.get('5/68')): + face_landmark_5_68 = target_face.landmark_set.get('5/68').astype(numpy.int32) + for index in range(face_landmark_5_68.shape[0]): + cv2.circle(temp_vision_frame, (face_landmark_5_68[index][0], face_landmark_5_68[index][1]), 3, tertiary_color if has_face_landmark_5_fallback else secondary_color, -1) + + if 'face-landmark-68' in face_debugger_items and numpy.any(target_face.landmark_set.get('68')): + face_landmark_68 = target_face.landmark_set.get('68').astype(numpy.int32) + for index in range(face_landmark_68.shape[0]): + cv2.circle(temp_vision_frame, (face_landmark_68[index][0], face_landmark_68[index][1]), 3, tertiary_color if has_face_landmark_68_fallback else secondary_color, -1) + + if 'face-landmark-68/5' in face_debugger_items and numpy.any(target_face.landmark_set.get('68')): + face_landmark_68 = target_face.landmark_set.get('68/5').astype(numpy.int32) + for index in range(face_landmark_68.shape[0]): + cv2.circle(temp_vision_frame, (face_landmark_68[index][0], face_landmark_68[index][1]), 3, tertiary_color, -1) + + if bounding_box[3] - bounding_box[1] > 50 and bounding_box[2] - bounding_box[0] > 50: + top = bounding_box[1] + left = bounding_box[0] - 20 + + if 'face-detector-score' in face_debugger_items: + face_score_text = str(round(target_face.score_set.get('detector'), 2)) + top = top + 20 + cv2.putText(temp_vision_frame, face_score_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, primary_color, 2) + + if 'face-landmarker-score' in face_debugger_items: + face_score_text = str(round(target_face.score_set.get('landmarker'), 2)) + top = top + 20 + cv2.putText(temp_vision_frame, face_score_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, tertiary_color if has_face_landmark_5_fallback else secondary_color, 2) + + if 'age' in face_debugger_items: + face_age_text = str(target_face.age.start) + '-' + str(target_face.age.stop) + top = top + 20 + cv2.putText(temp_vision_frame, face_age_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, primary_color, 2) + + if 'gender' in face_debugger_items: + face_gender_text = target_face.gender + top = top + 20 + cv2.putText(temp_vision_frame, face_gender_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, primary_color, 2) + + if 'race' in face_debugger_items: + face_race_text = target_face.race + top = top + 20 + cv2.putText(temp_vision_frame, face_race_text, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, primary_color, 2) + + return temp_vision_frame + + +def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + pass + + +def process_frame(inputs : FaceDebuggerInputs) -> VisionFrame: + reference_faces = inputs.get('reference_faces') + target_vision_frame = inputs.get('target_vision_frame') + many_faces = sort_and_filter_faces(get_many_faces([ target_vision_frame ])) + + if state_manager.get_item('face_selector_mode') == 'many': + if many_faces: + for target_face in many_faces: + target_vision_frame = debug_face(target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'one': + target_face = get_one_face(many_faces) + if target_face: + target_vision_frame = debug_face(target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'reference': + similar_faces = find_similar_faces(many_faces, reference_faces, state_manager.get_item('reference_face_distance')) + if similar_faces: + for similar_face in similar_faces: + target_vision_frame = debug_face(similar_face, target_vision_frame) + return target_vision_frame + + +def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + + for queue_payload in process_manager.manage(queue_payloads): + target_vision_path = queue_payload['frame_path'] + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + processors.multi_process_frames(source_paths, temp_frame_paths, process_frames) diff --git a/facefusion/processors/modules/face_editor.py b/facefusion/processors/modules/face_editor.py new file mode 100755 index 00000000..7e77d200 --- /dev/null +++ b/facefusion/processors/modules/face_editor.py @@ -0,0 +1,528 @@ +from argparse import ArgumentParser +from typing import List, Tuple + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, process_manager, state_manager, wording +from facefusion.common_helper import create_float_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.face_analyser import get_many_faces, get_one_face +from facefusion.face_helper import paste_back, scale_face_landmark_5, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_static_box_mask +from facefusion.face_selector import find_similar_faces, sort_and_filter_faces +from facefusion.face_store import get_reference_faces +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors import choices as processors_choices +from facefusion.processors.live_portrait import create_rotation, limit_euler_angles, limit_expression +from facefusion.processors.typing import FaceEditorInputs, LivePortraitExpression, LivePortraitFeatureVolume, LivePortraitMotionPoints, LivePortraitPitch, LivePortraitRoll, LivePortraitRotation, LivePortraitScale, LivePortraitTranslation, LivePortraitYaw +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore, thread_semaphore +from facefusion.typing import ApplyStateItem, Args, Face, FaceLandmark68, InferencePool, ModelOptions, ModelSet, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import read_image, read_static_image, write_image + +MODEL_SET : ModelSet =\ +{ + 'live_portrait': + { + 'hashes': + { + 'feature_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_feature_extractor.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_feature_extractor.hash') + }, + 'motion_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_motion_extractor.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_motion_extractor.hash') + }, + 'eye_retargeter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_eye_retargeter.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_eye_retargeter.hash') + }, + 'lip_retargeter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_lip_retargeter.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_lip_retargeter.hash') + }, + 'stitcher': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_stitcher.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_stitcher.hash') + }, + 'generator': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_generator.hash', + 'path': resolve_relative_path('../.assets/models/live_portrait_generator.hash') + } + }, + 'sources': + { + 'feature_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_feature_extractor.onnx', + 'path': resolve_relative_path('../.assets/models/live_portrait_feature_extractor.onnx') + }, + 'motion_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_motion_extractor.onnx', + 'path': resolve_relative_path('../.assets/models/live_portrait_motion_extractor.onnx') + }, + 'eye_retargeter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_eye_retargeter.onnx', + 'path': resolve_relative_path('../.assets/models/live_portrait_eye_retargeter.onnx') + }, + 'lip_retargeter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_lip_retargeter.onnx', + 'path': resolve_relative_path('../.assets/models/live_portrait_lip_retargeter.onnx') + }, + 'stitcher': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/live_portrait_stitcher.onnx', + 'path': resolve_relative_path('../.assets/models/live_portrait_stitcher.onnx') + }, + 'generator': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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_sources = get_model_options().get('sources') + model_context = __name__ + '.' + state_manager.get_item('face_editor_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('face_editor_model') + inference_manager.clear_inference_pool(model_context) + + +def get_model_options() -> ModelOptions: + face_editor_model = state_manager.get_item('face_editor_model') + return MODEL_SET.get(face_editor_model) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--face-editor-model', help = wording.get('help.face_editor_model'), default = config.get_str_value('processors.face_editor_model', 'live_portrait'), choices = processors_choices.face_editor_models) + group_processors.add_argument('--face-editor-eyebrow-direction', help = wording.get('help.face_editor_eyebrow_direction'), type = float, default = config.get_float_value('processors.face_editor_eyebrow_direction', '0'), choices = processors_choices.face_editor_eyebrow_direction_range, metavar = create_float_metavar(processors_choices.face_editor_eyebrow_direction_range)) + group_processors.add_argument('--face-editor-eye-gaze-horizontal', help = wording.get('help.face_editor_eye_gaze_horizontal'), type = float, default = config.get_float_value('processors.face_editor_eye_gaze_horizontal', '0'), choices = processors_choices.face_editor_eye_gaze_horizontal_range, metavar = create_float_metavar(processors_choices.face_editor_eye_gaze_horizontal_range)) + group_processors.add_argument('--face-editor-eye-gaze-vertical', help = wording.get('help.face_editor_eye_gaze_vertical'), type = float, default = config.get_float_value('processors.face_editor_eye_gaze_vertical', '0'), choices = processors_choices.face_editor_eye_gaze_vertical_range, metavar = create_float_metavar(processors_choices.face_editor_eye_gaze_vertical_range)) + group_processors.add_argument('--face-editor-eye-open-ratio', help = wording.get('help.face_editor_eye_open_ratio'), type = float, default = config.get_float_value('processors.face_editor_eye_open_ratio', '0'), choices = processors_choices.face_editor_eye_open_ratio_range, metavar = create_float_metavar(processors_choices.face_editor_eye_open_ratio_range)) + group_processors.add_argument('--face-editor-lip-open-ratio', help = wording.get('help.face_editor_lip_open_ratio'), type = float, default = config.get_float_value('processors.face_editor_lip_open_ratio', '0'), choices = processors_choices.face_editor_lip_open_ratio_range, metavar = create_float_metavar(processors_choices.face_editor_lip_open_ratio_range)) + group_processors.add_argument('--face-editor-mouth-grim', help = wording.get('help.face_editor_mouth_grim'), type = float, default = config.get_float_value('processors.face_editor_mouth_grim', '0'), choices = processors_choices.face_editor_mouth_grim_range, metavar = create_float_metavar(processors_choices.face_editor_mouth_grim_range)) + group_processors.add_argument('--face-editor-mouth-pout', help = wording.get('help.face_editor_mouth_pout'), type = float, default = config.get_float_value('processors.face_editor_mouth_pout', '0'), choices = processors_choices.face_editor_mouth_pout_range, metavar = create_float_metavar(processors_choices.face_editor_mouth_pout_range)) + group_processors.add_argument('--face-editor-mouth-purse', help = wording.get('help.face_editor_mouth_purse'), type = float, default = config.get_float_value('processors.face_editor_mouth_purse', '0'), choices = processors_choices.face_editor_mouth_purse_range, metavar = create_float_metavar(processors_choices.face_editor_mouth_purse_range)) + group_processors.add_argument('--face-editor-mouth-smile', help = wording.get('help.face_editor_mouth_smile'), type = float, default = config.get_float_value('processors.face_editor_mouth_smile', '0'), choices = processors_choices.face_editor_mouth_smile_range, metavar = create_float_metavar(processors_choices.face_editor_mouth_smile_range)) + group_processors.add_argument('--face-editor-mouth-position-horizontal', help = wording.get('help.face_editor_mouth_position_horizontal'), type = float, default = config.get_float_value('processors.face_editor_mouth_position_horizontal', '0'), choices = processors_choices.face_editor_mouth_position_horizontal_range, metavar = create_float_metavar(processors_choices.face_editor_mouth_position_horizontal_range)) + group_processors.add_argument('--face-editor-mouth-position-vertical', help = wording.get('help.face_editor_mouth_position_vertical'), type = float, default = config.get_float_value('processors.face_editor_mouth_position_vertical', '0'), choices = processors_choices.face_editor_mouth_position_vertical_range, metavar = create_float_metavar(processors_choices.face_editor_mouth_position_vertical_range)) + group_processors.add_argument('--face-editor-head-pitch', help = wording.get('help.face_editor_head_pitch'), type = float, default = config.get_float_value('processors.face_editor_head_pitch', '0'), choices = processors_choices.face_editor_head_pitch_range, metavar = create_float_metavar(processors_choices.face_editor_head_pitch_range)) + group_processors.add_argument('--face-editor-head-yaw', help=wording.get('help.face_editor_head_yaw'), type = float, default = config.get_float_value('processors.face_editor_head_yaw', '0'), choices = processors_choices.face_editor_head_yaw_range, metavar = create_float_metavar(processors_choices.face_editor_head_yaw_range)) + group_processors.add_argument('--face-editor-head-roll', help=wording.get('help.face_editor_head_roll'), type = float, default = config.get_float_value('processors.face_editor_head_roll', '0'), choices = processors_choices.face_editor_head_roll_range, metavar = create_float_metavar(processors_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: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +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(wording.get('choose_image_or_video_target') + wording.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + 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_static_box_mask(crop_vision_frame.shape[:2][::-1], 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) + temp_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, box_mask, affine_matrix) + return temp_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 = calc_distance_ratio(face_landmark_68, 37, 40, 39, 36) + right_eye_ratio = calc_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 = calc_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_euler_angles(pitch, yaw, roll, edit_pitch, edit_yaw, edit_roll) + rotation = create_rotation(edit_pitch, edit_yaw, edit_roll) + return rotation + + +def calc_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 get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + pass + + +def process_frame(inputs : FaceEditorInputs) -> VisionFrame: + reference_faces = inputs.get('reference_faces') + target_vision_frame = inputs.get('target_vision_frame') + many_faces = sort_and_filter_faces(get_many_faces([ target_vision_frame ])) + + if state_manager.get_item('face_selector_mode') == 'many': + if many_faces: + for target_face in many_faces: + target_vision_frame = edit_face(target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'one': + target_face = get_one_face(many_faces) + if target_face: + target_vision_frame = edit_face(target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'reference': + similar_faces = find_similar_faces(many_faces, reference_faces, state_manager.get_item('reference_face_distance')) + if similar_faces: + for similar_face in similar_faces: + target_vision_frame = edit_face(similar_face, target_vision_frame) + return target_vision_frame + + +def process_frames(source_path : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + + for queue_payload in process_manager.manage(queue_payloads): + target_vision_path = queue_payload['frame_path'] + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_path : str, target_path : str, output_path : str) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/modules/face_enhancer.py b/facefusion/processors/modules/face_enhancer.py new file mode 100755 index 00000000..da8d6498 --- /dev/null +++ b/facefusion/processors/modules/face_enhancer.py @@ -0,0 +1,397 @@ +from argparse import ArgumentParser +from typing import List + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, process_manager, state_manager, wording +from facefusion.common_helper import create_int_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.face_analyser import get_many_faces, get_one_face +from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_occlusion_mask, create_static_box_mask +from facefusion.face_selector import find_similar_faces, sort_and_filter_faces +from facefusion.face_store import get_reference_faces +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors import choices as processors_choices +from facefusion.processors.typing import FaceEnhancerInputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import thread_semaphore +from facefusion.typing import ApplyStateItem, Args, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import read_image, read_static_image, write_image + +MODEL_SET : ModelSet =\ +{ + 'codeformer': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/codeformer.hash', + 'path': resolve_relative_path('../.assets/models/codeformer.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/codeformer.onnx', + 'path': resolve_relative_path('../.assets/models/codeformer.onnx') + } + }, + 'template': 'ffhq_512', + 'size': (512, 512) + }, + 'gfpgan_1.2': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/gfpgan_1.2.hash', + 'path': resolve_relative_path('../.assets/models/gfpgan_1.2.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/gfpgan_1.3.hash', + 'path': resolve_relative_path('../.assets/models/gfpgan_1.3.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/gfpgan_1.4.hash', + 'path': resolve_relative_path('../.assets/models/gfpgan_1.4.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/gpen_bfr_256.hash', + 'path': resolve_relative_path('../.assets/models/gpen_bfr_256.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/gpen_bfr_256.onnx', + 'path': resolve_relative_path('../.assets/models/gpen_bfr_256.onnx') + } + }, + 'template': 'arcface_128_v2', + 'size': (256, 256) + }, + 'gpen_bfr_512': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/gpen_bfr_512.hash', + 'path': resolve_relative_path('../.assets/models/gpen_bfr_512.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/gpen_bfr_1024.hash', + 'path': resolve_relative_path('../.assets/models/gpen_bfr_1024.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/gpen_bfr_2048.hash', + 'path': resolve_relative_path('../.assets/models/gpen_bfr_2048.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/restoreformer_plus_plus.hash', + 'path': resolve_relative_path('../.assets/models/restoreformer_plus_plus.hash') + } + }, + 'sources': + { + 'face_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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_sources = get_model_options().get('sources') + model_context = __name__ + '.' + state_manager.get_item('face_enhancer_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('face_enhancer_model') + inference_manager.clear_inference_pool(model_context) + + +def get_model_options() -> ModelOptions: + face_enhancer_model = state_manager.get_item('face_enhancer_model') + return MODEL_SET.get(face_enhancer_model) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--face-enhancer-model', help = wording.get('help.face_enhancer_model'), default = config.get_str_value('processors.face_enhancer_model', 'gfpgan_1.4'), choices = processors_choices.face_enhancer_models) + group_processors.add_argument('--face-enhancer-blend', help = wording.get('help.face_enhancer_blend'), type = int, default = config.get_int_value('processors.face_enhancer_blend', '80'), choices = processors_choices.face_enhancer_blend_range, metavar = create_int_metavar(processors_choices.face_enhancer_blend_range)) + facefusion.jobs.job_store.register_step_keys([ 'face_enhancer_model', 'face_enhancer_blend' ]) + + +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')) + + +def pre_check() -> bool: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +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(wording.get('choose_image_or_video_target') + wording.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + 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_static_box_mask(crop_vision_frame.shape[:2][::-1], 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) + crop_vision_frame = forward(crop_vision_frame) + 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_frame(temp_vision_frame, paste_vision_frame) + return temp_vision_frame + + +def forward(crop_vision_frame : VisionFrame) -> 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': + weight = numpy.array([ 1 ]).astype(numpy.double) + face_enhancer_inputs[face_enhancer_input.name] = weight + + with thread_semaphore(): + crop_vision_frame = face_enhancer.run(None, face_enhancer_inputs)[0][0] + + return crop_vision_frame + + +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_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 = cv2.addWeighted(temp_vision_frame, face_enhancer_blend, paste_vision_frame, 1 - face_enhancer_blend, 0) + return temp_vision_frame + + +def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + return enhance_face(target_face, temp_vision_frame) + + +def process_frame(inputs : FaceEnhancerInputs) -> VisionFrame: + reference_faces = inputs.get('reference_faces') + target_vision_frame = inputs.get('target_vision_frame') + many_faces = sort_and_filter_faces(get_many_faces([ target_vision_frame ])) + + if state_manager.get_item('face_selector_mode') == 'many': + if many_faces: + for target_face in many_faces: + target_vision_frame = enhance_face(target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'one': + target_face = get_one_face(many_faces) + if target_face: + target_vision_frame = enhance_face(target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'reference': + similar_faces = find_similar_faces(many_faces, reference_faces, state_manager.get_item('reference_face_distance')) + if similar_faces: + for similar_face in similar_faces: + target_vision_frame = enhance_face(similar_face, target_vision_frame) + return target_vision_frame + + +def process_frames(source_path : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + + for queue_payload in process_manager.manage(queue_payloads): + target_vision_path = queue_payload['frame_path'] + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_path : str, target_path : str, output_path : str) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/modules/face_swapper.py b/facefusion/processors/modules/face_swapper.py new file mode 100755 index 00000000..6aa63733 --- /dev/null +++ b/facefusion/processors/modules/face_swapper.py @@ -0,0 +1,564 @@ +from argparse import ArgumentParser +from typing import List, Tuple + +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, process_manager, state_manager, wording +from facefusion.common_helper import get_first +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.execution import has_execution_provider +from facefusion.face_analyser import get_average_face, get_many_faces, get_one_face +from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 +from facefusion.face_masker import create_occlusion_mask, create_region_mask, create_static_box_mask +from facefusion.face_selector import find_similar_faces, sort_and_filter_faces +from facefusion.face_store import get_reference_faces +from facefusion.filesystem import filter_image_paths, has_image, in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.inference_manager import get_static_model_initializer +from facefusion.processors import choices as processors_choices +from facefusion.processors.pixel_boost import explode_pixel_boost, implode_pixel_boost +from facefusion.processors.typing import FaceSwapperInputs +from facefusion.program_helper import find_argument_group, suggest_face_swapper_pixel_boost_choices +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.typing import ApplyStateItem, Args, Embedding, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import read_image, read_static_image, read_static_images, unpack_resolution, write_image + +MODEL_SET : ModelSet =\ +{ + 'blendswap_256': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/blendswap_256.hash', + 'path': resolve_relative_path('../.assets/models/blendswap_256.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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_256_unet_1': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ghost_256_unet_1.hash', + 'path': resolve_relative_path('../.assets/models/ghost_256_unet_1.hash') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_ghost.hash', + 'path': resolve_relative_path('../.assets/models/arcface_converter_ghost.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ghost_256_unet_1.onnx', + 'path': resolve_relative_path('../.assets/models/ghost_256_unet_1.onnx') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_ghost.onnx', + 'path': resolve_relative_path('../.assets/models/arcface_converter_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_256_unet_2': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ghost_256_unet_2.hash', + 'path': resolve_relative_path('../.assets/models/ghost_256_unet_2.hash') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_ghost.hash', + 'path': resolve_relative_path('../.assets/models/arcface_converter_ghost.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ghost_256_unet_2.onnx', + 'path': resolve_relative_path('../.assets/models/ghost_256_unet_2.onnx') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_ghost.onnx', + 'path': resolve_relative_path('../.assets/models/arcface_converter_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_256_unet_3': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ghost_256_unet_3.hash', + 'path': resolve_relative_path('../.assets/models/ghost_256_unet_3.hash') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_ghost.hash', + 'path': resolve_relative_path('../.assets/models/arcface_converter_ghost.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ghost_256_unet_3.onnx', + 'path': resolve_relative_path('../.assets/models/ghost_256_unet_3.onnx') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_ghost.onnx', + 'path': resolve_relative_path('../.assets/models/arcface_converter_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 ] + }, + 'inswapper_128': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/inswapper_128.hash', + 'path': resolve_relative_path('../.assets/models/inswapper_128.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/inswapper_128.onnx', + 'path': resolve_relative_path('../.assets/models/inswapper_128.onnx') + } + }, + 'type': 'inswapper', + 'template': 'arcface_128_v2', + 'size': (128, 128), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'inswapper_128_fp16': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/inswapper_128_fp16.hash', + 'path': resolve_relative_path('../.assets/models/inswapper_128_fp16.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/inswapper_128_fp16.onnx', + 'path': resolve_relative_path('../.assets/models/inswapper_128_fp16.onnx') + } + }, + 'type': 'inswapper', + 'template': 'arcface_128_v2', + 'size': (128, 128), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'simswap_256': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/simswap_256.hash', + 'path': resolve_relative_path('../.assets/models/simswap_256.hash') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_simswap.hash', + 'path': resolve_relative_path('../.assets/models/arcface_converter_simswap.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/simswap_256.onnx', + 'path': resolve_relative_path('../.assets/models/simswap_256.onnx') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_simswap.onnx', + 'path': resolve_relative_path('../.assets/models/arcface_converter_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_512_unofficial': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/simswap_512_unofficial.hash', + 'path': resolve_relative_path('../.assets/models/simswap_512_unofficial.hash') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_simswap.hash', + 'path': resolve_relative_path('../.assets/models/arcface_converter_simswap.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/simswap_512_unofficial.onnx', + 'path': resolve_relative_path('../.assets/models/simswap_512_unofficial.onnx') + }, + 'embedding_converter': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/arcface_converter_simswap.onnx', + 'path': resolve_relative_path('../.assets/models/arcface_converter_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': + { + 'hashes': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/uniface_256.hash', + 'path': resolve_relative_path('../.assets/models/uniface_256.hash') + } + }, + 'sources': + { + 'face_swapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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_sources = get_model_options().get('sources') + model_context = __name__ + '.' + state_manager.get_item('face_swapper_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('face_swapper_model') + inference_manager.clear_inference_pool(model_context) + + +def get_model_options() -> ModelOptions: + face_swapper_model = state_manager.get_item('face_swapper_model') + face_swapper_model = 'inswapper_128' if has_execution_provider('coreml') and face_swapper_model == 'inswapper_128_fp16' else face_swapper_model + return MODEL_SET.get(face_swapper_model) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--face-swapper-model', help = wording.get('help.face_swapper_model'), default = config.get_str_value('processors.face_swapper_model', 'inswapper_128_fp16'), choices = processors_choices.face_swapper_set.keys()) + face_swapper_pixel_boost_choices = suggest_face_swapper_pixel_boost_choices(program) + group_processors.add_argument('--face-swapper-pixel-boost', help = wording.get('help.face_swapper_pixel_boost'), default = config.get_str_value('processors.face_swapper_pixel_boost', get_first(face_swapper_pixel_boost_choices)), choices = face_swapper_pixel_boost_choices) + facefusion.jobs.job_store.register_step_keys([ 'face_swapper_model', 'face_swapper_pixel_boost' ]) + + +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')) + + +def pre_check() -> bool: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +def pre_process(mode : ProcessMode) -> bool: + if not has_image(state_manager.get_item('source_paths')): + logger.error(wording.get('choose_image_source') + wording.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(wording.get('no_source_face_detected') + wording.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(wording.get('choose_image_or_video_target') + wording.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: + clear_inference_pool() + get_static_model_initializer.cache_clear() + 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) + crop_masks = [] + temp_vision_frames = [] + + if 'box' in state_manager.get_item('face_mask_types'): + box_mask = create_static_box_mask(crop_vision_frame.shape[:2][::-1], 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, 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 '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) + temp_vision_frame = paste_back(temp_vision_frame, crop_vision_frame, crop_mask, affine_matrix) + return temp_vision_frame + + +def forward_swap_face(source_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 = {} + + for face_swapper_input in face_swapper.get_inputs(): + if face_swapper_input.name == 'source': + if model_type == 'blendswap' or model_type == 'uniface': + face_swapper_inputs[face_swapper_input.name] = prepare_source_frame(source_face) + else: + face_swapper_inputs[face_swapper_input.name] = prepare_source_embedding(source_face) + 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(embedding : Embedding) -> Embedding: + embedding_converter = get_inference_pool().get('embedding_converter') + + with conditional_thread_semaphore(): + embedding = embedding_converter.run(None, + { + 'input': embedding + })[0] + + return 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, _ = convert_embedding(source_face) + source_embedding = source_embedding.reshape(1, -1) + elif 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) + else: + _, source_normed_embedding = convert_embedding(source_face) + source_embedding = source_normed_embedding.reshape(1, -1) + return source_embedding + + +def convert_embedding(source_face : Face) -> Tuple[Embedding, Embedding]: + embedding = source_face.embedding.reshape(-1, 512) + embedding = forward_convert_embedding(embedding) + embedding = embedding.ravel() + normed_embedding = embedding / numpy.linalg.norm(embedding) + return embedding, normed_embedding + + +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 == 'ghost' or model_type == '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 get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + return swap_face(source_face, target_face, temp_vision_frame) + + +def process_frame(inputs : FaceSwapperInputs) -> VisionFrame: + reference_faces = inputs.get('reference_faces') + source_face = inputs.get('source_face') + target_vision_frame = inputs.get('target_vision_frame') + many_faces = sort_and_filter_faces(get_many_faces([ target_vision_frame ])) + + if state_manager.get_item('face_selector_mode') == 'many': + if many_faces: + for target_face in many_faces: + target_vision_frame = swap_face(source_face, target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'one': + target_face = get_one_face(many_faces) + if target_face: + target_vision_frame = swap_face(source_face, target_face, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'reference': + similar_faces = find_similar_faces(many_faces, reference_faces, state_manager.get_item('reference_face_distance')) + if similar_faces: + for similar_face in similar_faces: + target_vision_frame = swap_face(source_face, similar_face, target_vision_frame) + return target_vision_frame + + +def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + source_frames = read_static_images(source_paths) + source_faces = get_many_faces(source_frames) + source_face = get_average_face(source_faces) + + for queue_payload in process_manager.manage(queue_payloads): + target_vision_path = queue_payload['frame_path'] + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'source_face': source_face, + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + source_frames = read_static_images(source_paths) + source_faces = get_many_faces(source_frames) + source_face = get_average_face(source_faces) + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'source_face': source_face, + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + processors.multi_process_frames(source_paths, temp_frame_paths, process_frames) diff --git a/facefusion/processors/modules/frame_colorizer.py b/facefusion/processors/modules/frame_colorizer.py new file mode 100644 index 00000000..43f6b3d7 --- /dev/null +++ b/facefusion/processors/modules/frame_colorizer.py @@ -0,0 +1,283 @@ +from argparse import ArgumentParser +from typing import List + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, inference_manager, logger, process_manager, state_manager, wording +from facefusion.common_helper import create_int_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors import choices as processors_choices +from facefusion.processors.typing import FrameColorizerInputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import thread_semaphore +from facefusion.typing import ApplyStateItem, Args, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import read_image, read_static_image, unpack_resolution, write_image + +MODEL_SET : ModelSet =\ +{ + 'ddcolor': + { + 'hashes': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ddcolor.hash', + 'path': resolve_relative_path('../.assets/models/ddcolor.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ddcolor.onnx', + 'path': resolve_relative_path('../.assets/models/ddcolor.onnx') + } + }, + 'type': 'ddcolor' + }, + 'ddcolor_artistic': + { + 'hashes': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ddcolor_artistic.hash', + 'path': resolve_relative_path('../.assets/models/ddcolor_artistic.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ddcolor_artistic.onnx', + 'path': resolve_relative_path('../.assets/models/ddcolor_artistic.onnx') + } + }, + 'type': 'ddcolor' + }, + 'deoldify': + { + 'hashes': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/deoldify.hash', + 'path': resolve_relative_path('../.assets/models/deoldify.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/deoldify.onnx', + 'path': resolve_relative_path('../.assets/models/deoldify.onnx') + } + }, + 'type': 'deoldify' + }, + 'deoldify_artistic': + { + 'hashes': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/deoldify_artistic.hash', + 'path': resolve_relative_path('../.assets/models/deoldify_artistic.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/deoldify_artistic.onnx', + 'path': resolve_relative_path('../.assets/models/deoldify_artistic.onnx') + } + }, + 'type': 'deoldify' + }, + 'deoldify_stable': + { + 'hashes': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/deoldify_stable.hash', + 'path': resolve_relative_path('../.assets/models/deoldify_stable.hash') + } + }, + 'sources': + { + 'frame_colorizer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/deoldify_stable.onnx', + 'path': resolve_relative_path('../.assets/models/deoldify_stable.onnx') + } + }, + 'type': 'deoldify' + } +} + + +def get_inference_pool() -> InferencePool: + model_sources = get_model_options().get('sources') + model_context = __name__ + '.' + state_manager.get_item('frame_colorizer_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('frame_colorizer_model') + inference_manager.clear_inference_pool(model_context) + + +def get_model_options() -> ModelOptions: + frame_colorizer_model = state_manager.get_item('frame_colorizer_model') + return MODEL_SET.get(frame_colorizer_model) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--frame-colorizer-model', help = wording.get('help.frame_colorizer_model'), default = config.get_str_value('processors.frame_colorizer_model', 'ddcolor'), choices = processors_choices.frame_colorizer_models) + group_processors.add_argument('--frame-colorizer-blend', help = wording.get('help.frame_colorizer_blend'), type = int, default = config.get_int_value('processors.frame_colorizer_blend', '100'), choices = processors_choices.frame_colorizer_blend_range, metavar = create_int_metavar(processors_choices.frame_colorizer_blend_range)) + group_processors.add_argument('--frame-colorizer-size', help = wording.get('help.frame_colorizer_size'), type = str, default = config.get_str_value('processors.frame_colorizer_size', '256x256'), choices = processors_choices.frame_colorizer_sizes) + 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: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +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(wording.get('choose_image_or_video_target') + wording.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + 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_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_frame(temp_vision_frame : VisionFrame, paste_vision_frame : VisionFrame) -> VisionFrame: + frame_colorizer_blend = 1 - (state_manager.get_item('frame_colorizer_blend') / 100) + temp_vision_frame = cv2.addWeighted(temp_vision_frame, frame_colorizer_blend, paste_vision_frame, 1 - frame_colorizer_blend, 0) + return temp_vision_frame + + +def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + pass + + +def process_frame(inputs : FrameColorizerInputs) -> VisionFrame: + target_vision_frame = inputs.get('target_vision_frame') + return colorize_frame(target_vision_frame) + + +def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + for queue_payload in process_manager.manage(queue_payloads): + target_vision_path = queue_payload['frame_path'] + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/modules/frame_enhancer.py b/facefusion/processors/modules/frame_enhancer.py new file mode 100644 index 00000000..8e9c9df2 --- /dev/null +++ b/facefusion/processors/modules/frame_enhancer.py @@ -0,0 +1,415 @@ +from argparse import ArgumentParser +from typing import List + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, inference_manager, logger, process_manager, state_manager, wording +from facefusion.common_helper import create_int_metavar +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors import choices as processors_choices +from facefusion.processors.typing import FrameEnhancerInputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.typing import ApplyStateItem, Args, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import create_tile_frames, merge_tile_frames, read_image, read_static_image, write_image + +MODEL_SET : ModelSet =\ +{ + 'clear_reality_x4': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/clear_reality_x4.hash', + 'path': resolve_relative_path('../.assets/models/clear_reality_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/clear_reality_x4.onnx', + 'path': resolve_relative_path('../.assets/models/clear_reality_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'lsdir_x4': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/lsdir_x4.hash', + 'path': resolve_relative_path('../.assets/models/lsdir_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/lsdir_x4.onnx', + 'path': resolve_relative_path('../.assets/models/lsdir_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'nomos8k_sc_x4': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/nomos8k_sc_x4.hash', + 'path': resolve_relative_path('../.assets/models/nomos8k_sc_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/real_esrgan_x2.hash', + 'path': resolve_relative_path('../.assets/models/real_esrgan_x2.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/real_esrgan_x2_fp16.hash', + 'path': resolve_relative_path('../.assets/models/real_esrgan_x2_fp16.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/real_esrgan_x4.hash', + 'path': resolve_relative_path('../.assets/models/real_esrgan_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/real_esrgan_x4_fp16.hash', + 'path': resolve_relative_path('../.assets/models/real_esrgan_x4_fp16.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/real_esrgan_x8.hash', + 'path': resolve_relative_path('../.assets/models/real_esrgan_x8.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/real_esrgan_x8_fp16.hash', + 'path': resolve_relative_path('../.assets/models/real_esrgan_x8_fp16.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/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': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/real_hatgan_x4.hash', + 'path': resolve_relative_path('../.assets/models/real_hatgan_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/real_hatgan_x4.onnx', + 'path': resolve_relative_path('../.assets/models/real_hatgan_x4.onnx') + } + }, + 'size': (256, 16, 8), + 'scale': 4 + }, + 'span_kendata_x4': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/span_kendata_x4.hash', + 'path': resolve_relative_path('../.assets/models/span_kendata_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/span_kendata_x4.onnx', + 'path': resolve_relative_path('../.assets/models/span_kendata_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + }, + 'ultra_sharp_x4': + { + 'hashes': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ultra_sharp_x4.hash', + 'path': resolve_relative_path('../.assets/models/ultra_sharp_x4.hash') + } + }, + 'sources': + { + 'frame_enhancer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/ultra_sharp_x4.onnx', + 'path': resolve_relative_path('../.assets/models/ultra_sharp_x4.onnx') + } + }, + 'size': (128, 8, 4), + 'scale': 4 + } +} + + +def get_inference_pool() -> InferencePool: + model_sources = get_model_options().get('sources') + model_context = __name__ + '.' + state_manager.get_item('frame_enhancer_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('frame_enhancer_model') + inference_manager.clear_inference_pool(model_context) + + +def get_model_options() -> ModelOptions: + frame_enhancer_model = state_manager.get_item('frame_enhancer_model') + return MODEL_SET.get(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 = wording.get('help.frame_enhancer_model'), default = config.get_str_value('processors.frame_enhancer_model', 'span_kendata_x4'), choices = processors_choices.frame_enhancer_models) + group_processors.add_argument('--frame-enhancer-blend', help = wording.get('help.frame_enhancer_blend'), type = int, default = config.get_int_value('processors.frame_enhancer_blend', '80'), choices = processors_choices.frame_enhancer_blend_range, metavar = create_int_metavar(processors_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: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +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(wording.get('choose_image_or_video_target') + wording.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + 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_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(vision_tile_frame : VisionFrame) -> VisionFrame: + vision_tile_frame = numpy.expand_dims(vision_tile_frame[:, :, ::-1], axis = 0) + vision_tile_frame = vision_tile_frame.transpose(0, 3, 1, 2) + vision_tile_frame = vision_tile_frame.astype(numpy.float32) / 255 + return vision_tile_frame + + +def normalize_tile_frame(vision_tile_frame : VisionFrame) -> VisionFrame: + vision_tile_frame = vision_tile_frame.transpose(0, 2, 3, 1).squeeze(0) * 255 + vision_tile_frame = vision_tile_frame.clip(0, 255).astype(numpy.uint8)[:, :, ::-1] + return vision_tile_frame + + +def blend_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 = cv2.addWeighted(temp_vision_frame, frame_enhancer_blend, merge_vision_frame, 1 - frame_enhancer_blend, 0) + return temp_vision_frame + + +def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + pass + + +def process_frame(inputs : FrameEnhancerInputs) -> VisionFrame: + target_vision_frame = inputs.get('target_vision_frame') + return enhance_frame(target_vision_frame) + + +def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + for queue_payload in process_manager.manage(queue_payloads): + target_vision_path = queue_payload['frame_path'] + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/modules/lip_syncer.py b/facefusion/processors/modules/lip_syncer.py new file mode 100755 index 00000000..53f2a380 --- /dev/null +++ b/facefusion/processors/modules/lip_syncer.py @@ -0,0 +1,270 @@ +from argparse import ArgumentParser +from typing import List + +import cv2 +import numpy + +import facefusion.jobs.job_manager +import facefusion.jobs.job_store +import facefusion.processors.core as processors +from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, process_manager, state_manager, voice_extractor, wording +from facefusion.audio import create_empty_audio_frame, get_voice_frame, read_static_voice +from facefusion.common_helper import get_first +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.face_analyser import get_many_faces, get_one_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_mouth_mask, create_occlusion_mask, create_static_box_mask +from facefusion.face_selector import find_similar_faces, sort_and_filter_faces +from facefusion.face_store import get_reference_faces +from facefusion.filesystem import filter_audio_paths, has_audio, in_directory, is_image, is_video, resolve_relative_path, same_file_extension +from facefusion.processors import choices as processors_choices +from facefusion.processors.typing import LipSyncerInputs +from facefusion.program_helper import find_argument_group +from facefusion.thread_helper import conditional_thread_semaphore +from facefusion.typing import ApplyStateItem, Args, AudioFrame, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, QueuePayload, UpdateProgress, VisionFrame +from facefusion.vision import read_image, read_static_image, restrict_video_fps, write_image + +MODEL_SET : ModelSet =\ +{ + 'wav2lip': + { + 'hashes': + { + 'lip_syncer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/wav2lip.hash', + 'path': resolve_relative_path('../.assets/models/wav2lip.hash') + } + }, + 'sources': + { + 'lip_syncer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/wav2lip.onnx', + 'path': resolve_relative_path('../.assets/models/wav2lip.onnx') + } + }, + 'size': (96, 96) + }, + 'wav2lip_gan': + { + 'hashes': + { + 'lip_syncer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/wav2lip_gan.hash', + 'path': resolve_relative_path('../.assets/models/wav2lip_gan.hash') + } + }, + 'sources': + { + 'lip_syncer': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/wav2lip_gan.onnx', + 'path': resolve_relative_path('../.assets/models/wav2lip_gan.onnx') + } + }, + 'size': (96, 96) + } +} + + +def get_inference_pool() -> InferencePool: + model_sources = get_model_options().get('sources') + model_context = __name__ + '.' + state_manager.get_item('lip_syncer_model') + return inference_manager.get_inference_pool(model_context, model_sources) + + +def clear_inference_pool() -> None: + model_context = __name__ + '.' + state_manager.get_item('lip_syncer_model') + inference_manager.clear_inference_pool(model_context) + + +def get_model_options() -> ModelOptions: + lip_syncer_model = state_manager.get_item('lip_syncer_model') + return MODEL_SET.get(lip_syncer_model) + + +def register_args(program : ArgumentParser) -> None: + group_processors = find_argument_group(program, 'processors') + if group_processors: + group_processors.add_argument('--lip-syncer-model', help = wording.get('help.lip_syncer_model'), default = config.get_str_value('processors.lip_syncer_model', 'wav2lip_gan'), choices = processors_choices.lip_syncer_models) + facefusion.jobs.job_store.register_step_keys([ 'lip_syncer_model' ]) + + +def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: + apply_state_item('lip_syncer_model', args.get('lip_syncer_model')) + + +def pre_check() -> bool: + download_directory_path = resolve_relative_path('../.assets/models') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') + + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) + + +def pre_process(mode : ProcessMode) -> bool: + if not has_audio(state_manager.get_item('source_paths')): + logger.error(wording.get('choose_audio_source') + wording.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(wording.get('choose_image_or_video_target') + wording.get('exclamation_mark'), __name__) + return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(wording.get('specify_image_or_video_output') + wording.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(wording.get('match_target_and_output_extension') + wording.get('exclamation_mark'), __name__) + return False + return True + + +def post_process() -> None: + read_static_image.cache_clear() + read_static_voice.cache_clear() + 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, temp_audio_frame : AudioFrame, temp_vision_frame : VisionFrame) -> VisionFrame: + model_size = get_model_options().get('size') + temp_audio_frame = prepare_audio_frame(temp_audio_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)) + face_landmark_68 = cv2.transform(target_face.landmark_set.get('68').reshape(1, -1, 2), affine_matrix).reshape(-1, 2) + bounding_box = create_bounding_box(face_landmark_68) + bounding_box[1] -= numpy.abs(bounding_box[3] - bounding_box[1]) * 0.125 + mouth_mask = create_mouth_mask(face_landmark_68) + box_mask = create_static_box_mask(crop_vision_frame.shape[:2][::-1], state_manager.get_item('face_mask_blur'), state_manager.get_item('face_mask_padding')) + crop_masks =\ + [ + mouth_mask, + 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) + + close_vision_frame, close_matrix = warp_face_by_bounding_box(crop_vision_frame, bounding_box, model_size) + close_vision_frame = prepare_crop_frame(close_vision_frame) + close_vision_frame = forward(temp_audio_frame, close_vision_frame) + close_vision_frame = normalize_close_frame(close_vision_frame) + crop_vision_frame = cv2.warpAffine(close_vision_frame, cv2.invertAffineTransform(close_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(temp_audio_frame : AudioFrame, close_vision_frame : VisionFrame) -> VisionFrame: + lip_syncer = get_inference_pool().get('lip_syncer') + + with conditional_thread_semaphore(): + close_vision_frame = lip_syncer.run(None, + { + 'source': temp_audio_frame, + 'target': close_vision_frame + })[0] + + return close_vision_frame + + +def prepare_audio_frame(temp_audio_frame : AudioFrame) -> AudioFrame: + 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) + 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: + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0) + prepare_vision_frame = crop_vision_frame.copy() + prepare_vision_frame[:, 48:] = 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('float32') / 255.0 + return crop_vision_frame + + +def normalize_close_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + 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) + return crop_vision_frame + + +def get_reference_frame(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + pass + + +def process_frame(inputs : LipSyncerInputs) -> VisionFrame: + reference_faces = inputs.get('reference_faces') + source_audio_frame = inputs.get('source_audio_frame') + target_vision_frame = inputs.get('target_vision_frame') + many_faces = sort_and_filter_faces(get_many_faces([ target_vision_frame ])) + + if state_manager.get_item('face_selector_mode') == 'many': + if many_faces: + for target_face in many_faces: + target_vision_frame = sync_lip(target_face, source_audio_frame, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'one': + target_face = get_one_face(many_faces) + if target_face: + target_vision_frame = sync_lip(target_face, source_audio_frame, target_vision_frame) + if state_manager.get_item('face_selector_mode') == 'reference': + similar_faces = find_similar_faces(many_faces, reference_faces, state_manager.get_item('reference_face_distance')) + if similar_faces: + for similar_face in similar_faces: + target_vision_frame = sync_lip(similar_face, source_audio_frame, target_vision_frame) + return target_vision_frame + + +def process_frames(source_paths : List[str], queue_payloads : List[QueuePayload], update_progress : UpdateProgress) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + source_audio_path = get_first(filter_audio_paths(source_paths)) + temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + + for queue_payload in process_manager.manage(queue_payloads): + frame_number = queue_payload.get('frame_number') + target_vision_path = queue_payload.get('frame_path') + source_audio_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() + target_vision_frame = read_image(target_vision_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'source_audio_frame': source_audio_frame, + 'target_vision_frame': target_vision_frame + }) + write_image(target_vision_path, output_vision_frame) + update_progress(1) + + +def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + source_audio_frame = create_empty_audio_frame() + target_vision_frame = read_static_image(target_path) + output_vision_frame = process_frame( + { + 'reference_faces': reference_faces, + 'source_audio_frame': source_audio_frame, + 'target_vision_frame': target_vision_frame + }) + write_image(output_path, output_vision_frame) + + +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + source_audio_paths = 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')) + for source_audio_path in source_audio_paths: + read_static_voice(source_audio_path, temp_video_fps) + processors.multi_process_frames(source_paths, temp_frame_paths, process_frames) diff --git a/facefusion/processors/pixel_boost.py b/facefusion/processors/pixel_boost.py new file mode 100644 index 00000000..13665c01 --- /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.typing 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, axis = 0).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/typing.py b/facefusion/processors/typing.py new file mode 100644 index 00000000..ae57587e --- /dev/null +++ b/facefusion/processors/typing.py @@ -0,0 +1,125 @@ +from typing import Any, Dict, List, Literal, TypedDict + +from numpy._typing import NDArray + +from facefusion.typing import AppContext, AudioFrame, Face, FaceSet, VisionFrame + +AgeModifierModel = Literal['styleganex_age'] +ExpressionRestorerModel = Literal['live_portrait'] +FaceDebuggerItem = Literal['bounding-box', 'face-landmark-5', 'face-landmark-5/68', 'face-landmark-68', 'face-landmark-68/5', 'face-mask', 'face-detector-score', 'face-landmarker-score', 'age', 'gender', 'race'] +FaceEditorModel = Literal['live_portrait'] +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'] +FaceSwapperModel = Literal['blendswap_256', 'ghost_256_unet_1', 'ghost_256_unet_2', 'ghost_256_unet_3', 'inswapper_128', 'inswapper_128_fp16', 'simswap_256', 'simswap_512_unofficial', 'uniface_256'] +FrameColorizerModel = Literal['ddcolor', 'ddcolor_artistic', 'deoldify', 'deoldify_artistic', 'deoldify_stable'] +FrameEnhancerModel = Literal['clear_reality_x4', 'lsdir_x4', 'nomos8k_sc_x4', 'real_esrgan_x2', 'real_esrgan_x2_fp16', 'real_esrgan_x4', 'real_esrgan_x4_fp16', 'real_hatgan_x4', 'real_esrgan_x8', 'real_esrgan_x8_fp16', 'span_kendata_x4', 'ultra_sharp_x4'] +LipSyncerModel = Literal['wav2lip', 'wav2lip_gan'] + +FaceSwapperSet = Dict[FaceSwapperModel, List[str]] + +AgeModifierInputs = TypedDict('AgeModifierInputs', +{ + 'reference_faces' : FaceSet, + 'target_vision_frame' : VisionFrame +}) +ExpressionRestorerInputs = TypedDict('ExpressionRestorerInputs', +{ + 'reference_faces' : FaceSet, + 'source_vision_frame' : VisionFrame, + 'target_vision_frame' : VisionFrame +}) +FaceDebuggerInputs = TypedDict('FaceDebuggerInputs', +{ + 'reference_faces' : FaceSet, + 'target_vision_frame' : VisionFrame +}) +FaceEditorInputs = TypedDict('FaceEditorInputs', +{ + 'reference_faces' : FaceSet, + 'target_vision_frame' : VisionFrame +}) +FaceEnhancerInputs = TypedDict('FaceEnhancerInputs', +{ + 'reference_faces' : FaceSet, + 'target_vision_frame' : VisionFrame +}) +FaceSwapperInputs = TypedDict('FaceSwapperInputs', +{ + 'reference_faces' : FaceSet, + 'source_face' : Face, + 'target_vision_frame' : VisionFrame +}) +FrameColorizerInputs = TypedDict('FrameColorizerInputs', +{ + 'target_vision_frame' : VisionFrame +}) +FrameEnhancerInputs = TypedDict('FrameEnhancerInputs', +{ + 'target_vision_frame' : VisionFrame +}) +LipSyncerInputs = TypedDict('LipSyncerInputs', +{ + 'reference_faces' : FaceSet, + 'source_audio_frame' : AudioFrame, + 'target_vision_frame' : VisionFrame +}) + +ProcessorStateKey = Literal\ +[ + 'age_modifier_model', + 'age_modifier_direction', + 'expression_restorer_model', + 'expression_restorer_factor', + '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_swapper_model', + 'face_swapper_pixel_boost', + 'frame_colorizer_model', + 'frame_colorizer_blend', + 'frame_colorizer_size', + 'frame_enhancer_model', + 'frame_enhancer_blend', + 'lip_syncer_model' +] +ProcessorState = TypedDict('ProcessorState', +{ + 'age_modifier_model': AgeModifierModel, + 'age_modifier_direction': int, + 'face_debugger_items' : List[FaceDebuggerItem], + 'face_enhancer_model' : FaceEnhancerModel, + 'face_enhancer_blend' : int, + 'face_swapper_model' : FaceSwapperModel, + 'face_swapper_pixel_boost' : str, + 'frame_colorizer_model' : FrameColorizerModel, + 'frame_colorizer_blend' : int, + 'frame_colorizer_size' : str, + 'frame_enhancer_model' : FrameEnhancerModel, + 'frame_enhancer_blend' : int, + 'lip_syncer_model' : LipSyncerModel +}) +ProcessorStateSet = Dict[AppContext, ProcessorState] + +LivePortraitPitch = float +LivePortraitYaw = float +LivePortraitRoll = float +LivePortraitExpression = NDArray[Any] +LivePortraitFeatureVolume = NDArray[Any] +LivePortraitMotionPoints = NDArray[Any] +LivePortraitRotation = NDArray[Any] +LivePortraitScale = NDArray[Any] +LivePortraitTranslation = NDArray[Any] diff --git a/facefusion/program.py b/facefusion/program.py new file mode 100755 index 00000000..5e0a3761 --- /dev/null +++ b/facefusion/program.py @@ -0,0 +1,234 @@ +from argparse import ArgumentParser, HelpFormatter + +import facefusion.choices +from facefusion import config, metadata, state_manager, wording +from facefusion.common_helper import create_float_metavar, create_int_metavar +from facefusion.execution import get_execution_provider_choices +from facefusion.filesystem import list_directory +from facefusion.jobs import job_store +from facefusion.processors.core import get_processors_modules +from facefusion.program_helper import suggest_face_detector_choices + + +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_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + program.add_argument('-c', '--config-path', help = wording.get('help.config_path'), default = 'facefusion.ini') + job_store.register_job_keys([ 'config-path' ]) + apply_config_path(program) + return program + + +def create_jobs_path_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + program.add_argument('-j', '--jobs-path', help = wording.get('help.jobs_path'), default = config.get_str_value('paths.jobs_path', '.jobs')) + job_store.register_job_keys([ 'jobs_path' ]) + return program + + +def create_paths_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + program.add_argument('-s', '--source-paths', help = wording.get('help.source_paths'), action = 'append', default = config.get_str_list('paths.source_paths')) + program.add_argument('-t', '--target-path', help = wording.get('help.target_path'), default = config.get_str_value('paths.target_path')) + program.add_argument('-o', '--output-path', help = wording.get('help.output_path'), default = config.get_str_value('paths.output_path')) + job_store.register_step_keys([ 'source_paths', 'target_path', 'output_path' ]) + 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 = wording.get('help.face_detector_model'), default = config.get_str_value('face_detector.face_detector_model', 'yoloface'), choices = facefusion.choices.face_detector_set.keys()) + group_face_detector.add_argument('--face-detector-size', help = wording.get('help.face_detector_size'), default = config.get_str_value('face_detector.face_detector_size', '640x640'), choices = suggest_face_detector_choices(program)) + group_face_detector.add_argument('--face-detector-angles', help = wording.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 = wording.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_angles', 'face_detector_size', '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 = wording.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 = wording.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 = wording.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 = wording.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-gender', help = wording.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 = wording.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('--face-selector-age-start', help = wording.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 = wording.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('--reference-face-position', help = wording.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 = wording.get('help.reference_face_distance'), type = float, default = config.get_float_value('face_selector.reference_face_distance', '0.6'), 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 = wording.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-mask-types', help = wording.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-blur', help = wording.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 = wording.get('help.face_mask_padding'), type = int, default = config.get_int_list('face_masker.face_mask_padding', '0 0 0 0'), nargs = '+') + group_face_masker.add_argument('--face-mask-regions', help = wording.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') + job_store.register_step_keys([ 'face_mask_types', 'face_mask_blur', 'face_mask_padding', 'face_mask_regions' ]) + 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 = wording.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 = wording.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 = wording.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 = wording.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) + group_output_creation = program.add_argument_group('output creation') + group_output_creation.add_argument('--output-image-quality', help = wording.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-resolution', help = wording.get('help.output_image_resolution'), default = config.get_str_value('output_creation.output_image_resolution')) + group_output_creation.add_argument('--output-audio-encoder', help = wording.get('help.output_audio_encoder'), default = config.get_str_value('output_creation.output_audio_encoder', 'aac'), choices = facefusion.choices.output_audio_encoders) + group_output_creation.add_argument('--output-video-encoder', help = wording.get('help.output_video_encoder'), default = config.get_str_value('output_creation.output_video_encoder', 'libx264'), choices = facefusion.choices.output_video_encoders) + group_output_creation.add_argument('--output-video-preset', help = wording.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 = wording.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-resolution', help = wording.get('help.output_video_resolution'), default = config.get_str_value('output_creation.output_video_resolution')) + group_output_creation.add_argument('--output-video-fps', help = wording.get('help.output_video_fps'), type = float, default = config.get_str_value('output_creation.output_video_fps')) + group_output_creation.add_argument('--skip-audio', help = wording.get('help.skip_audio'), action = 'store_true', default = config.get_bool_value('output_creation.skip_audio')) + job_store.register_step_keys([ 'output_image_quality', 'output_image_resolution', 'output_audio_encoder', 'output_video_encoder', 'output_video_preset', 'output_video_quality', 'output_video_resolution', 'output_video_fps', 'skip_audio' ]) + return program + + +def create_processors_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + available_processors = list_directory('facefusion/processors/modules') + group_processors = program.add_argument_group('processors') + group_processors.add_argument('--processors', help = wording.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 = list_directory('facefusion/uis/layouts') + group_uis = program.add_argument_group('uis') + group_uis.add_argument('--open-browser', help = wording.get('help.open_browser'), action = 'store_true', default = config.get_bool_value('uis.open_browser')) + group_uis.add_argument('--ui-layouts', help = wording.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 = wording.get('help.ui_workflow'), default = config.get_str_value('uis.ui_workflow', 'instant_runner'), choices = facefusion.choices.ui_workflows) + return program + + +def create_execution_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + execution_providers = get_execution_provider_choices() + group_execution = program.add_argument_group('execution') + group_execution.add_argument('--execution-device-id', help = wording.get('help.execution_device_id'), default = config.get_str_value('execution.execution_device_id', '0')) + group_execution.add_argument('--execution-providers', help = wording.get('help.execution_providers').format(choices = ', '.join(execution_providers)), default = config.get_str_list('execution.execution_providers', 'cpu'), choices = execution_providers, nargs = '+', metavar = 'EXECUTION_PROVIDERS') + group_execution.add_argument('--execution-thread-count', help = wording.get('help.execution_thread_count'), type = int, default = config.get_int_value('execution.execution_thread_count', '4'), choices = facefusion.choices.execution_thread_count_range, metavar = create_int_metavar(facefusion.choices.execution_thread_count_range)) + group_execution.add_argument('--execution-queue-count', help = wording.get('help.execution_queue_count'), type = int, default = config.get_int_value('execution.execution_queue_count', '1'), choices = facefusion.choices.execution_queue_count_range, metavar = create_int_metavar(facefusion.choices.execution_queue_count_range)) + job_store.register_job_keys([ 'execution_device_id', 'execution_providers', 'execution_thread_count', 'execution_queue_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 = wording.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 = wording.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_skip_download_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_misc = program.add_argument_group('misc') + group_misc.add_argument('--skip-download', help = wording.get('help.skip_download'), action = 'store_true', default = config.get_bool_value('misc.skip_download')) + job_store.register_job_keys([ 'skip_download' ]) + 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 = wording.get('help.log_level'), default = config.get_str_value('misc.log_level', 'info'), choices = facefusion.choices.log_level_set.keys()) + job_store.register_job_keys([ 'log_level' ]) + return program + + +def create_job_id_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + program.add_argument('job_id', help = wording.get('help.job_id')) + job_store.register_job_keys([ 'job_id' ]) + return program + + +def create_job_status_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + program.add_argument('job_status', help = wording.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 = wording.get('help.step_index'), type = int) + return program + + +def collect_step_program() -> ArgumentParser: + return ArgumentParser(parents= [ create_config_program(), create_jobs_path_program(), create_paths_program(), create_face_detector_program(), create_face_landmarker_program(), create_face_selector_program(), create_face_masker_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_memory_program(), create_skip_download_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 = wording.get('help.run'), parents = [ collect_step_program(), create_uis_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('headless-run', help = wording.get('help.headless_run'), parents = [ collect_step_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('force-download', help = wording.get('help.force_download'), parents = [ create_log_level_program() ], formatter_class = create_help_formatter_large) + # job manager + sub_program.add_parser('job-create', help = wording.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 = wording.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 = wording.get('help.job_submit_all'), parents = [ create_jobs_path_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-delete', help = wording.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 = wording.get('help.job_delete_all'), parents = [ create_jobs_path_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-list', help = wording.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-add-step', help = wording.get('help.job_add_step'), parents = [ create_job_id_program(), collect_step_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-remix-step', help = wording.get('help.job_remix_step'), parents = [ create_job_id_program(), create_step_index_program(), collect_step_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-insert-step', help = wording.get('help.job_insert_step'), parents = [ create_job_id_program(), create_step_index_program(), collect_step_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-remove-step', help = wording.get('help.job_remove_step'), parents = [ create_job_id_program(), create_step_index_program(), collect_step_program(), create_log_level_program() ], formatter_class = create_help_formatter_large) + # job runner + sub_program.add_parser('job-run', help = wording.get('help.job_run'), parents = [ create_job_id_program(), create_config_program(), create_jobs_path_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-run-all', help = wording.get('help.job_run_all'), parents = [ create_config_program(), create_jobs_path_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-retry', help = wording.get('help.job_retry'), parents = [ create_job_id_program(), create_config_program(), create_jobs_path_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + sub_program.add_parser('job-retry-all', help = wording.get('help.job_retry_all'), parents = [ create_config_program(), create_jobs_path_program(), collect_job_program() ], formatter_class = create_help_formatter_large) + return ArgumentParser(parents = [ program ], formatter_class = create_help_formatter_small, add_help = True) + + +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 00000000..28648fb2 --- /dev/null +++ b/facefusion/program_helper.py @@ -0,0 +1,45 @@ +from argparse import ArgumentParser, _ArgumentGroup, _SubParsersAction +from typing import List, Optional + +import facefusion.choices +from facefusion.processors import choices as processors_choices + + +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 not validate_actions(program): + return False + + 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 + + +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 + + +def suggest_face_detector_choices(program : ArgumentParser) -> List[str]: + known_args, _ = program.parse_known_args() + return facefusion.choices.face_detector_set.get(known_args.face_detector_model) #type:ignore[call-overload] + + +def suggest_face_swapper_pixel_boost_choices(program : ArgumentParser) -> List[str]: + known_args, _ = program.parse_known_args() + return processors_choices.face_swapper_set.get(known_args.face_swapper_model) #type:ignore[call-overload] diff --git a/facefusion/state_manager.py b/facefusion/state_manager.py new file mode 100644 index 00000000..d0e5f1a1 --- /dev/null +++ b/facefusion/state_manager.py @@ -0,0 +1,38 @@ +from typing import Any, Union + +from facefusion.app_context import detect_app_context +from facefusion.processors.typing import ProcessorState, ProcessorStateKey +from facefusion.typing import State, StateKey, StateSet + +STATES : Union[StateSet, ProcessorState] =\ +{ + 'cli': {}, #type:ignore[typeddict-item] + 'ui': {} #type:ignore[typeddict-item] +} + + +def get_state() -> Union[State, ProcessorState]: + app_context = detect_app_context() + return STATES.get(app_context) #type:ignore + + +def init_item(key : Union[StateKey, ProcessorStateKey], value : Any) -> None: + STATES['cli'][key] = value #type:ignore + STATES['ui'][key] = value #type:ignore + + +def get_item(key : Union[StateKey, ProcessorStateKey]) -> Any: + return get_state().get(key) #type:ignore + + +def set_item(key : Union[StateKey, ProcessorStateKey], value : Any) -> None: + app_context = detect_app_context() + STATES[app_context][key] = value #type:ignore + + +def sync_item(key : Union[StateKey, ProcessorStateKey]) -> None: + STATES['cli'][key] = STATES.get('ui').get(key) #type:ignore + + +def clear_item(key : Union[StateKey, ProcessorStateKey]) -> None: + set_item(key, None) diff --git a/facefusion/statistics.py b/facefusion/statistics.py index f67c32b9..5f500a0b 100644 --- a/facefusion/statistics.py +++ b/facefusion/statistics.py @@ -1,15 +1,15 @@ from typing import Any, Dict + import numpy -import facefusion.globals -from facefusion.face_store import FACE_STORE +from facefusion import logger, state_manager +from facefusion.face_store import get_face_store from facefusion.typing import FaceSet -from facefusion import logger def create_statistics(static_faces : FaceSet) -> Dict[str, Any]: - face_detector_score_list = [] - face_landmarker_score_list = [] + face_detector_scores = [] + face_landmarker_scores = [] statistics =\ { 'min_face_detector_score': 0, @@ -27,25 +27,25 @@ def create_statistics(static_faces : FaceSet) -> Dict[str, Any]: statistics['total_frames_with_faces'] = statistics.get('total_frames_with_faces') + 1 for face in faces: statistics['total_faces'] = statistics.get('total_faces') + 1 - face_detector_score_list.append(face.scores.get('detector')) - face_landmarker_score_list.append(face.scores.get('landmarker')) - if numpy.array_equal(face.landmarks.get('5'), face.landmarks.get('5/68')): + face_detector_scores.append(face.score_set.get('detector')) + face_landmarker_scores.append(face.score_set.get('landmarker')) + if numpy.array_equal(face.landmark_set.get('5'), face.landmark_set.get('5/68')): statistics['total_face_landmark_5_fallbacks'] = statistics.get('total_face_landmark_5_fallbacks') + 1 - if face_detector_score_list: - statistics['min_face_detector_score'] = round(min(face_detector_score_list), 2) - statistics['max_face_detector_score'] = round(max(face_detector_score_list), 2) - statistics['average_face_detector_score'] = round(numpy.mean(face_detector_score_list), 2) - if face_landmarker_score_list: - statistics['min_face_landmarker_score'] = round(min(face_landmarker_score_list), 2) - statistics['max_face_landmarker_score'] = round(max(face_landmarker_score_list), 2) - statistics['average_face_landmarker_score'] = round(numpy.mean(face_landmarker_score_list), 2) + if face_detector_scores: + statistics['min_face_detector_score'] = round(min(face_detector_scores), 2) + statistics['max_face_detector_score'] = round(max(face_detector_scores), 2) + statistics['average_face_detector_score'] = round(numpy.mean(face_detector_scores), 2) + if face_landmarker_scores: + statistics['min_face_landmarker_score'] = round(min(face_landmarker_scores), 2) + statistics['max_face_landmarker_score'] = round(max(face_landmarker_scores), 2) + statistics['average_face_landmarker_score'] = round(numpy.mean(face_landmarker_scores), 2) return statistics def conditional_log_statistics() -> None: - if facefusion.globals.log_level == 'debug': - statistics = create_statistics(FACE_STORE.get('static_faces')) + if state_manager.get_item('log_level') == 'debug': + statistics = create_statistics(get_face_store().get('static_faces')) for name, value in statistics.items(): - logger.debug(str(name) + ': ' + str(value), __name__.upper()) + logger.debug(str(name) + ': ' + str(value), __name__) diff --git a/facefusion/temp_helper.py b/facefusion/temp_helper.py new file mode 100644 index 00000000..c1798366 --- /dev/null +++ b/facefusion/temp_helper.py @@ -0,0 +1,60 @@ +import glob +import os +import tempfile +from typing import List + +from facefusion import state_manager +from facefusion.filesystem import create_directory, move_file, remove_directory + + +def get_temp_file_path(file_path : str) -> str: + _, temp_file_extension = os.path.splitext(os.path.basename(file_path)) + temp_directory_path = get_temp_directory_path(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 get_temp_frame_paths(target_path : str) -> List[str]: + temp_frames_pattern = get_temp_frames_pattern(target_path, '*') + return sorted(glob.glob(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_base_directory_path() -> str: + return os.path.join(tempfile.gettempdir(), 'facefusion') + + +def create_base_directory() -> bool: + base_directory_path = get_base_directory_path() + return create_directory(base_directory_path) + + +def clear_base_directory() -> bool: + base_directory_path = get_base_directory_path() + return remove_directory(base_directory_path) + + +def get_temp_directory_path(file_path : str) -> str: + temp_file_name, _ = os.path.splitext(os.path.basename(file_path)) + base_directory_path = get_base_directory_path() + return os.path.join(base_directory_path, 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 index c08c6f17..84717f9d 100644 --- a/facefusion/thread_helper.py +++ b/facefusion/thread_helper.py @@ -1,6 +1,8 @@ -from typing import List, Union, ContextManager import threading from contextlib import nullcontext +from typing import ContextManager, Union + +from facefusion.execution import has_execution_provider THREAD_LOCK : threading.Lock = threading.Lock() THREAD_SEMAPHORE : threading.Semaphore = threading.Semaphore() @@ -15,7 +17,7 @@ def thread_semaphore() -> threading.Semaphore: return THREAD_SEMAPHORE -def conditional_thread_semaphore(execution_providers : List[str]) -> Union[threading.Semaphore, ContextManager[None]]: - if 'DmlExecutionProvider' in execution_providers: +def conditional_thread_semaphore() -> Union[threading.Semaphore, ContextManager[None]]: + if has_execution_provider('directml') or has_execution_provider('rocm'): return THREAD_SEMAPHORE return NULL_CONTEXT diff --git a/facefusion/typing.py b/facefusion/typing.py index bc05f801..f8ee1b54 100755 --- a/facefusion/typing.py +++ b/facefusion/typing.py @@ -1,10 +1,20 @@ -from typing import Any, Literal, Callable, List, Tuple, Dict, TypedDict from collections import namedtuple -import numpy +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, TypedDict -BoundingBox = numpy.ndarray[Any, Any] -FaceLandmark5 = numpy.ndarray[Any, Any] -FaceLandmark68 = numpy.ndarray[Any, Any] +import numpy +from numpy.typing import NDArray +from onnxruntime import InferenceSession + +Scale = float +Score = float +Angle = int + +Detection = NDArray[Any] +Prediction = NDArray[Any] + +BoundingBox = NDArray[Any] +FaceLandmark5 = NDArray[Any] +FaceLandmark68 = NDArray[Any] FaceLandmarkSet = TypedDict('FaceLandmarkSet', { '5' : FaceLandmark5, #type:ignore[valid-type] @@ -12,22 +22,26 @@ FaceLandmarkSet = TypedDict('FaceLandmarkSet', '68' : FaceLandmark68, #type:ignore[valid-type] '68/5' : FaceLandmark68 #type:ignore[valid-type] }) -Score = float FaceScoreSet = TypedDict('FaceScoreSet', { 'detector' : Score, 'landmarker' : Score }) -Embedding = numpy.ndarray[Any, Any] +Embedding = NDArray[numpy.float64] +Gender = Literal['female', 'male'] +Age = range +Race = Literal['white', 'black', 'latino', 'asian', 'indian', 'arabic'] Face = namedtuple('Face', [ 'bounding_box', - 'landmarks', - 'scores', + 'score_set', + 'landmark_set', + 'angle', 'embedding', 'normed_embedding', 'gender', - 'age' + 'age', + 'race' ]) FaceSet = Dict[str, List[Face]] FaceStore = TypedDict('FaceStore', @@ -36,20 +50,25 @@ FaceStore = TypedDict('FaceStore', 'reference_faces': FaceSet }) -VisionFrame = numpy.ndarray[Any, Any] -Mask = numpy.ndarray[Any, Any] -Matrix = numpy.ndarray[Any, Any] -Translation = numpy.ndarray[Any, Any] +VisionFrame = NDArray[Any] +Mask = NDArray[Any] +Points = NDArray[Any] +Distance = NDArray[Any] +Matrix = NDArray[Any] +Anchors = NDArray[Any] +Translation = NDArray[Any] AudioBuffer = bytes -Audio = numpy.ndarray[Any, Any] -AudioChunk = numpy.ndarray[Any, Any] -AudioFrame = numpy.ndarray[Any, Any] -Spectrogram = numpy.ndarray[Any, Any] -MelFilterBank = numpy.ndarray[Any, Any] +Audio = NDArray[Any] +AudioChunk = NDArray[Any] +AudioFrame = NDArray[Any] +Spectrogram = NDArray[Any] +Mel = NDArray[Any] +MelFilterBank = NDArray[Any] Fps = float Padding = Tuple[int, int, int, int] +Orientation = Literal['landscape', 'portrait'] Resolution = Tuple[int, int] ProcessState = Literal['checking', 'processing', 'stopping', 'pending'] @@ -58,38 +77,55 @@ QueuePayload = TypedDict('QueuePayload', 'frame_number' : int, 'frame_path' : str }) +Args = Dict[str, Any] UpdateProgress = Callable[[int], None] ProcessFrames = Callable[[List[str], List[QueuePayload], UpdateProgress], None] +ProcessStep = Callable[[str, int, Args], bool] + +Content = Dict[str, Any] WarpTemplate = Literal['arcface_112_v1', 'arcface_112_v2', 'arcface_128_v2', 'ffhq_512'] -WarpTemplateSet = Dict[WarpTemplate, numpy.ndarray[Any, Any]] +WarpTemplateSet = Dict[WarpTemplate, NDArray[Any]] ProcessMode = Literal['output', 'preview', 'stream'] +ErrorCode = Literal[0, 1, 2, 3, 4] LogLevel = Literal['error', 'warn', 'info', 'debug'] +LogLevelSet = Dict[LogLevel, int] + +TableHeaders = List[str] +TableContents = List[List[Any]] + VideoMemoryStrategy = Literal['strict', 'moderate', 'tolerant'] +FaceDetectorModel = Literal['many', 'retinaface', 'scrfd', 'yoloface'] +FaceLandmarkerModel = Literal['many', '2dfan4', 'peppa_wutz'] +FaceDetectorSet = Dict[FaceDetectorModel, List[str]] FaceSelectorMode = Literal['many', 'one', 'reference'] -FaceAnalyserOrder = Literal['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best'] -FaceAnalyserAge = Literal['child', 'teen', 'adult', 'senior'] -FaceAnalyserGender = Literal['female', 'male'] -FaceDetectorModel = Literal['many', 'retinaface', 'scrfd', 'yoloface', 'yunet'] -FaceDetectorTweak = Literal['low-luminance', 'high-luminance'] -FaceRecognizerModel = Literal['arcface_blendswap', 'arcface_inswapper', 'arcface_simswap', 'arcface_uniface'] +FaceSelectorOrder = Literal['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best'] FaceMaskType = Literal['box', 'occlusion', 'region'] FaceMaskRegion = Literal['skin', 'left-eyebrow', 'right-eyebrow', 'left-eye', 'right-eye', 'glasses', 'nose', 'mouth', 'upper-lip', 'lower-lip'] TempFrameFormat = Literal['jpg', 'png', 'bmp'] -OutputVideoEncoder = Literal['libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc', 'h264_amf', 'hevc_amf'] +OutputAudioEncoder = Literal['aac', 'libmp3lame', 'libopus', 'libvorbis'] +OutputVideoEncoder = Literal['libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc', 'h264_amf', 'hevc_amf', 'h264_videotoolbox', 'hevc_videotoolbox'] OutputVideoPreset = Literal['ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow'] -ModelValue = Dict[str, Any] -ModelSet = Dict[str, ModelValue] -OptionsWithModel = TypedDict('OptionsWithModel', -{ - 'model' : ModelValue +Download = TypedDict('Download', + { + 'url' : str, + 'path' : str }) +DownloadSet = Dict[str, Download] + +ModelOptions = Dict[str, Any] +ModelSet = Dict[str, ModelOptions] +ModelInitializer = NDArray[Any] + +ExecutionProviderKey = Literal['cpu', 'coreml', 'cuda', 'directml', 'openvino', 'rocm', 'tensorrt'] +ExecutionProviderValue = Literal['CPUExecutionProvider', 'CoreMLExecutionProvider', 'CUDAExecutionProvider', 'DmlExecutionProvider', 'OpenVINOExecutionProvider', 'ROCMExecutionProvider', 'TensorrtExecutionProvider'] +ExecutionProviderSet = Dict[ExecutionProviderKey, ExecutionProviderValue] ValueAndUnit = TypedDict('ValueAndUnit', { - 'value' : str, + 'value' : int, 'unit' : str }) ExecutionDeviceFramework = TypedDict('ExecutionDeviceFramework', @@ -120,3 +156,147 @@ ExecutionDevice = TypedDict('ExecutionDevice', 'video_memory' : ExecutionDeviceVideoMemory, 'utilization' : ExecutionDeviceUtilization }) + +AppContext = Literal['cli', 'ui'] + +InferencePool = Dict[str, InferenceSession] +InferencePoolSet = 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 = 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 = Dict[str, Job] + +ApplyStateItem = Callable[[Any, Any], None] +StateKey = Literal\ +[ + 'command', + 'config_path', + 'jobs_path', + 'source_paths', + 'target_path', + 'output_path', + 'face_detector_model', + 'face_detector_size', + '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_mask_types', + 'face_mask_blur', + 'face_mask_padding', + 'face_mask_regions', + 'trim_frame_start', + 'trim_frame_end', + 'temp_frame_format', + 'keep_temp', + 'output_image_quality', + 'output_image_resolution', + 'output_audio_encoder', + 'output_video_encoder', + 'output_video_preset', + 'output_video_quality', + 'output_video_resolution', + 'output_video_fps', + 'skip_audio', + 'processors', + 'open_browser', + 'ui_layouts', + 'ui_workflow', + 'execution_device_id', + 'execution_providers', + 'execution_thread_count', + 'execution_queue_count', + 'video_memory_strategy', + 'system_memory_limit', + 'skip_download', + 'log_level', + 'job_id', + 'job_status', + 'step_index' +] +State = TypedDict('State', +{ + 'command' : str, + 'config_path' : str, + 'jobs_path' : str, + 'source_paths' : List[str], + 'target_path' : str, + 'output_path' : str, + 'face_detector_model' : FaceDetectorModel, + 'face_detector_size' : str, + '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_mask_types' : List[FaceMaskType], + 'face_mask_blur' : float, + 'face_mask_padding' : Padding, + 'face_mask_regions' : List[FaceMaskRegion], + 'trim_frame_start' : int, + 'trim_frame_end' : int, + 'temp_frame_format' : TempFrameFormat, + 'keep_temp' : bool, + 'output_image_quality' : int, + 'output_image_resolution' : str, + 'output_audio_encoder' : OutputAudioEncoder, + 'output_video_encoder' : OutputVideoEncoder, + 'output_video_preset' : OutputVideoPreset, + 'output_video_quality' : int, + 'output_video_resolution' : str, + 'output_video_fps' : float, + 'skip_audio' : bool, + 'processors' : List[str], + 'open_browser' : bool, + 'ui_layouts' : List[str], + 'ui_workflow' : UiWorkflow, + 'execution_device_id': str, + 'execution_providers': List[ExecutionProviderKey], + 'execution_thread_count': int, + 'execution_queue_count': int, + 'video_memory_strategy': VideoMemoryStrategy, + 'system_memory_limit': int, + 'skip_download': bool, + 'log_level': LogLevel, + 'job_id': str, + 'job_status': JobStatus, + 'step_index': int +}) +StateSet = Dict[AppContext, State] diff --git a/facefusion/uis/assets/fixes.css b/facefusion/uis/assets/fixes.css deleted file mode 100644 index f65a7cfd..00000000 --- a/facefusion/uis/assets/fixes.css +++ /dev/null @@ -1,7 +0,0 @@ -:root:root:root button:not([class]) -{ - border-radius: 0.375rem; - float: left; - overflow: hidden; - width: 100%; -} diff --git a/facefusion/uis/assets/overrides.css b/facefusion/uis/assets/overrides.css index 744ed3ba..639c9870 100644 --- a/facefusion/uis/assets/overrides.css +++ b/facefusion/uis/assets/overrides.css @@ -1,23 +1,31 @@ -:root:root:root input[type="number"] +:root:root:root:root .gradio-container +{ + max-width: 110em; + overflow: unset; +} + +:root:root:root:root input[type="number"] { max-width: 6rem; } -:root:root:root [type="checkbox"], -:root:root:root [type="radio"] +:root:root:root:root [type="checkbox"], +:root:root:root:root [type="radio"] { border-radius: 50%; height: 1.125rem; width: 1.125rem; } -:root:root:root input[type="range"] +:root:root:root:root input[type="range"], +:root:root:root:root .range-slider div { height: 0.5rem; + border-radius: 0.5rem; } -:root:root:root input[type="range"]::-moz-range-thumb, -:root:root:root input[type="range"]::-webkit-slider-thumb +:root:root:root:root input[type="range"]::-moz-range-thumb, +:root:root:root:root input[type="range"]::-webkit-slider-thumb { background: var(--neutral-300); border: unset; @@ -26,33 +34,63 @@ width: 1.125rem; } -:root:root:root input[type="range"]::-webkit-slider-thumb +:root:root:root:root input[type="range"]::-webkit-slider-thumb { margin-top: 0.375rem; } -:root:root:root .grid-wrap.fixed-height +:root:root:root:root .range-slider input[type="range"]::-webkit-slider-thumb +{ + margin-top: 0.125rem; +} + +:root:root:root:root .range-slider div, +:root:root:root:root .range-slider input[type="range"] +{ + bottom: 50%; + margin-top: -0.25rem; + top: 50%; +} + +:root:root:root:root .grid-wrap.fixed-height { min-height: unset; } -:root:root:root .grid-container -{ - grid-auto-rows: minmax(5em, 1fr); - grid-template-columns: repeat(var(--grid-cols), minmax(5em, 1fr)); - grid-template-rows: repeat(var(--grid-rows), minmax(5em, 1fr)); -} - -:root:root:root .tab-nav > button +:root:root:root:root .generating, +:root:root:root:root .thumbnail-item { border: unset; - border-bottom: 0.125rem solid transparent; - font-size: 1.125em; - margin: 0.5rem 1rem; - padding: 0; } -:root:root:root .tab-nav > button.selected +:root:root:root:root .feather-upload, +:root:root:root:root footer { - border-bottom: 0.125rem solid; + display: none; +} + +:root:root:root:root .tab-nav > button +{ + border: unset; + box-shadow: 0 0.125rem; + font-size: 1.125em; + margin: 0.5rem 0.75rem; + padding: unset; +} + +:root:root:root:root .image-frame +{ + 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; } diff --git a/facefusion/uis/choices.py b/facefusion/uis/choices.py index cae02569..3650365e 100644 --- a/facefusion/uis/choices.py +++ b/facefusion/uis/choices.py @@ -1,7 +1,11 @@ from typing import List -from facefusion.uis.typing import WebcamMode +from facefusion.uis.typing import JobManagerAction, JobRunnerAction, WebcamMode + +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', 'skip-audio', 'skip-download' ] + webcam_modes : List[WebcamMode] = [ 'inline', 'udp', 'v4l2' ] webcam_resolutions : List[str] = [ '320x240', '640x480', '800x600', '1024x768', '1280x720', '1280x960', '1920x1080', '2560x1440', '3840x2160' ] diff --git a/facefusion/uis/components/about.py b/facefusion/uis/components/about.py index 01ecb00d..fedaba72 100644 --- a/facefusion/uis/components/about.py +++ b/facefusion/uis/components/about.py @@ -1,23 +1,41 @@ +import random from typing import Optional + import gradio from facefusion import metadata, wording -ABOUT_BUTTON : Optional[gradio.HTML] = None -DONATE_BUTTON : Optional[gradio.HTML] = None +METADATA_BUTTON : Optional[gradio.Button] = None +ACTION_BUTTON : Optional[gradio.Button] = None def render() -> None: - global ABOUT_BUTTON - global DONATE_BUTTON + global METADATA_BUTTON + global ACTION_BUTTON - ABOUT_BUTTON = gradio.Button( + action = random.choice( + [ + { + 'wording': wording.get('about.become_a_member'), + 'url': 'https://subscribe.facefusion.io' + }, + { + 'wording': wording.get('about.join_our_community'), + 'url': 'https://join.facefusion.io' + }, + { + 'wording': wording.get('about.read_the_documentation'), + 'url': 'https://docs.facefusion.io' + } + ]) + + METADATA_BUTTON = gradio.Button( value = metadata.get('name') + ' ' + metadata.get('version'), variant = 'primary', link = metadata.get('url') ) - DONATE_BUTTON = gradio.Button( - value = wording.get('uis.donate_button'), - link = 'https://donate.facefusion.io', + ACTION_BUTTON = gradio.Button( + value = action.get('wording'), + 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 100755 index 00000000..73b1de23 --- /dev/null +++ b/facefusion/uis/components/age_modifier_options.py @@ -0,0 +1,63 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, wording +from facefusion.common_helper import calc_float_step +from facefusion.processors import choices as processors_choices +from facefusion.processors.core import load_processor_module +from facefusion.processors.typing 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 + + AGE_MODIFIER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.age_modifier_model_dropdown'), + choices = processors_choices.age_modifier_models, + value = state_manager.get_item('age_modifier_model'), + visible = 'age_modifier' in state_manager.get_item('processors') + ) + AGE_MODIFIER_DIRECTION_SLIDER = gradio.Slider( + label = wording.get('uis.age_modifier_direction_slider'), + value = state_manager.get_item('age_modifier_direction'), + step = calc_float_step(processors_choices.age_modifier_direction_range), + minimum = processors_choices.age_modifier_direction_range[0], + maximum = processors_choices.age_modifier_direction_range[-1], + visible = 'age_modifier' in state_manager.get_item('processors') + ) + 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/benchmark.py b/facefusion/uis/components/benchmark.py index 6322febd..9fbdbd81 100644 --- a/facefusion/uis/components/benchmark.py +++ b/facefusion/uis/components/benchmark.py @@ -1,20 +1,20 @@ -from typing import Any, Optional, List, Dict, Generator -from time import sleep, perf_counter -import tempfile +import hashlib +import os import statistics +import tempfile +from time import perf_counter +from typing import Any, Dict, Generator, List, Optional + import gradio -import facefusion.globals -from facefusion import process_manager, wording -from facefusion.face_store import clear_static_faces -from facefusion.processors.frame.core import get_frame_processors_modules -from facefusion.vision import count_video_frame_total, detect_video_resolution, detect_video_fps, pack_resolution +from facefusion import state_manager, wording from facefusion.core import conditional_process +from facefusion.filesystem import is_video from facefusion.memory import limit_system_memory -from facefusion.filesystem import clear_temp from facefusion.uis.core import get_ui_component +from facefusion.vision import count_video_frame_total, detect_video_fps, detect_video_resolution, pack_resolution -BENCHMARK_RESULTS_DATAFRAME : Optional[gradio.Dataframe] = None +BENCHMARK_BENCHMARKS_DATAFRAME : Optional[gradio.Dataframe] = None BENCHMARK_START_BUTTON : Optional[gradio.Button] = None BENCHMARK_CLEAR_BUTTON : Optional[gradio.Button] = None BENCHMARKS : Dict[str, str] =\ @@ -30,12 +30,11 @@ BENCHMARKS : Dict[str, str] =\ def render() -> None: - global BENCHMARK_RESULTS_DATAFRAME + global BENCHMARK_BENCHMARKS_DATAFRAME global BENCHMARK_START_BUTTON global BENCHMARK_CLEAR_BUTTON - BENCHMARK_RESULTS_DATAFRAME = gradio.Dataframe( - label = wording.get('uis.benchmark_results_dataframe'), + BENCHMARK_BENCHMARKS_DATAFRAME = gradio.Dataframe( headers = [ 'target_path', @@ -53,17 +52,14 @@ def render() -> None: 'number', 'number', 'number' - ] + ], + show_label = False ) BENCHMARK_START_BUTTON = gradio.Button( value = wording.get('uis.start_button'), variant = 'primary', size = 'sm' ) - BENCHMARK_CLEAR_BUTTON = gradio.Button( - value = wording.get('uis.clear_button'), - size = 'sm' - ) def listen() -> None: @@ -71,46 +67,51 @@ def listen() -> None: benchmark_cycles_slider = get_ui_component('benchmark_cycles_slider') if benchmark_runs_checkbox_group and benchmark_cycles_slider: - BENCHMARK_START_BUTTON.click(start, inputs = [ benchmark_runs_checkbox_group, benchmark_cycles_slider ], outputs = BENCHMARK_RESULTS_DATAFRAME) - BENCHMARK_CLEAR_BUTTON.click(clear, outputs = BENCHMARK_RESULTS_DATAFRAME) + BENCHMARK_START_BUTTON.click(start, inputs = [ benchmark_runs_checkbox_group, benchmark_cycles_slider ], outputs = BENCHMARK_BENCHMARKS_DATAFRAME) + + +def suggest_output_path(target_path : str) -> Optional[str]: + if is_video(target_path): + _, target_extension = os.path.splitext(target_path) + return os.path.join(tempfile.gettempdir(), hashlib.sha1().hexdigest()[:8] + target_extension) + return None def start(benchmark_runs : List[str], benchmark_cycles : int) -> Generator[List[Any], None, None]: - facefusion.globals.source_paths = [ '.assets/examples/source.jpg', '.assets/examples/source.mp3' ] - facefusion.globals.output_path = tempfile.gettempdir() - facefusion.globals.face_landmarker_score = 0 - facefusion.globals.temp_frame_format = 'bmp' - facefusion.globals.output_video_preset = 'ultrafast' + 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_video_preset', 'ultrafast') + state_manager.sync_item('execution_providers') + state_manager.sync_item('execution_thread_count') + state_manager.sync_item('execution_queue_count') + state_manager.sync_item('system_memory_limit') benchmark_results = [] target_paths = [ BENCHMARKS[benchmark_run] for benchmark_run in benchmark_runs if benchmark_run in BENCHMARKS ] if target_paths: pre_process() for target_path in target_paths: - facefusion.globals.target_path = target_path + state_manager.init_item('target_path', target_path) + state_manager.init_item('output_path', suggest_output_path(state_manager.get_item('target_path'))) benchmark_results.append(benchmark(benchmark_cycles)) yield benchmark_results - post_process() def pre_process() -> None: - if facefusion.globals.system_memory_limit > 0: - limit_system_memory(facefusion.globals.system_memory_limit) - for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - frame_processor_module.get_frame_processor() - - -def post_process() -> None: - clear_static_faces() + 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) def benchmark(benchmark_cycles : int) -> List[Any]: process_times = [] - video_frame_total = count_video_frame_total(facefusion.globals.target_path) - output_video_resolution = detect_video_resolution(facefusion.globals.target_path) - facefusion.globals.output_video_resolution = pack_resolution(output_video_resolution) - facefusion.globals.output_video_fps = detect_video_fps(facefusion.globals.target_path) + video_frame_total = count_video_frame_total(state_manager.get_item('target_path')) + output_video_resolution = detect_video_resolution(state_manager.get_item('target_path')) + state_manager.init_item('output_video_resolution', pack_resolution(output_video_resolution)) + state_manager.init_item('output_video_fps', detect_video_fps(state_manager.get_item('target_path'))) + conditional_process() for index in range(benchmark_cycles): start_time = perf_counter() conditional_process() @@ -123,18 +124,10 @@ def benchmark(benchmark_cycles : int) -> List[Any]: return\ [ - facefusion.globals.target_path, + state_manager.get_item('target_path'), benchmark_cycles, average_run, fastest_run, slowest_run, relative_fps ] - - -def clear() -> gradio.Dataframe: - while process_manager.is_processing(): - sleep(0.5) - if facefusion.globals.target_path: - clear_temp(facefusion.globals.target_path) - return gradio.Dataframe(value = None) diff --git a/facefusion/uis/components/benchmark_options.py b/facefusion/uis/components/benchmark_options.py index 6748dd9b..5b5cda02 100644 --- a/facefusion/uis/components/benchmark_options.py +++ b/facefusion/uis/components/benchmark_options.py @@ -1,9 +1,10 @@ from typing import Optional + import gradio from facefusion import wording -from facefusion.uis.core import register_ui_component from facefusion.uis.components.benchmark import BENCHMARKS +from facefusion.uis.core import register_ui_component BENCHMARK_RUNS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None BENCHMARK_CYCLES_SLIDER : Optional[gradio.Button] = None diff --git a/facefusion/uis/components/common_options.py b/facefusion/uis/components/common_options.py index 43817413..0352ff34 100644 --- a/facefusion/uis/components/common_options.py +++ b/facefusion/uis/components/common_options.py @@ -1,8 +1,8 @@ -from typing import Optional, List +from typing import List, Optional + import gradio -import facefusion.globals -from facefusion import wording +from facefusion import state_manager, wording from facefusion.uis import choices as uis_choices COMMON_OPTIONS_CHECKBOX_GROUP : Optional[gradio.Checkboxgroup] = None @@ -11,17 +11,19 @@ COMMON_OPTIONS_CHECKBOX_GROUP : Optional[gradio.Checkboxgroup] = None def render() -> None: global COMMON_OPTIONS_CHECKBOX_GROUP - value = [] - if facefusion.globals.keep_temp: - value.append('keep-temp') - if facefusion.globals.skip_audio: - value.append('skip-audio') - if facefusion.globals.skip_download: - value.append('skip-download') + common_options = [] + + if state_manager.get_item('skip_download'): + common_options.append('skip-download') + if state_manager.get_item('keep_temp'): + common_options.append('keep-temp') + if state_manager.get_item('skip_audio'): + common_options.append('skip-audio') + COMMON_OPTIONS_CHECKBOX_GROUP = gradio.Checkboxgroup( label = wording.get('uis.common_options_checkbox_group'), choices = uis_choices.common_options, - value = value + value = common_options ) @@ -30,6 +32,9 @@ def listen() -> None: def update(common_options : List[str]) -> None: - facefusion.globals.keep_temp = 'keep-temp' in common_options - facefusion.globals.skip_audio = 'skip-audio' in common_options - facefusion.globals.skip_download = 'skip-download' in common_options + skip_temp = 'skip-download' in common_options + keep_temp = 'keep-temp' in common_options + skip_audio = 'skip-audio' in common_options + state_manager.set_item('skip_download', skip_temp) + state_manager.set_item('keep_temp', keep_temp) + state_manager.set_item('skip_audio', skip_audio) diff --git a/facefusion/uis/components/execution.py b/facefusion/uis/components/execution.py index 083727de..baf1888c 100644 --- a/facefusion/uis/components/execution.py +++ b/facefusion/uis/components/execution.py @@ -1,12 +1,11 @@ from typing import List, Optional -import gradio -import onnxruntime -import facefusion.globals -from facefusion import wording -from facefusion.face_analyser import clear_face_analyser -from facefusion.processors.frame.core import clear_frame_processors_modules -from facefusion.execution import encode_execution_providers, decode_execution_providers +import gradio + +from facefusion import content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, state_manager, voice_extractor, wording +from facefusion.execution import get_execution_provider_choices +from facefusion.processors.core import clear_processors_modules +from facefusion.typing import ExecutionProviderKey EXECUTION_PROVIDERS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None @@ -16,8 +15,8 @@ def render() -> None: EXECUTION_PROVIDERS_CHECKBOX_GROUP = gradio.CheckboxGroup( label = wording.get('uis.execution_providers_checkbox_group'), - choices = encode_execution_providers(onnxruntime.get_available_providers()), - value = encode_execution_providers(facefusion.globals.execution_providers) + choices = get_execution_provider_choices(), + value = state_manager.get_item('execution_providers') ) @@ -25,9 +24,15 @@ 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[str]) -> gradio.CheckboxGroup: - clear_face_analyser() - clear_frame_processors_modules() - execution_providers = execution_providers or encode_execution_providers(onnxruntime.get_available_providers()) - facefusion.globals.execution_providers = decode_execution_providers(execution_providers) - return gradio.CheckboxGroup(value = execution_providers) +def update_execution_providers(execution_providers : List[ExecutionProviderKey]) -> gradio.CheckboxGroup: + 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() + clear_processors_modules(state_manager.get_item('processors')) + execution_providers = execution_providers or get_execution_provider_choices() + 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_queue_count.py b/facefusion/uis/components/execution_queue_count.py index 1b6725e5..b5ab5dad 100644 --- a/facefusion/uis/components/execution_queue_count.py +++ b/facefusion/uis/components/execution_queue_count.py @@ -1,9 +1,10 @@ from typing import Optional + import gradio -import facefusion.globals import facefusion.choices -from facefusion import wording +from facefusion import state_manager, wording +from facefusion.common_helper import calc_int_step EXECUTION_QUEUE_COUNT_SLIDER : Optional[gradio.Slider] = None @@ -13,8 +14,8 @@ def render() -> None: EXECUTION_QUEUE_COUNT_SLIDER = gradio.Slider( label = wording.get('uis.execution_queue_count_slider'), - value = facefusion.globals.execution_queue_count, - step = facefusion.choices.execution_queue_count_range[1] - facefusion.choices.execution_queue_count_range[0], + value = state_manager.get_item('execution_queue_count'), + step = calc_int_step(facefusion.choices.execution_queue_count_range), minimum = facefusion.choices.execution_queue_count_range[0], maximum = facefusion.choices.execution_queue_count_range[-1] ) @@ -24,5 +25,5 @@ def listen() -> None: EXECUTION_QUEUE_COUNT_SLIDER.release(update_execution_queue_count, inputs = EXECUTION_QUEUE_COUNT_SLIDER) -def update_execution_queue_count(execution_queue_count : int = 1) -> None: - facefusion.globals.execution_queue_count = execution_queue_count +def update_execution_queue_count(execution_queue_count : float) -> None: + state_manager.set_item('execution_queue_count', int(execution_queue_count)) diff --git a/facefusion/uis/components/execution_thread_count.py b/facefusion/uis/components/execution_thread_count.py index 4a1f4646..f5716a99 100644 --- a/facefusion/uis/components/execution_thread_count.py +++ b/facefusion/uis/components/execution_thread_count.py @@ -1,9 +1,10 @@ from typing import Optional + import gradio -import facefusion.globals import facefusion.choices -from facefusion import wording +from facefusion import state_manager, wording +from facefusion.common_helper import calc_int_step EXECUTION_THREAD_COUNT_SLIDER : Optional[gradio.Slider] = None @@ -13,8 +14,8 @@ def render() -> None: EXECUTION_THREAD_COUNT_SLIDER = gradio.Slider( label = wording.get('uis.execution_thread_count_slider'), - value = facefusion.globals.execution_thread_count, - step = facefusion.choices.execution_thread_count_range[1] - facefusion.choices.execution_thread_count_range[0], + value = state_manager.get_item('execution_thread_count'), + step = calc_int_step(facefusion.choices.execution_thread_count_range), minimum = facefusion.choices.execution_thread_count_range[0], maximum = facefusion.choices.execution_thread_count_range[-1] ) @@ -24,6 +25,5 @@ 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 : int = 1) -> None: - facefusion.globals.execution_thread_count = execution_thread_count - +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 100755 index 00000000..eaa48ecf --- /dev/null +++ b/facefusion/uis/components/expression_restorer_options.py @@ -0,0 +1,63 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, wording +from facefusion.common_helper import calc_float_step +from facefusion.processors import choices as processors_choices +from facefusion.processors.core import load_processor_module +from facefusion.processors.typing import 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 + + +def render() -> None: + global EXPRESSION_RESTORER_MODEL_DROPDOWN + global EXPRESSION_RESTORER_FACTOR_SLIDER + + EXPRESSION_RESTORER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.expression_restorer_model_dropdown'), + choices = processors_choices.expression_restorer_models, + value = state_manager.get_item('expression_restorer_model'), + visible = 'expression_restorer' in state_manager.get_item('processors') + ) + EXPRESSION_RESTORER_FACTOR_SLIDER = gradio.Slider( + label = wording.get('uis.expression_restorer_factor_slider'), + value = state_manager.get_item('expression_restorer_factor'), + step = calc_float_step(processors_choices.expression_restorer_factor_range), + minimum = processors_choices.expression_restorer_factor_range[0], + maximum = processors_choices.expression_restorer_factor_range[-1], + visible = 'expression_restorer' in state_manager.get_item('processors'), + ) + register_ui_component('expression_restorer_model_dropdown', EXPRESSION_RESTORER_MODEL_DROPDOWN) + register_ui_component('expression_restorer_factor_slider', EXPRESSION_RESTORER_FACTOR_SLIDER) + + +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) + + 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 ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider]: + has_expression_restorer = 'expression_restorer' in processors + return gradio.Dropdown(visible = has_expression_restorer), gradio.Slider(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)) diff --git a/facefusion/uis/components/face_analyser.py b/facefusion/uis/components/face_analyser.py deleted file mode 100644 index aed04182..00000000 --- a/facefusion/uis/components/face_analyser.py +++ /dev/null @@ -1,123 +0,0 @@ -from typing import Optional, Dict, Any, Tuple - -import gradio - -import facefusion.globals -import facefusion.choices -from facefusion import face_analyser, wording -from facefusion.typing import FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, FaceDetectorModel -from facefusion.uis.core import register_ui_component - -FACE_ANALYSER_ORDER_DROPDOWN : Optional[gradio.Dropdown] = None -FACE_ANALYSER_AGE_DROPDOWN : Optional[gradio.Dropdown] = None -FACE_ANALYSER_GENDER_DROPDOWN : Optional[gradio.Dropdown] = None -FACE_DETECTOR_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None -FACE_DETECTOR_SIZE_DROPDOWN : Optional[gradio.Dropdown] = None -FACE_DETECTOR_SCORE_SLIDER : Optional[gradio.Slider] = None -FACE_LANDMARKER_SCORE_SLIDER : Optional[gradio.Slider] = None - - -def render() -> None: - global FACE_ANALYSER_ORDER_DROPDOWN - global FACE_ANALYSER_AGE_DROPDOWN - global FACE_ANALYSER_GENDER_DROPDOWN - global FACE_DETECTOR_MODEL_DROPDOWN - global FACE_DETECTOR_SIZE_DROPDOWN - global FACE_DETECTOR_SCORE_SLIDER - global FACE_LANDMARKER_SCORE_SLIDER - - face_detector_size_dropdown_args : Dict[str, Any] =\ - { - 'label': wording.get('uis.face_detector_size_dropdown'), - 'value': facefusion.globals.face_detector_size - } - if facefusion.globals.face_detector_size in facefusion.choices.face_detector_set[facefusion.globals.face_detector_model]: - face_detector_size_dropdown_args['choices'] = facefusion.choices.face_detector_set[facefusion.globals.face_detector_model] - with gradio.Row(): - FACE_ANALYSER_ORDER_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.face_analyser_order_dropdown'), - choices = facefusion.choices.face_analyser_orders, - value = facefusion.globals.face_analyser_order - ) - FACE_ANALYSER_AGE_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.face_analyser_age_dropdown'), - choices = [ 'none' ] + facefusion.choices.face_analyser_ages, - value = facefusion.globals.face_analyser_age or 'none' - ) - FACE_ANALYSER_GENDER_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.face_analyser_gender_dropdown'), - choices = [ 'none' ] + facefusion.choices.face_analyser_genders, - value = facefusion.globals.face_analyser_gender or 'none' - ) - FACE_DETECTOR_MODEL_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.face_detector_model_dropdown'), - choices = facefusion.choices.face_detector_set.keys(), - value = facefusion.globals.face_detector_model - ) - FACE_DETECTOR_SIZE_DROPDOWN = gradio.Dropdown(**face_detector_size_dropdown_args) - with gradio.Row(): - FACE_DETECTOR_SCORE_SLIDER = gradio.Slider( - label = wording.get('uis.face_detector_score_slider'), - value = facefusion.globals.face_detector_score, - step = facefusion.choices.face_detector_score_range[1] - facefusion.choices.face_detector_score_range[0], - minimum = facefusion.choices.face_detector_score_range[0], - maximum = facefusion.choices.face_detector_score_range[-1] - ) - FACE_LANDMARKER_SCORE_SLIDER = gradio.Slider( - label = wording.get('uis.face_landmarker_score_slider'), - value = facefusion.globals.face_landmarker_score, - step = facefusion.choices.face_landmarker_score_range[1] - facefusion.choices.face_landmarker_score_range[0], - minimum = facefusion.choices.face_landmarker_score_range[0], - maximum = facefusion.choices.face_landmarker_score_range[-1] - ) - register_ui_component('face_analyser_order_dropdown', FACE_ANALYSER_ORDER_DROPDOWN) - register_ui_component('face_analyser_age_dropdown', FACE_ANALYSER_AGE_DROPDOWN) - register_ui_component('face_analyser_gender_dropdown', FACE_ANALYSER_GENDER_DROPDOWN) - 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_score_slider', FACE_DETECTOR_SCORE_SLIDER) - register_ui_component('face_landmarker_score_slider', FACE_LANDMARKER_SCORE_SLIDER) - - -def listen() -> None: - FACE_ANALYSER_ORDER_DROPDOWN.change(update_face_analyser_order, inputs = FACE_ANALYSER_ORDER_DROPDOWN) - FACE_ANALYSER_AGE_DROPDOWN.change(update_face_analyser_age, inputs = FACE_ANALYSER_AGE_DROPDOWN) - FACE_ANALYSER_GENDER_DROPDOWN.change(update_face_analyser_gender, inputs = FACE_ANALYSER_GENDER_DROPDOWN) - 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_SCORE_SLIDER.release(update_face_detector_score, inputs = FACE_DETECTOR_SCORE_SLIDER) - FACE_LANDMARKER_SCORE_SLIDER.release(update_face_landmarker_score, inputs = FACE_LANDMARKER_SCORE_SLIDER) - - -def update_face_analyser_order(face_analyser_order : FaceAnalyserOrder) -> None: - facefusion.globals.face_analyser_order = face_analyser_order if face_analyser_order != 'none' else None - - -def update_face_analyser_age(face_analyser_age : FaceAnalyserAge) -> None: - facefusion.globals.face_analyser_age = face_analyser_age if face_analyser_age != 'none' else None - - -def update_face_analyser_gender(face_analyser_gender : FaceAnalyserGender) -> None: - facefusion.globals.face_analyser_gender = face_analyser_gender if face_analyser_gender != 'none' else None - - -def update_face_detector_model(face_detector_model : FaceDetectorModel) -> Tuple[gradio.Dropdown, gradio.Dropdown]: - facefusion.globals.face_detector_model = face_detector_model - update_face_detector_size('640x640') - if face_analyser.pre_check(): - if facefusion.globals.face_detector_size in facefusion.choices.face_detector_set[face_detector_model]: - return gradio.Dropdown(value = facefusion.globals.face_detector_model), gradio.Dropdown(value = facefusion.globals.face_detector_size, choices = facefusion.choices.face_detector_set[face_detector_model]) - return gradio.Dropdown(value = facefusion.globals.face_detector_model), gradio.Dropdown(value = facefusion.globals.face_detector_size, choices = [ facefusion.globals.face_detector_size ]) - return gradio.Dropdown(), gradio.Dropdown() - - -def update_face_detector_size(face_detector_size : str) -> None: - facefusion.globals.face_detector_size = face_detector_size - - -def update_face_detector_score(face_detector_score : float) -> None: - facefusion.globals.face_detector_score = face_detector_score - - -def update_face_landmarker_score(face_landmarker_score : float) -> None: - facefusion.globals.face_landmarker_score = face_landmarker_score diff --git a/facefusion/uis/components/face_debugger_options.py b/facefusion/uis/components/face_debugger_options.py new file mode 100755 index 00000000..088b0877 --- /dev/null +++ b/facefusion/uis/components/face_debugger_options.py @@ -0,0 +1,39 @@ +from typing import List, Optional + +import gradio + +from facefusion import state_manager, wording +from facefusion.processors import choices as processors_choices +from facefusion.processors.typing 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 + + FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = wording.get('uis.face_debugger_items_checkbox_group'), + choices = processors_choices.face_debugger_items, + value = state_manager.get_item('face_debugger_items'), + visible = 'face_debugger' in state_manager.get_item('processors') + ) + 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 00000000..eb0e20fc --- /dev/null +++ b/facefusion/uis/components/face_detector.py @@ -0,0 +1,85 @@ +from typing import Optional, Sequence, Tuple + +import gradio + +import facefusion.choices +from facefusion import choices, face_detector, state_manager, wording +from facefusion.common_helper import calc_float_step, get_last +from facefusion.typing import Angle, FaceDetectorModel, Score +from facefusion.uis.core import register_ui_component +from facefusion.uis.typing import ComponentOptions + +FACE_DETECTOR_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_DETECTOR_SIZE_DROPDOWN : Optional[gradio.Dropdown] = 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_ANGLES_CHECKBOX_GROUP + global FACE_DETECTOR_SCORE_SLIDER + + face_detector_size_dropdown_options : ComponentOptions =\ + { + 'label': wording.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 = wording.get('uis.face_detector_model_dropdown'), + choices = facefusion.choices.face_detector_set.keys(), + value = state_manager.get_item('face_detector_model') + ) + FACE_DETECTOR_SIZE_DROPDOWN = gradio.Dropdown(**face_detector_size_dropdown_options) + FACE_DETECTOR_ANGLES_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = wording.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 = wording.get('uis.face_detector_score_slider'), + value = state_manager.get_item('face_detector_score'), + step = calc_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_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_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 = 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_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 100755 index 00000000..cf4efa0f --- /dev/null +++ b/facefusion/uis/components/face_editor_options.py @@ -0,0 +1,271 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, wording +from facefusion.common_helper import calc_float_step +from facefusion.processors import choices as processors_choices +from facefusion.processors.core import load_processor_module +from facefusion.processors.typing 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 + + FACE_EDITOR_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.face_editor_model_dropdown'), + choices = processors_choices.face_editor_models, + value = state_manager.get_item('face_editor_model'), + visible = 'face_editor' in state_manager.get_item('processors') + ) + FACE_EDITOR_EYEBROW_DIRECTION_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_eyebrow_direction_slider'), + value = state_manager.get_item('face_editor_eyebrow_direction'), + step = calc_float_step(processors_choices.face_editor_eyebrow_direction_range), + minimum = processors_choices.face_editor_eyebrow_direction_range[0], + maximum = processors_choices.face_editor_eyebrow_direction_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_EYE_GAZE_HORIZONTAL_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_eye_gaze_horizontal_slider'), + value = state_manager.get_item('face_editor_eye_gaze_horizontal'), + step = calc_float_step(processors_choices.face_editor_eye_gaze_horizontal_range), + minimum = processors_choices.face_editor_eye_gaze_horizontal_range[0], + maximum = processors_choices.face_editor_eye_gaze_horizontal_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_EYE_GAZE_VERTICAL_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_eye_gaze_vertical_slider'), + value = state_manager.get_item('face_editor_eye_gaze_vertical'), + step = calc_float_step(processors_choices.face_editor_eye_gaze_vertical_range), + minimum = processors_choices.face_editor_eye_gaze_vertical_range[0], + maximum = processors_choices.face_editor_eye_gaze_vertical_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_EYE_OPEN_RATIO_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_eye_open_ratio_slider'), + value = state_manager.get_item('face_editor_eye_open_ratio'), + step = calc_float_step(processors_choices.face_editor_eye_open_ratio_range), + minimum = processors_choices.face_editor_eye_open_ratio_range[0], + maximum = processors_choices.face_editor_eye_open_ratio_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_LIP_OPEN_RATIO_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_lip_open_ratio_slider'), + value = state_manager.get_item('face_editor_lip_open_ratio'), + step = calc_float_step(processors_choices.face_editor_lip_open_ratio_range), + minimum = processors_choices.face_editor_lip_open_ratio_range[0], + maximum = processors_choices.face_editor_lip_open_ratio_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_MOUTH_GRIM_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_mouth_grim_slider'), + value = state_manager.get_item('face_editor_mouth_grim'), + step = calc_float_step(processors_choices.face_editor_mouth_grim_range), + minimum = processors_choices.face_editor_mouth_grim_range[0], + maximum = processors_choices.face_editor_mouth_grim_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_MOUTH_POUT_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_mouth_pout_slider'), + value = state_manager.get_item('face_editor_mouth_pout'), + step = calc_float_step(processors_choices.face_editor_mouth_pout_range), + minimum = processors_choices.face_editor_mouth_pout_range[0], + maximum = processors_choices.face_editor_mouth_pout_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_MOUTH_PURSE_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_mouth_purse_slider'), + value = state_manager.get_item('face_editor_mouth_purse'), + step = calc_float_step(processors_choices.face_editor_mouth_purse_range), + minimum = processors_choices.face_editor_mouth_purse_range[0], + maximum = processors_choices.face_editor_mouth_purse_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_MOUTH_SMILE_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_mouth_smile_slider'), + value = state_manager.get_item('face_editor_mouth_smile'), + step = calc_float_step(processors_choices.face_editor_mouth_smile_range), + minimum = processors_choices.face_editor_mouth_smile_range[0], + maximum = processors_choices.face_editor_mouth_smile_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_MOUTH_POSITION_HORIZONTAL_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_mouth_position_horizontal_slider'), + value = state_manager.get_item('face_editor_mouth_position_horizontal'), + step = calc_float_step(processors_choices.face_editor_mouth_position_horizontal_range), + minimum = processors_choices.face_editor_mouth_position_horizontal_range[0], + maximum = processors_choices.face_editor_mouth_position_horizontal_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_MOUTH_POSITION_VERTICAL_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_mouth_position_vertical_slider'), + value = state_manager.get_item('face_editor_mouth_position_vertical'), + step = calc_float_step(processors_choices.face_editor_mouth_position_vertical_range), + minimum = processors_choices.face_editor_mouth_position_vertical_range[0], + maximum = processors_choices.face_editor_mouth_position_vertical_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_HEAD_PITCH_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_head_pitch_slider'), + value = state_manager.get_item('face_editor_head_pitch'), + step = calc_float_step(processors_choices.face_editor_head_pitch_range), + minimum = processors_choices.face_editor_head_pitch_range[0], + maximum = processors_choices.face_editor_head_pitch_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_HEAD_YAW_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_head_yaw_slider'), + value = state_manager.get_item('face_editor_head_yaw'), + step = calc_float_step(processors_choices.face_editor_head_yaw_range), + minimum = processors_choices.face_editor_head_yaw_range[0], + maximum = processors_choices.face_editor_head_yaw_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + FACE_EDITOR_HEAD_ROLL_SLIDER = gradio.Slider( + label = wording.get('uis.face_editor_head_roll_slider'), + value = state_manager.get_item('face_editor_head_roll'), + step = calc_float_step(processors_choices.face_editor_head_roll_range), + minimum = processors_choices.face_editor_head_roll_range[0], + maximum = processors_choices.face_editor_head_roll_range[-1], + visible = 'face_editor' in state_manager.get_item('processors'), + ) + 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 100755 index 00000000..5ce4c11a --- /dev/null +++ b/facefusion/uis/components/face_enhancer_options.py @@ -0,0 +1,63 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, wording +from facefusion.common_helper import calc_int_step +from facefusion.processors import choices as processors_choices +from facefusion.processors.core import load_processor_module +from facefusion.processors.typing import FaceEnhancerModel +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 + + +def render() -> None: + global FACE_ENHANCER_MODEL_DROPDOWN + global FACE_ENHANCER_BLEND_SLIDER + + FACE_ENHANCER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.face_enhancer_model_dropdown'), + choices = processors_choices.face_enhancer_models, + value = state_manager.get_item('face_enhancer_model'), + visible = 'face_enhancer' in state_manager.get_item('processors') + ) + FACE_ENHANCER_BLEND_SLIDER = gradio.Slider( + label = wording.get('uis.face_enhancer_blend_slider'), + value = state_manager.get_item('face_enhancer_blend'), + step = calc_int_step(processors_choices.face_enhancer_blend_range), + minimum = processors_choices.face_enhancer_blend_range[0], + maximum = processors_choices.face_enhancer_blend_range[-1], + visible = 'face_enhancer' in state_manager.get_item('processors') + ) + register_ui_component('face_enhancer_model_dropdown', FACE_ENHANCER_MODEL_DROPDOWN) + register_ui_component('face_enhancer_blend_slider', FACE_ENHANCER_BLEND_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_BLEND_SLIDER.release(update_face_enhancer_blend, inputs = FACE_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 = [ FACE_ENHANCER_MODEL_DROPDOWN, FACE_ENHANCER_BLEND_SLIDER ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Slider]: + has_face_enhancer = 'face_enhancer' in processors + return gradio.Dropdown(visible = has_face_enhancer), gradio.Slider(visible = has_face_enhancer) + + +def update_face_enhancer_model(face_enhancer_model : FaceEnhancerModel) -> gradio.Dropdown: + 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')) + return gradio.Dropdown() + + +def update_face_enhancer_blend(face_enhancer_blend : float) -> None: + state_manager.set_item('face_enhancer_blend', int(face_enhancer_blend)) diff --git a/facefusion/uis/components/face_landmarker.py b/facefusion/uis/components/face_landmarker.py new file mode 100644 index 00000000..f6a179d7 --- /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, wording +from facefusion.common_helper import calc_float_step +from facefusion.typing 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 = wording.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 = wording.get('uis.face_landmarker_score_slider'), + value = state_manager.get_item('face_landmarker_score'), + step = calc_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 index bb1c28c9..7579cf80 100755 --- a/facefusion/uis/components/face_masker.py +++ b/facefusion/uis/components/face_masker.py @@ -1,16 +1,15 @@ -from typing import Optional, Tuple, List +from typing import List, Optional, Tuple + import gradio -import facefusion.globals import facefusion.choices -from facefusion import wording -from facefusion.typing import FaceMaskType, FaceMaskRegion +from facefusion import state_manager, wording +from facefusion.common_helper import calc_float_step, calc_int_step +from facefusion.typing import FaceMaskRegion, FaceMaskType from facefusion.uis.core import register_ui_component FACE_MASK_TYPES_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None FACE_MASK_BLUR_SLIDER : Optional[gradio.Slider] = None -FACE_MASK_BOX_GROUP : Optional[gradio.Group] = None -FACE_MASK_REGION_GROUP : Optional[gradio.Group] = 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 @@ -20,100 +19,105 @@ FACE_MASK_REGION_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None def render() -> None: global FACE_MASK_TYPES_CHECKBOX_GROUP + global FACE_MASK_REGION_CHECKBOX_GROUP global FACE_MASK_BLUR_SLIDER - global FACE_MASK_BOX_GROUP - global FACE_MASK_REGION_GROUP global FACE_MASK_PADDING_TOP_SLIDER global FACE_MASK_PADDING_RIGHT_SLIDER global FACE_MASK_PADDING_BOTTOM_SLIDER global FACE_MASK_PADDING_LEFT_SLIDER - global FACE_MASK_REGION_CHECKBOX_GROUP - has_box_mask = 'box' in facefusion.globals.face_mask_types - has_region_mask = 'region' in facefusion.globals.face_mask_types + has_box_mask = 'box' in state_manager.get_item('face_mask_types') + has_region_mask = 'region' in state_manager.get_item('face_mask_types') FACE_MASK_TYPES_CHECKBOX_GROUP = gradio.CheckboxGroup( label = wording.get('uis.face_mask_types_checkbox_group'), choices = facefusion.choices.face_mask_types, - value = facefusion.globals.face_mask_types + value = state_manager.get_item('face_mask_types') ) - with gradio.Group(visible = has_box_mask) as FACE_MASK_BOX_GROUP: - FACE_MASK_BLUR_SLIDER = gradio.Slider( - label = wording.get('uis.face_mask_blur_slider'), - step = facefusion.choices.face_mask_blur_range[1] - facefusion.choices.face_mask_blur_range[0], - minimum = facefusion.choices.face_mask_blur_range[0], - maximum = facefusion.choices.face_mask_blur_range[-1], - value = facefusion.globals.face_mask_blur - ) + FACE_MASK_REGION_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = wording.get('uis.face_mask_region_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 = wording.get('uis.face_mask_blur_slider'), + step = calc_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(): with gradio.Row(): FACE_MASK_PADDING_TOP_SLIDER = gradio.Slider( label = wording.get('uis.face_mask_padding_top_slider'), - step = facefusion.choices.face_mask_padding_range[1] - facefusion.choices.face_mask_padding_range[0], + step = calc_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 = facefusion.globals.face_mask_padding[0] + value = state_manager.get_item('face_mask_padding')[0], + visible = has_box_mask ) FACE_MASK_PADDING_RIGHT_SLIDER = gradio.Slider( label = wording.get('uis.face_mask_padding_right_slider'), - step = facefusion.choices.face_mask_padding_range[1] - facefusion.choices.face_mask_padding_range[0], + step = calc_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 = facefusion.globals.face_mask_padding[1] + value = state_manager.get_item('face_mask_padding')[1], + visible = has_box_mask ) with gradio.Row(): FACE_MASK_PADDING_BOTTOM_SLIDER = gradio.Slider( label = wording.get('uis.face_mask_padding_bottom_slider'), - step = facefusion.choices.face_mask_padding_range[1] - facefusion.choices.face_mask_padding_range[0], + step = calc_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 = facefusion.globals.face_mask_padding[2] + value = state_manager.get_item('face_mask_padding')[2], + visible = has_box_mask ) FACE_MASK_PADDING_LEFT_SLIDER = gradio.Slider( label = wording.get('uis.face_mask_padding_left_slider'), - step = facefusion.choices.face_mask_padding_range[1] - facefusion.choices.face_mask_padding_range[0], + step = calc_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 = facefusion.globals.face_mask_padding[3] + value = state_manager.get_item('face_mask_padding')[3], + visible = has_box_mask ) - with gradio.Row(): - FACE_MASK_REGION_CHECKBOX_GROUP = gradio.CheckboxGroup( - label = wording.get('uis.face_mask_region_checkbox_group'), - choices = facefusion.choices.face_mask_regions, - value = facefusion.globals.face_mask_regions, - visible = has_region_mask - ) register_ui_component('face_mask_types_checkbox_group', FACE_MASK_TYPES_CHECKBOX_GROUP) + register_ui_component('face_mask_region_checkbox_group', FACE_MASK_REGION_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) - register_ui_component('face_mask_region_checkbox_group', FACE_MASK_REGION_CHECKBOX_GROUP) def listen() -> None: - FACE_MASK_TYPES_CHECKBOX_GROUP.change(update_face_mask_type, inputs = FACE_MASK_TYPES_CHECKBOX_GROUP, outputs = [ FACE_MASK_TYPES_CHECKBOX_GROUP, FACE_MASK_BOX_GROUP, FACE_MASK_REGION_CHECKBOX_GROUP ]) - FACE_MASK_BLUR_SLIDER.release(update_face_mask_blur, inputs = FACE_MASK_BLUR_SLIDER) + FACE_MASK_TYPES_CHECKBOX_GROUP.change(update_face_mask_type, inputs = FACE_MASK_TYPES_CHECKBOX_GROUP, outputs = [ FACE_MASK_TYPES_CHECKBOX_GROUP, FACE_MASK_REGION_CHECKBOX_GROUP, FACE_MASK_BLUR_SLIDER, FACE_MASK_PADDING_TOP_SLIDER, FACE_MASK_PADDING_RIGHT_SLIDER, FACE_MASK_PADDING_BOTTOM_SLIDER, FACE_MASK_PADDING_LEFT_SLIDER ]) FACE_MASK_REGION_CHECKBOX_GROUP.change(update_face_mask_regions, inputs = FACE_MASK_REGION_CHECKBOX_GROUP, outputs = FACE_MASK_REGION_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_mask_type(face_mask_types : List[FaceMaskType]) -> Tuple[gradio.CheckboxGroup, gradio.Group, gradio.CheckboxGroup]: - facefusion.globals.face_mask_types = face_mask_types or facefusion.choices.face_mask_types +def update_face_mask_type(face_mask_types : List[FaceMaskType]) -> Tuple[gradio.CheckboxGroup, gradio.CheckboxGroup, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider]: + 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_region_mask = 'region' in face_mask_types - return gradio.CheckboxGroup(value = facefusion.globals.face_mask_types), gradio.Group(visible = has_box_mask), gradio.CheckboxGroup(visible = has_region_mask) - - -def update_face_mask_blur(face_mask_blur : float) -> None: - facefusion.globals.face_mask_blur = face_mask_blur - - -def update_face_mask_padding(face_mask_padding_top : int, face_mask_padding_right : int, face_mask_padding_bottom : int, face_mask_padding_left : int) -> None: - facefusion.globals.face_mask_padding = (face_mask_padding_top, face_mask_padding_right, face_mask_padding_bottom, face_mask_padding_left) + return gradio.CheckboxGroup(value = state_manager.get_item('face_mask_types')), gradio.CheckboxGroup(visible = has_region_mask), gradio.Slider(visible = has_box_mask), gradio.Slider(visible = has_box_mask), gradio.Slider(visible = has_box_mask), gradio.Slider(visible = has_box_mask), gradio.Slider(visible = has_box_mask) def update_face_mask_regions(face_mask_regions : List[FaceMaskRegion]) -> gradio.CheckboxGroup: - facefusion.globals.face_mask_regions = face_mask_regions or facefusion.choices.face_mask_regions - return gradio.CheckboxGroup(value = facefusion.globals.face_mask_regions) + 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 = (int(face_mask_padding_top), int(face_mask_padding_right), int(face_mask_padding_bottom), int(face_mask_padding_left)) + state_manager.set_item('face_mask_padding', face_mask_padding) diff --git a/facefusion/uis/components/face_selector.py b/facefusion/uis/components/face_selector.py index 19fef55b..2cb31577 100644 --- a/facefusion/uis/components/face_selector.py +++ b/facefusion/uis/components/face_selector.py @@ -1,62 +1,109 @@ -from typing import List, Optional, Tuple, Any, Dict +from typing import List, Optional, Tuple import gradio +from gradio_rangeslider import RangeSlider -import facefusion.globals import facefusion.choices -from facefusion import wording -from facefusion.face_store import clear_static_faces, clear_reference_faces -from facefusion.vision import get_video_frame, read_static_image, normalize_frame_color -from facefusion.filesystem import is_image, is_video +from facefusion import state_manager, wording +from facefusion.common_helper import calc_float_step, calc_int_step from facefusion.face_analyser import get_many_faces -from facefusion.typing import VisionFrame, FaceSelectorMode +from facefusion.face_selector import sort_and_filter_faces +from facefusion.face_store import clear_reference_faces, clear_static_faces +from facefusion.filesystem import is_image, is_video +from facefusion.typing import FaceSelectorMode, FaceSelectorOrder, Gender, Race, VisionFrame from facefusion.uis.core import get_ui_component, get_ui_components, register_ui_component +from facefusion.uis.typing import ComponentOptions +from facefusion.uis.ui_helper import convert_str_none +from facefusion.vision import get_video_frame, normalize_frame_color, read_static_image 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_args : Dict[str, Any] =\ + reference_face_gallery_options : ComponentOptions =\ { 'label': wording.get('uis.reference_face_gallery'), 'object_fit': 'cover', 'columns': 8, 'allow_preview': False, - 'visible': 'reference' in facefusion.globals.face_selector_mode + 'visible': 'reference' in state_manager.get_item('face_selector_mode') } - if is_image(facefusion.globals.target_path): - reference_frame = read_static_image(facefusion.globals.target_path) - reference_face_gallery_args['value'] = extract_gallery_frames(reference_frame) - if is_video(facefusion.globals.target_path): - reference_frame = get_video_frame(facefusion.globals.target_path, facefusion.globals.reference_frame_number) - reference_face_gallery_args['value'] = extract_gallery_frames(reference_frame) + if is_image(state_manager.get_item('target_path')): + reference_frame = read_static_image(state_manager.get_item('target_path')) + reference_face_gallery_options['value'] = extract_gallery_frames(reference_frame) + if is_video(state_manager.get_item('target_path')): + reference_frame = get_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) + reference_face_gallery_options['value'] = extract_gallery_frames(reference_frame) FACE_SELECTOR_MODE_DROPDOWN = gradio.Dropdown( label = wording.get('uis.face_selector_mode_dropdown'), choices = facefusion.choices.face_selector_modes, - value = facefusion.globals.face_selector_mode + value = state_manager.get_item('face_selector_mode') ) - REFERENCE_FACE_POSITION_GALLERY = gradio.Gallery(**reference_face_gallery_args) + REFERENCE_FACE_POSITION_GALLERY = gradio.Gallery(**reference_face_gallery_options) + with gradio.Group(): + with gradio.Row(): + FACE_SELECTOR_ORDER_DROPDOWN = gradio.Dropdown( + label = wording.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 = wording.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 = wording.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 = wording.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 = calc_int_step(facefusion.choices.face_selector_age_range) + ) REFERENCE_FACE_DISTANCE_SLIDER = gradio.Slider( label = wording.get('uis.reference_face_distance_slider'), - value = facefusion.globals.reference_face_distance, - step = facefusion.choices.reference_face_distance_range[1] - facefusion.choices.reference_face_distance_range[0], + value = state_manager.get_item('reference_face_distance'), + step = calc_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 facefusion.globals.face_selector_mode + 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_POSITION_GALLERY.select(clear_and_update_reference_face_position) REFERENCE_FACE_DISTANCE_SLIDER.release(update_reference_face_distance, inputs = REFERENCE_FACE_DISTANCE_SLIDER) @@ -69,46 +116,56 @@ def listen() -> None: getattr(ui_component, method)(update_reference_face_position) getattr(ui_component, method)(update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY) - for ui_component in get_ui_components( - [ - 'face_analyser_order_dropdown', - 'face_analyser_age_dropdown', - 'face_analyser_gender_dropdown' - ]): - ui_component.change(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_size_dropdown', + 'face_detector_angles_checkbox_group' ]): ui_component.change(clear_and_update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY) - for ui_component in get_ui_components( - [ - 'face_detector_score_slider', - 'face_landmarker_score_slider' - ]): - ui_component.release(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(clear_and_update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY) preview_frame_slider = get_ui_component('preview_frame_slider') if preview_frame_slider: - preview_frame_slider.change(update_reference_frame_number, inputs = preview_frame_slider) + preview_frame_slider.release(update_reference_frame_number, inputs = preview_frame_slider) preview_frame_slider.release(update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY) 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': - facefusion.globals.face_selector_mode = face_selector_mode return gradio.Gallery(visible = False), gradio.Slider(visible = False) if face_selector_mode == 'one': - facefusion.globals.face_selector_mode = face_selector_mode return gradio.Gallery(visible = False), gradio.Slider(visible = False) if face_selector_mode == 'reference': - facefusion.globals.face_selector_mode = face_selector_mode 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 clear_and_update_reference_face_position(event : gradio.SelectData) -> gradio.Gallery: clear_reference_faces() clear_static_faces() @@ -117,15 +174,15 @@ def clear_and_update_reference_face_position(event : gradio.SelectData) -> gradi def update_reference_face_position(reference_face_position : int = 0) -> None: - facefusion.globals.reference_face_position = reference_face_position + state_manager.set_item('reference_face_position', reference_face_position) def update_reference_face_distance(reference_face_distance : float) -> None: - facefusion.globals.reference_face_distance = reference_face_distance + state_manager.set_item('reference_face_distance', reference_face_distance) def update_reference_frame_number(reference_frame_number : int) -> None: - facefusion.globals.reference_frame_number = reference_frame_number + state_manager.set_item('reference_frame_number', reference_frame_number) def clear_and_update_reference_position_gallery() -> gradio.Gallery: @@ -136,11 +193,11 @@ def clear_and_update_reference_position_gallery() -> gradio.Gallery: def update_reference_position_gallery() -> gradio.Gallery: gallery_vision_frames = [] - if is_image(facefusion.globals.target_path): - temp_vision_frame = read_static_image(facefusion.globals.target_path) + if is_image(state_manager.get_item('target_path')): + temp_vision_frame = read_static_image(state_manager.get_item('target_path')) gallery_vision_frames = extract_gallery_frames(temp_vision_frame) - if is_video(facefusion.globals.target_path): - temp_vision_frame = get_video_frame(facefusion.globals.target_path, facefusion.globals.reference_frame_number) + if is_video(state_manager.get_item('target_path')): + temp_vision_frame = get_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) gallery_vision_frames = extract_gallery_frames(temp_vision_frame) if gallery_vision_frames: return gradio.Gallery(value = gallery_vision_frames) @@ -149,7 +206,7 @@ def update_reference_position_gallery() -> gradio.Gallery: def extract_gallery_frames(temp_vision_frame : VisionFrame) -> List[VisionFrame]: gallery_vision_frames = [] - faces = get_many_faces(temp_vision_frame) + faces = sort_and_filter_faces(get_many_faces([ temp_vision_frame ])) for face in faces: start_x, start_y, end_x, end_y = map(int, face.bounding_box) diff --git a/facefusion/uis/components/face_swapper_options.py b/facefusion/uis/components/face_swapper_options.py new file mode 100755 index 00000000..7eb4b713 --- /dev/null +++ b/facefusion/uis/components/face_swapper_options.py @@ -0,0 +1,63 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, wording +from facefusion.common_helper import get_first +from facefusion.processors import choices as processors_choices +from facefusion.processors.core import load_processor_module +from facefusion.processors.typing import FaceSwapperModel +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 + + +def render() -> None: + global FACE_SWAPPER_MODEL_DROPDOWN + global FACE_SWAPPER_PIXEL_BOOST_DROPDOWN + + FACE_SWAPPER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.face_swapper_model_dropdown'), + choices = processors_choices.face_swapper_set.keys(), + value = state_manager.get_item('face_swapper_model'), + visible = 'face_swapper' in state_manager.get_item('processors') + ) + FACE_SWAPPER_PIXEL_BOOST_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.face_swapper_pixel_boost_dropdown'), + choices = processors_choices.face_swapper_set.get(state_manager.get_item('face_swapper_model')), + value = state_manager.get_item('face_swapper_pixel_boost'), + visible = 'face_swapper' in state_manager.get_item('processors') + ) + register_ui_component('face_swapper_model_dropdown', FACE_SWAPPER_MODEL_DROPDOWN) + register_ui_component('face_swapper_pixel_boost_dropdown', FACE_SWAPPER_PIXEL_BOOST_DROPDOWN) + + +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_PIXEL_BOOST_DROPDOWN.change(update_face_swapper_pixel_boost, inputs = FACE_SWAPPER_PIXEL_BOOST_DROPDOWN) + + 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 ]) + + +def remote_update(processors : List[str]) -> Tuple[gradio.Dropdown, gradio.Dropdown]: + has_face_swapper = 'face_swapper' in processors + return gradio.Dropdown(visible = has_face_swapper), gradio.Dropdown(visible = has_face_swapper) + + +def update_face_swapper_model(face_swapper_model : FaceSwapperModel) -> Tuple[gradio.Dropdown, gradio.Dropdown]: + 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_choices = processors_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_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_choices) + return gradio.Dropdown(), gradio.Dropdown() + + +def update_face_swapper_pixel_boost(face_swapper_pixel_boost : str) -> None: + state_manager.set_item('face_swapper_pixel_boost', face_swapper_pixel_boost) diff --git a/facefusion/uis/components/frame_colorizer_options.py b/facefusion/uis/components/frame_colorizer_options.py new file mode 100755 index 00000000..f038392b --- /dev/null +++ b/facefusion/uis/components/frame_colorizer_options.py @@ -0,0 +1,77 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, wording +from facefusion.common_helper import calc_int_step +from facefusion.processors import choices as processors_choices +from facefusion.processors.core import load_processor_module +from facefusion.processors.typing import FrameColorizerModel +from facefusion.uis.core import get_ui_component, register_ui_component + +FRAME_COLORIZER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FRAME_COLORIZER_BLEND_SLIDER : Optional[gradio.Slider] = None +FRAME_COLORIZER_SIZE_DROPDOWN : Optional[gradio.Dropdown] = None + + +def render() -> None: + global FRAME_COLORIZER_MODEL_DROPDOWN + global FRAME_COLORIZER_BLEND_SLIDER + global FRAME_COLORIZER_SIZE_DROPDOWN + + FRAME_COLORIZER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.frame_colorizer_model_dropdown'), + choices = processors_choices.frame_colorizer_models, + value = state_manager.get_item('frame_colorizer_model'), + visible = 'frame_colorizer' in state_manager.get_item('processors') + ) + FRAME_COLORIZER_BLEND_SLIDER = gradio.Slider( + label = wording.get('uis.frame_colorizer_blend_slider'), + value = state_manager.get_item('frame_colorizer_blend'), + step = calc_int_step(processors_choices.frame_colorizer_blend_range), + minimum = processors_choices.frame_colorizer_blend_range[0], + maximum = processors_choices.frame_colorizer_blend_range[-1], + visible = 'frame_colorizer' in state_manager.get_item('processors') + ) + FRAME_COLORIZER_SIZE_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.frame_colorizer_size_dropdown'), + choices = processors_choices.frame_colorizer_sizes, + value = state_manager.get_item('frame_colorizer_size'), + visible = 'frame_colorizer' in state_manager.get_item('processors') + ) + register_ui_component('frame_colorizer_model_dropdown', FRAME_COLORIZER_MODEL_DROPDOWN) + register_ui_component('frame_colorizer_blend_slider', FRAME_COLORIZER_BLEND_SLIDER) + register_ui_component('frame_colorizer_size_dropdown', FRAME_COLORIZER_SIZE_DROPDOWN) + + +def listen() -> None: + FRAME_COLORIZER_MODEL_DROPDOWN.change(update_frame_colorizer_model, inputs = FRAME_COLORIZER_MODEL_DROPDOWN, outputs = FRAME_COLORIZER_MODEL_DROPDOWN) + FRAME_COLORIZER_BLEND_SLIDER.release(update_frame_colorizer_blend, inputs = FRAME_COLORIZER_BLEND_SLIDER) + FRAME_COLORIZER_SIZE_DROPDOWN.change(update_frame_colorizer_size, inputs = FRAME_COLORIZER_SIZE_DROPDOWN) + + 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_blend(frame_colorizer_blend : float) -> None: + state_manager.set_item('frame_colorizer_blend', int(frame_colorizer_blend)) + + +def update_frame_colorizer_size(frame_colorizer_size : str) -> None: + state_manager.set_item('frame_colorizer_size', frame_colorizer_size) diff --git a/facefusion/uis/components/frame_enhancer_options.py b/facefusion/uis/components/frame_enhancer_options.py new file mode 100755 index 00000000..f8629ef7 --- /dev/null +++ b/facefusion/uis/components/frame_enhancer_options.py @@ -0,0 +1,63 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import state_manager, wording +from facefusion.common_helper import calc_int_step +from facefusion.processors import choices as processors_choices +from facefusion.processors.core import load_processor_module +from facefusion.processors.typing 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 + + FRAME_ENHANCER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.frame_enhancer_model_dropdown'), + choices = processors_choices.frame_enhancer_models, + value = state_manager.get_item('frame_enhancer_model'), + visible = 'frame_enhancer' in state_manager.get_item('processors') + ) + FRAME_ENHANCER_BLEND_SLIDER = gradio.Slider( + label = wording.get('uis.frame_enhancer_blend_slider'), + value = state_manager.get_item('frame_enhancer_blend'), + step = calc_int_step(processors_choices.frame_enhancer_blend_range), + minimum = processors_choices.frame_enhancer_blend_range[0], + maximum = processors_choices.frame_enhancer_blend_range[-1], + visible = 'frame_enhancer' in state_manager.get_item('processors') + ) + 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/frame_processors.py b/facefusion/uis/components/frame_processors.py deleted file mode 100644 index 4195c63b..00000000 --- a/facefusion/uis/components/frame_processors.py +++ /dev/null @@ -1,40 +0,0 @@ -from typing import List, Optional -import gradio - -import facefusion.globals -from facefusion import wording -from facefusion.processors.frame.core import load_frame_processor_module, clear_frame_processors_modules -from facefusion.filesystem import list_directory -from facefusion.uis.core import register_ui_component - -FRAME_PROCESSORS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None - - -def render() -> None: - global FRAME_PROCESSORS_CHECKBOX_GROUP - - FRAME_PROCESSORS_CHECKBOX_GROUP = gradio.CheckboxGroup( - label = wording.get('uis.frame_processors_checkbox_group'), - choices = sort_frame_processors(facefusion.globals.frame_processors), - value = facefusion.globals.frame_processors - ) - register_ui_component('frame_processors_checkbox_group', FRAME_PROCESSORS_CHECKBOX_GROUP) - - -def listen() -> None: - FRAME_PROCESSORS_CHECKBOX_GROUP.change(update_frame_processors, inputs = FRAME_PROCESSORS_CHECKBOX_GROUP, outputs = FRAME_PROCESSORS_CHECKBOX_GROUP) - - -def update_frame_processors(frame_processors : List[str]) -> gradio.CheckboxGroup: - facefusion.globals.frame_processors = frame_processors - clear_frame_processors_modules() - for frame_processor in frame_processors: - frame_processor_module = load_frame_processor_module(frame_processor) - if not frame_processor_module.pre_check(): - return gradio.CheckboxGroup() - return gradio.CheckboxGroup(value = facefusion.globals.frame_processors, choices = sort_frame_processors(facefusion.globals.frame_processors)) - - -def sort_frame_processors(frame_processors : List[str]) -> list[str]: - available_frame_processors = list_directory('facefusion/processors/frame/modules') - return sorted(available_frame_processors, key = lambda frame_processor : frame_processors.index(frame_processor) if frame_processor in frame_processors else len(frame_processors)) diff --git a/facefusion/uis/components/frame_processors_options.py b/facefusion/uis/components/frame_processors_options.py deleted file mode 100755 index a705b53b..00000000 --- a/facefusion/uis/components/frame_processors_options.py +++ /dev/null @@ -1,216 +0,0 @@ -from typing import List, Optional, Tuple -import gradio - -import facefusion.globals -from facefusion import face_analyser, wording -from facefusion.processors.frame.core import load_frame_processor_module -from facefusion.processors.frame import globals as frame_processors_globals, choices as frame_processors_choices -from facefusion.processors.frame.typings import FaceDebuggerItem, FaceEnhancerModel, FaceSwapperModel, FrameColorizerModel, FrameEnhancerModel, LipSyncerModel -from facefusion.uis.core import get_ui_component, register_ui_component - -FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None -FACE_ENHANCER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None -FACE_ENHANCER_BLEND_SLIDER : Optional[gradio.Slider] = None -FACE_SWAPPER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None -FRAME_COLORIZER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None -FRAME_COLORIZER_BLEND_SLIDER : Optional[gradio.Slider] = None -FRAME_COLORIZER_SIZE_DROPDOWN : Optional[gradio.Dropdown] = None -FRAME_ENHANCER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None -FRAME_ENHANCER_BLEND_SLIDER : Optional[gradio.Slider] = None -LIP_SYNCER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None - - -def render() -> None: - global FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP - global FACE_ENHANCER_MODEL_DROPDOWN - global FACE_ENHANCER_BLEND_SLIDER - global FACE_SWAPPER_MODEL_DROPDOWN - global FRAME_COLORIZER_MODEL_DROPDOWN - global FRAME_COLORIZER_BLEND_SLIDER - global FRAME_COLORIZER_SIZE_DROPDOWN - global FRAME_ENHANCER_MODEL_DROPDOWN - global FRAME_ENHANCER_BLEND_SLIDER - global LIP_SYNCER_MODEL_DROPDOWN - - FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP = gradio.CheckboxGroup( - label = wording.get('uis.face_debugger_items_checkbox_group'), - choices = frame_processors_choices.face_debugger_items, - value = frame_processors_globals.face_debugger_items, - visible = 'face_debugger' in facefusion.globals.frame_processors - ) - FACE_ENHANCER_MODEL_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.face_enhancer_model_dropdown'), - choices = frame_processors_choices.face_enhancer_models, - value = frame_processors_globals.face_enhancer_model, - visible = 'face_enhancer' in facefusion.globals.frame_processors - ) - FACE_ENHANCER_BLEND_SLIDER = gradio.Slider( - label = wording.get('uis.face_enhancer_blend_slider'), - value = frame_processors_globals.face_enhancer_blend, - step = frame_processors_choices.face_enhancer_blend_range[1] - frame_processors_choices.face_enhancer_blend_range[0], - minimum = frame_processors_choices.face_enhancer_blend_range[0], - maximum = frame_processors_choices.face_enhancer_blend_range[-1], - visible = 'face_enhancer' in facefusion.globals.frame_processors - ) - FACE_SWAPPER_MODEL_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.face_swapper_model_dropdown'), - choices = frame_processors_choices.face_swapper_models, - value = frame_processors_globals.face_swapper_model, - visible = 'face_swapper' in facefusion.globals.frame_processors - ) - FRAME_COLORIZER_MODEL_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.frame_colorizer_model_dropdown'), - choices = frame_processors_choices.frame_colorizer_models, - value = frame_processors_globals.frame_colorizer_model, - visible = 'frame_colorizer' in facefusion.globals.frame_processors - ) - FRAME_COLORIZER_BLEND_SLIDER = gradio.Slider( - label = wording.get('uis.frame_colorizer_blend_slider'), - value = frame_processors_globals.frame_colorizer_blend, - step = frame_processors_choices.frame_colorizer_blend_range[1] - frame_processors_choices.frame_colorizer_blend_range[0], - minimum = frame_processors_choices.frame_colorizer_blend_range[0], - maximum = frame_processors_choices.frame_colorizer_blend_range[-1], - visible = 'frame_colorizer' in facefusion.globals.frame_processors - ) - FRAME_COLORIZER_SIZE_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.frame_colorizer_size_dropdown'), - choices = frame_processors_choices.frame_colorizer_sizes, - value = frame_processors_globals.frame_colorizer_size, - visible = 'frame_colorizer' in facefusion.globals.frame_processors - ) - FRAME_ENHANCER_MODEL_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.frame_enhancer_model_dropdown'), - choices = frame_processors_choices.frame_enhancer_models, - value = frame_processors_globals.frame_enhancer_model, - visible = 'frame_enhancer' in facefusion.globals.frame_processors - ) - FRAME_ENHANCER_BLEND_SLIDER = gradio.Slider( - label = wording.get('uis.frame_enhancer_blend_slider'), - value = frame_processors_globals.frame_enhancer_blend, - step = frame_processors_choices.frame_enhancer_blend_range[1] - frame_processors_choices.frame_enhancer_blend_range[0], - minimum = frame_processors_choices.frame_enhancer_blend_range[0], - maximum = frame_processors_choices.frame_enhancer_blend_range[-1], - visible = 'frame_enhancer' in facefusion.globals.frame_processors - ) - LIP_SYNCER_MODEL_DROPDOWN = gradio.Dropdown( - label = wording.get('uis.lip_syncer_model_dropdown'), - choices = frame_processors_choices.lip_syncer_models, - value = frame_processors_globals.lip_syncer_model, - visible = 'lip_syncer' in facefusion.globals.frame_processors - ) - register_ui_component('face_debugger_items_checkbox_group', FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP) - 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_swapper_model_dropdown', FACE_SWAPPER_MODEL_DROPDOWN) - register_ui_component('frame_colorizer_model_dropdown', FRAME_COLORIZER_MODEL_DROPDOWN) - register_ui_component('frame_colorizer_blend_slider', FRAME_COLORIZER_BLEND_SLIDER) - register_ui_component('frame_colorizer_size_dropdown', FRAME_COLORIZER_SIZE_DROPDOWN) - register_ui_component('frame_enhancer_model_dropdown', FRAME_ENHANCER_MODEL_DROPDOWN) - register_ui_component('frame_enhancer_blend_slider', FRAME_ENHANCER_BLEND_SLIDER) - register_ui_component('lip_syncer_model_dropdown', LIP_SYNCER_MODEL_DROPDOWN) - - -def listen() -> None: - FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP.change(update_face_debugger_items, inputs = FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP) - FACE_ENHANCER_MODEL_DROPDOWN.change(update_face_enhancer_model, inputs = FACE_ENHANCER_MODEL_DROPDOWN, outputs = FACE_ENHANCER_MODEL_DROPDOWN) - FACE_ENHANCER_BLEND_SLIDER.release(update_face_enhancer_blend, inputs = FACE_ENHANCER_BLEND_SLIDER) - FACE_SWAPPER_MODEL_DROPDOWN.change(update_face_swapper_model, inputs = FACE_SWAPPER_MODEL_DROPDOWN, outputs = FACE_SWAPPER_MODEL_DROPDOWN) - FRAME_COLORIZER_MODEL_DROPDOWN.change(update_frame_colorizer_model, inputs = FRAME_COLORIZER_MODEL_DROPDOWN, outputs = FRAME_COLORIZER_MODEL_DROPDOWN) - FRAME_COLORIZER_BLEND_SLIDER.release(update_frame_colorizer_blend, inputs = FRAME_COLORIZER_BLEND_SLIDER) - FRAME_COLORIZER_SIZE_DROPDOWN.change(update_frame_colorizer_size, inputs = FRAME_COLORIZER_SIZE_DROPDOWN, outputs = FRAME_COLORIZER_SIZE_DROPDOWN) - 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) - LIP_SYNCER_MODEL_DROPDOWN.change(update_lip_syncer_model, inputs = LIP_SYNCER_MODEL_DROPDOWN, outputs = LIP_SYNCER_MODEL_DROPDOWN) - frame_processors_checkbox_group = get_ui_component('frame_processors_checkbox_group') - if frame_processors_checkbox_group: - frame_processors_checkbox_group.change(update_frame_processors, inputs = frame_processors_checkbox_group, outputs = [ FACE_DEBUGGER_ITEMS_CHECKBOX_GROUP, FACE_ENHANCER_MODEL_DROPDOWN, FACE_ENHANCER_BLEND_SLIDER, FACE_SWAPPER_MODEL_DROPDOWN, FRAME_COLORIZER_MODEL_DROPDOWN, FRAME_COLORIZER_BLEND_SLIDER, FRAME_COLORIZER_SIZE_DROPDOWN, FRAME_ENHANCER_MODEL_DROPDOWN, FRAME_ENHANCER_BLEND_SLIDER, LIP_SYNCER_MODEL_DROPDOWN ]) - - -def update_frame_processors(frame_processors : List[str]) -> Tuple[gradio.CheckboxGroup, gradio.Dropdown, gradio.Slider, gradio.Dropdown, gradio.Dropdown, gradio.Slider, gradio.Dropdown, gradio.Dropdown, gradio.Slider, gradio.Dropdown]: - has_face_debugger = 'face_debugger' in frame_processors - has_face_enhancer = 'face_enhancer' in frame_processors - has_face_swapper = 'face_swapper' in frame_processors - has_frame_colorizer = 'frame_colorizer' in frame_processors - has_frame_enhancer = 'frame_enhancer' in frame_processors - has_lip_syncer = 'lip_syncer' in frame_processors - return gradio.CheckboxGroup(visible = has_face_debugger), gradio.Dropdown(visible = has_face_enhancer), gradio.Slider(visible = has_face_enhancer), gradio.Dropdown(visible = has_face_swapper), gradio.Dropdown(visible = has_frame_colorizer), gradio.Slider(visible = has_frame_colorizer), gradio.Dropdown(visible = has_frame_colorizer), gradio.Dropdown(visible = has_frame_enhancer), gradio.Slider(visible = has_frame_enhancer), gradio.Dropdown(visible = has_lip_syncer) - - -def update_face_debugger_items(face_debugger_items : List[FaceDebuggerItem]) -> None: - frame_processors_globals.face_debugger_items = face_debugger_items - - -def update_face_enhancer_model(face_enhancer_model : FaceEnhancerModel) -> gradio.Dropdown: - frame_processors_globals.face_enhancer_model = face_enhancer_model - face_enhancer_module = load_frame_processor_module('face_enhancer') - face_enhancer_module.clear_frame_processor() - face_enhancer_module.set_options('model', face_enhancer_module.MODELS[face_enhancer_model]) - if face_enhancer_module.pre_check(): - return gradio.Dropdown(value = frame_processors_globals.face_enhancer_model) - return gradio.Dropdown() - - -def update_face_enhancer_blend(face_enhancer_blend : int) -> None: - frame_processors_globals.face_enhancer_blend = face_enhancer_blend - - -def update_face_swapper_model(face_swapper_model : FaceSwapperModel) -> gradio.Dropdown: - frame_processors_globals.face_swapper_model = face_swapper_model - if face_swapper_model == 'blendswap_256': - facefusion.globals.face_recognizer_model = 'arcface_blendswap' - if face_swapper_model == 'inswapper_128' or face_swapper_model == 'inswapper_128_fp16': - facefusion.globals.face_recognizer_model = 'arcface_inswapper' - if face_swapper_model == 'simswap_256' or face_swapper_model == 'simswap_512_unofficial': - facefusion.globals.face_recognizer_model = 'arcface_simswap' - if face_swapper_model == 'uniface_256': - facefusion.globals.face_recognizer_model = 'arcface_uniface' - face_swapper_module = load_frame_processor_module('face_swapper') - face_swapper_module.clear_model_initializer() - face_swapper_module.clear_frame_processor() - face_swapper_module.set_options('model', face_swapper_module.MODELS[face_swapper_model]) - if face_analyser.pre_check() and face_swapper_module.pre_check(): - return gradio.Dropdown(value = frame_processors_globals.face_swapper_model) - return gradio.Dropdown() - - -def update_frame_colorizer_model(frame_colorizer_model : FrameColorizerModel) -> gradio.Dropdown: - frame_processors_globals.frame_colorizer_model = frame_colorizer_model - frame_colorizer_module = load_frame_processor_module('frame_colorizer') - frame_colorizer_module.clear_frame_processor() - frame_colorizer_module.set_options('model', frame_colorizer_module.MODELS[frame_colorizer_model]) - if frame_colorizer_module.pre_check(): - return gradio.Dropdown(value = frame_processors_globals.frame_colorizer_model) - return gradio.Dropdown() - - -def update_frame_colorizer_blend(frame_colorizer_blend : int) -> None: - frame_processors_globals.frame_colorizer_blend = frame_colorizer_blend - - -def update_frame_colorizer_size(frame_colorizer_size : str) -> gradio.Dropdown: - frame_processors_globals.frame_colorizer_size = frame_colorizer_size - return gradio.Dropdown(value = frame_processors_globals.frame_colorizer_size) - - -def update_frame_enhancer_model(frame_enhancer_model : FrameEnhancerModel) -> gradio.Dropdown: - frame_processors_globals.frame_enhancer_model = frame_enhancer_model - frame_enhancer_module = load_frame_processor_module('frame_enhancer') - frame_enhancer_module.clear_frame_processor() - frame_enhancer_module.set_options('model', frame_enhancer_module.MODELS[frame_enhancer_model]) - if frame_enhancer_module.pre_check(): - return gradio.Dropdown(value = frame_processors_globals.frame_enhancer_model) - return gradio.Dropdown() - - -def update_frame_enhancer_blend(frame_enhancer_blend : int) -> None: - frame_processors_globals.frame_enhancer_blend = frame_enhancer_blend - - -def update_lip_syncer_model(lip_syncer_model : LipSyncerModel) -> gradio.Dropdown: - frame_processors_globals.lip_syncer_model = lip_syncer_model - lip_syncer_module = load_frame_processor_module('lip_syncer') - lip_syncer_module.clear_frame_processor() - lip_syncer_module.set_options('model', lip_syncer_module.MODELS[lip_syncer_model]) - if lip_syncer_module.pre_check(): - return gradio.Dropdown(value = frame_processors_globals.lip_syncer_model) - return gradio.Dropdown() diff --git a/facefusion/uis/components/instant_runner.py b/facefusion/uis/components/instant_runner.py new file mode 100644 index 00000000..1d3ab136 --- /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, wording +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.typing 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 = wording.get('uis.start_button'), + variant = 'primary', + size = 'sm' + ) + INSTANT_RUNNER_STOP_BUTTON = gradio.Button( + value = wording.get('uis.stop_button'), + variant = 'primary', + size = 'sm', + visible = False + ) + INSTANT_RUNNER_CLEAR_BUTTON = gradio.Button( + value = wording.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 ]) + 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 + + 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]: + process_manager.stop() + return gradio.Button(visible = True), gradio.Button(visible = False) + + +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 00000000..ae808881 --- /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, wording +from facefusion.common_helper import get_first +from facefusion.jobs import job_list, job_manager +from facefusion.typing 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 = wording.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 00000000..87636267 --- /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, wording +from facefusion.common_helper import get_first +from facefusion.jobs import job_manager +from facefusion.typing 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 = wording.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 00000000..893e433a --- /dev/null +++ b/facefusion/uis/components/job_manager.py @@ -0,0 +1,184 @@ +from typing import List, Optional, Tuple + +import gradio + +from facefusion import logger, state_manager, wording +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.typing import UiWorkflow +from facefusion.uis import choices as uis_choices +from facefusion.uis.core import get_ui_component +from facefusion.uis.typing 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 = wording.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 = wording.get('uis.job_manager_job_id_dropdown'), + max_lines = 1, + interactive = True + ) + JOB_MANAGER_JOB_ID_DROPDOWN = gradio.Dropdown( + label = wording.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 = wording.get('uis.job_manager_step_index_dropdown'), + choices = [ 'none' ], + value = 'none', + interactive = True, + visible = False + ) + JOB_MANAGER_APPLY_BUTTON = gradio.Button( + value = wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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(wording.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 00000000..540a1025 --- /dev/null +++ b/facefusion/uis/components/job_runner.py @@ -0,0 +1,136 @@ +from time import sleep +from typing import Optional, Tuple + +import gradio + +from facefusion import logger, process_manager, state_manager, wording +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.typing import UiWorkflow +from facefusion.uis import choices as uis_choices +from facefusion.uis.core import get_ui_component +from facefusion.uis.typing 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 = wording.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 = wording.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 = wording.get('uis.start_button'), + variant = 'primary', + size = 'sm' + ) + JOB_RUNNER_STOP_BUTTON = gradio.Button( + value = wording.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 + + if job_action == 'job-run': + logger.info(wording.get('running_job').format(job_id = job_id), __name__) + if job_id and job_runner.run_job(job_id, process_step): + logger.info(wording.get('processing_job_succeed').format(job_id = job_id), __name__) + else: + logger.info(wording.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(wording.get('running_jobs'), __name__) + if job_runner.run_jobs(process_step): + logger.info(wording.get('processing_jobs_succeed'), __name__) + else: + logger.info(wording.get('processing_jobs_failed'), __name__) + if job_action == 'job-retry': + logger.info(wording.get('retrying_job').format(job_id = job_id), __name__) + if job_id and job_runner.retry_job(job_id, process_step): + logger.info(wording.get('processing_job_succeed').format(job_id = job_id), __name__) + else: + logger.info(wording.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(wording.get('retrying_jobs'), __name__) + if job_runner.retry_jobs(process_step): + logger.info(wording.get('processing_jobs_succeed'), __name__) + else: + logger.info(wording.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 100755 index 00000000..7c195d93 --- /dev/null +++ b/facefusion/uis/components/lip_syncer_options.py @@ -0,0 +1,46 @@ +from typing import List, Optional + +import gradio + +from facefusion import state_manager, wording +from facefusion.processors import choices as processors_choices +from facefusion.processors.core import load_processor_module +from facefusion.processors.typing import LipSyncerModel +from facefusion.uis.core import get_ui_component, register_ui_component + +LIP_SYNCER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None + + +def render() -> None: + global LIP_SYNCER_MODEL_DROPDOWN + + LIP_SYNCER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.lip_syncer_model_dropdown'), + choices = processors_choices.lip_syncer_models, + value = state_manager.get_item('lip_syncer_model'), + visible = 'lip_syncer' in state_manager.get_item('processors') + ) + register_ui_component('lip_syncer_model_dropdown', LIP_SYNCER_MODEL_DROPDOWN) + + +def listen() -> None: + LIP_SYNCER_MODEL_DROPDOWN.change(update_lip_syncer_model, inputs = LIP_SYNCER_MODEL_DROPDOWN, outputs = LIP_SYNCER_MODEL_DROPDOWN) + + 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) + + +def remote_update(processors : List[str]) -> gradio.Dropdown: + has_lip_syncer = 'lip_syncer' in processors + return gradio.Dropdown(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() diff --git a/facefusion/uis/components/memory.py b/facefusion/uis/components/memory.py index f67c27ae..1c461621 100644 --- a/facefusion/uis/components/memory.py +++ b/facefusion/uis/components/memory.py @@ -1,10 +1,11 @@ from typing import Optional + import gradio -import facefusion.globals import facefusion.choices +from facefusion import state_manager, wording +from facefusion.common_helper import calc_int_step from facefusion.typing import VideoMemoryStrategy -from facefusion import wording VIDEO_MEMORY_STRATEGY_DROPDOWN : Optional[gradio.Dropdown] = None SYSTEM_MEMORY_LIMIT_SLIDER : Optional[gradio.Slider] = None @@ -17,14 +18,14 @@ def render() -> None: VIDEO_MEMORY_STRATEGY_DROPDOWN = gradio.Dropdown( label = wording.get('uis.video_memory_strategy_dropdown'), choices = facefusion.choices.video_memory_strategies, - value = facefusion.globals.video_memory_strategy + value = state_manager.get_item('video_memory_strategy') ) SYSTEM_MEMORY_LIMIT_SLIDER = gradio.Slider( label = wording.get('uis.system_memory_limit_slider'), - step =facefusion.choices.system_memory_limit_range[1] - facefusion.choices.system_memory_limit_range[0], + step = calc_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 = facefusion.globals.system_memory_limit + value = state_manager.get_item('system_memory_limit') ) @@ -34,8 +35,8 @@ def listen() -> None: def update_video_memory_strategy(video_memory_strategy : VideoMemoryStrategy) -> None: - facefusion.globals.video_memory_strategy = video_memory_strategy + state_manager.set_item('video_memory_strategy', video_memory_strategy) -def update_system_memory_limit(system_memory_limit : int) -> None: - facefusion.globals.system_memory_limit = system_memory_limit +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 index cfba2a6d..84fd0891 100644 --- a/facefusion/uis/components/output.py +++ b/facefusion/uis/components/output.py @@ -1,29 +1,28 @@ -from typing import Tuple, Optional -from time import sleep +import tempfile +from typing import Optional + import gradio -import facefusion.globals -from facefusion import process_manager, wording -from facefusion.core import conditional_process -from facefusion.memory import limit_system_memory -from facefusion.normalizer import normalize_output_path -from facefusion.uis.core import get_ui_component -from facefusion.filesystem import clear_temp, is_image, is_video +from facefusion import state_manager, wording +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 -OUTPUT_START_BUTTON : Optional[gradio.Button] = None -OUTPUT_CLEAR_BUTTON : Optional[gradio.Button] = None -OUTPUT_STOP_BUTTON : Optional[gradio.Button] = None def render() -> None: + global OUTPUT_PATH_TEXTBOX global OUTPUT_IMAGE global OUTPUT_VIDEO - global OUTPUT_START_BUTTON - global OUTPUT_STOP_BUTTON - global OUTPUT_CLEAR_BUTTON + if not state_manager.get_item('output_path'): + state_manager.set_item('output_path', tempfile.gettempdir()) + OUTPUT_PATH_TEXTBOX = gradio.Textbox( + label = wording.get('uis.output_path_textbox'), + value = state_manager.get_item('output_path'), + max_lines = 1 + ) OUTPUT_IMAGE = gradio.Image( label = wording.get('uis.output_image_or_video'), visible = False @@ -31,58 +30,13 @@ def render() -> None: OUTPUT_VIDEO = gradio.Video( label = wording.get('uis.output_image_or_video') ) - OUTPUT_START_BUTTON = gradio.Button( - value = wording.get('uis.start_button'), - variant = 'primary', - size = 'sm' - ) - OUTPUT_STOP_BUTTON = gradio.Button( - value = wording.get('uis.stop_button'), - variant = 'primary', - size = 'sm', - visible = False - ) - OUTPUT_CLEAR_BUTTON = gradio.Button( - value = wording.get('uis.clear_button'), - size = 'sm' - ) def listen() -> None: - output_path_textbox = get_ui_component('output_path_textbox') - if output_path_textbox: - OUTPUT_START_BUTTON.click(start, outputs = [ OUTPUT_START_BUTTON, OUTPUT_STOP_BUTTON ]) - OUTPUT_START_BUTTON.click(process, outputs = [ OUTPUT_IMAGE, OUTPUT_VIDEO, OUTPUT_START_BUTTON, OUTPUT_STOP_BUTTON ]) - OUTPUT_STOP_BUTTON.click(stop, outputs = [ OUTPUT_START_BUTTON, OUTPUT_STOP_BUTTON ]) - OUTPUT_CLEAR_BUTTON.click(clear, outputs = [ OUTPUT_IMAGE, OUTPUT_VIDEO ]) + 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 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 process() -> Tuple[gradio.Image, gradio.Video, gradio.Button, gradio.Button]: - normed_output_path = normalize_output_path(facefusion.globals.target_path, facefusion.globals.output_path) - if facefusion.globals.system_memory_limit > 0: - limit_system_memory(facefusion.globals.system_memory_limit) - conditional_process() - if is_image(normed_output_path): - return gradio.Image(value = normed_output_path, visible = True), gradio.Video(value = None, visible = False), gradio.Button(visible = True), gradio.Button(visible = False) - if is_video(normed_output_path): - return gradio.Image(value = None, visible = False), gradio.Video(value = normed_output_path, visible = True), gradio.Button(visible = True), gradio.Button(visible = False) - return gradio.Image(value = None), gradio.Video(value = None), gradio.Button(visible = True), gradio.Button(visible = False) - - -def stop() -> Tuple[gradio.Button, gradio.Button]: - process_manager.stop() - return gradio.Button(visible = True), gradio.Button(visible = False) - - -def clear() -> Tuple[gradio.Image, gradio.Video]: - while process_manager.is_processing(): - sleep(0.5) - if facefusion.globals.target_path: - clear_temp(facefusion.globals.target_path) - return gradio.Image(value = None), gradio.Video(value = None) +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 index 4919920a..31fe154f 100644 --- a/facefusion/uis/components/output_options.py +++ b/facefusion/uis/components/output_options.py @@ -1,17 +1,18 @@ from typing import Optional, Tuple + import gradio -import facefusion.globals import facefusion.choices -from facefusion import wording -from facefusion.typing import OutputVideoEncoder, OutputVideoPreset, Fps +from facefusion import state_manager, wording +from facefusion.common_helper import calc_int_step from facefusion.filesystem import is_image, is_video +from facefusion.typing import Fps, OutputAudioEncoder, OutputVideoEncoder, OutputVideoPreset from facefusion.uis.core import get_ui_components, register_ui_component -from facefusion.vision import detect_image_resolution, create_image_resolutions, detect_video_fps, detect_video_resolution, create_video_resolutions, pack_resolution +from facefusion.vision import create_image_resolutions, create_video_resolutions, detect_image_resolution, detect_video_fps, detect_video_resolution, pack_resolution -OUTPUT_PATH_TEXTBOX : Optional[gradio.Textbox] = None OUTPUT_IMAGE_QUALITY_SLIDER : Optional[gradio.Slider] = None OUTPUT_IMAGE_RESOLUTION_DROPDOWN : Optional[gradio.Dropdown] = None +OUTPUT_AUDIO_ENCODER_DROPDOWN : Optional[gradio.Dropdown] = None OUTPUT_VIDEO_ENCODER_DROPDOWN : Optional[gradio.Dropdown] = None OUTPUT_VIDEO_PRESET_DROPDOWN : Optional[gradio.Dropdown] = None OUTPUT_VIDEO_RESOLUTION_DROPDOWN : Optional[gradio.Dropdown] = None @@ -20,9 +21,9 @@ OUTPUT_VIDEO_FPS_SLIDER : Optional[gradio.Slider] = None def render() -> None: - global OUTPUT_PATH_TEXTBOX global OUTPUT_IMAGE_QUALITY_SLIDER global OUTPUT_IMAGE_RESOLUTION_DROPDOWN + global OUTPUT_AUDIO_ENCODER_DROPDOWN global OUTPUT_VIDEO_ENCODER_DROPDOWN global OUTPUT_VIDEO_PRESET_DROPDOWN global OUTPUT_VIDEO_RESOLUTION_DROPDOWN @@ -31,74 +32,73 @@ def render() -> None: output_image_resolutions = [] output_video_resolutions = [] - if is_image(facefusion.globals.target_path): - output_image_resolution = detect_image_resolution(facefusion.globals.target_path) + if is_image(state_manager.get_item('target_path')): + output_image_resolution = detect_image_resolution(state_manager.get_item('target_path')) output_image_resolutions = create_image_resolutions(output_image_resolution) - if is_video(facefusion.globals.target_path): - output_video_resolution = detect_video_resolution(facefusion.globals.target_path) + if is_video(state_manager.get_item('target_path')): + output_video_resolution = detect_video_resolution(state_manager.get_item('target_path')) output_video_resolutions = create_video_resolutions(output_video_resolution) - facefusion.globals.output_path = facefusion.globals.output_path or '.' - OUTPUT_PATH_TEXTBOX = gradio.Textbox( - label = wording.get('uis.output_path_textbox'), - value = facefusion.globals.output_path, - max_lines = 1 - ) OUTPUT_IMAGE_QUALITY_SLIDER = gradio.Slider( label = wording.get('uis.output_image_quality_slider'), - value = facefusion.globals.output_image_quality, - step = facefusion.choices.output_image_quality_range[1] - facefusion.choices.output_image_quality_range[0], + value = state_manager.get_item('output_image_quality'), + step = calc_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(facefusion.globals.target_path) + visible = is_image(state_manager.get_item('target_path')) ) OUTPUT_IMAGE_RESOLUTION_DROPDOWN = gradio.Dropdown( label = wording.get('uis.output_image_resolution_dropdown'), choices = output_image_resolutions, - value = facefusion.globals.output_image_resolution, - visible = is_image(facefusion.globals.target_path) + value = state_manager.get_item('output_image_resolution'), + visible = is_image(state_manager.get_item('target_path')) + ) + OUTPUT_AUDIO_ENCODER_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.output_audio_encoder_dropdown'), + choices = facefusion.choices.output_audio_encoders, + value = state_manager.get_item('output_audio_encoder'), + visible = is_video(state_manager.get_item('target_path')) ) OUTPUT_VIDEO_ENCODER_DROPDOWN = gradio.Dropdown( label = wording.get('uis.output_video_encoder_dropdown'), choices = facefusion.choices.output_video_encoders, - value = facefusion.globals.output_video_encoder, - visible = is_video(facefusion.globals.target_path) + value = state_manager.get_item('output_video_encoder'), + visible = is_video(state_manager.get_item('target_path')) ) OUTPUT_VIDEO_PRESET_DROPDOWN = gradio.Dropdown( label = wording.get('uis.output_video_preset_dropdown'), choices = facefusion.choices.output_video_presets, - value = facefusion.globals.output_video_preset, - visible = is_video(facefusion.globals.target_path) + value = state_manager.get_item('output_video_preset'), + visible = is_video(state_manager.get_item('target_path')) ) OUTPUT_VIDEO_QUALITY_SLIDER = gradio.Slider( label = wording.get('uis.output_video_quality_slider'), - value = facefusion.globals.output_video_quality, - step = facefusion.choices.output_video_quality_range[1] - facefusion.choices.output_video_quality_range[0], + value = state_manager.get_item('output_video_quality'), + step = calc_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(facefusion.globals.target_path) + visible = is_video(state_manager.get_item('target_path')) ) OUTPUT_VIDEO_RESOLUTION_DROPDOWN = gradio.Dropdown( label = wording.get('uis.output_video_resolution_dropdown'), choices = output_video_resolutions, - value = facefusion.globals.output_video_resolution, - visible = is_video(facefusion.globals.target_path) + value = state_manager.get_item('output_video_resolution'), + visible = is_video(state_manager.get_item('target_path')) ) OUTPUT_VIDEO_FPS_SLIDER = gradio.Slider( label = wording.get('uis.output_video_fps_slider'), - value = facefusion.globals.output_video_fps, + value = state_manager.get_item('output_video_fps'), step = 0.01, minimum = 1, maximum = 60, - visible = is_video(facefusion.globals.target_path) + visible = is_video(state_manager.get_item('target_path')) ) - register_ui_component('output_path_textbox', OUTPUT_PATH_TEXTBOX) register_ui_component('output_video_fps_slider', OUTPUT_VIDEO_FPS_SLIDER) def listen() -> None: - OUTPUT_PATH_TEXTBOX.change(update_output_path, inputs = OUTPUT_PATH_TEXTBOX) OUTPUT_IMAGE_QUALITY_SLIDER.release(update_output_image_quality, inputs = OUTPUT_IMAGE_QUALITY_SLIDER) OUTPUT_IMAGE_RESOLUTION_DROPDOWN.change(update_output_image_resolution, inputs = OUTPUT_IMAGE_RESOLUTION_DROPDOWN) + OUTPUT_AUDIO_ENCODER_DROPDOWN.change(update_output_audio_encoder, inputs = OUTPUT_AUDIO_ENCODER_DROPDOWN) 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) @@ -111,51 +111,51 @@ def listen() -> None: 'target_video' ]): for method in [ 'upload', 'change', 'clear' ]: - getattr(ui_component, method)(remote_update, outputs = [ OUTPUT_IMAGE_QUALITY_SLIDER, OUTPUT_IMAGE_RESOLUTION_DROPDOWN, OUTPUT_VIDEO_ENCODER_DROPDOWN, OUTPUT_VIDEO_PRESET_DROPDOWN, OUTPUT_VIDEO_QUALITY_SLIDER, OUTPUT_VIDEO_RESOLUTION_DROPDOWN, OUTPUT_VIDEO_FPS_SLIDER ]) + getattr(ui_component, method)(remote_update, outputs = [ OUTPUT_IMAGE_QUALITY_SLIDER, OUTPUT_IMAGE_RESOLUTION_DROPDOWN, OUTPUT_AUDIO_ENCODER_DROPDOWN, OUTPUT_VIDEO_ENCODER_DROPDOWN, OUTPUT_VIDEO_PRESET_DROPDOWN, OUTPUT_VIDEO_QUALITY_SLIDER, OUTPUT_VIDEO_RESOLUTION_DROPDOWN, OUTPUT_VIDEO_FPS_SLIDER ]) -def remote_update() -> Tuple[gradio.Slider, gradio.Dropdown, gradio.Dropdown, gradio.Dropdown, gradio.Slider, gradio.Dropdown, gradio.Slider]: - if is_image(facefusion.globals.target_path): - output_image_resolution = detect_image_resolution(facefusion.globals.target_path) +def remote_update() -> Tuple[gradio.Slider, gradio.Dropdown, gradio.Dropdown, gradio.Dropdown, gradio.Dropdown, gradio.Slider, gradio.Dropdown, gradio.Slider]: + if is_image(state_manager.get_item('target_path')): + output_image_resolution = detect_image_resolution(state_manager.get_item('target_path')) output_image_resolutions = create_image_resolutions(output_image_resolution) - facefusion.globals.output_image_resolution = pack_resolution(output_image_resolution) - return gradio.Slider(visible = True), gradio.Dropdown(visible = True, value = facefusion.globals.output_image_resolution, choices = output_image_resolutions), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False), gradio.Dropdown(visible = False, value = None, choices = None), gradio.Slider(visible = False, value = None) - if is_video(facefusion.globals.target_path): - output_video_resolution = detect_video_resolution(facefusion.globals.target_path) + state_manager.set_item('output_image_resolution', pack_resolution(output_image_resolution)) + return gradio.Slider(visible = True), gradio.Dropdown(value = state_manager.get_item('output_image_resolution'), choices = output_image_resolutions, visible = True), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False) + if is_video(state_manager.get_item('target_path')): + output_video_resolution = detect_video_resolution(state_manager.get_item('target_path')) output_video_resolutions = create_video_resolutions(output_video_resolution) - facefusion.globals.output_video_resolution = pack_resolution(output_video_resolution) - facefusion.globals.output_video_fps = detect_video_fps(facefusion.globals.target_path) - return gradio.Slider(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = True), gradio.Dropdown(visible = True), gradio.Slider(visible = True), gradio.Dropdown(visible = True, value = facefusion.globals.output_video_resolution, choices = output_video_resolutions), gradio.Slider(visible = True, value = facefusion.globals.output_video_fps) - return gradio.Slider(visible = False), gradio.Dropdown(visible = False, value = None, choices = None), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False), gradio.Dropdown(visible = False, value = None, choices = None), gradio.Slider(visible = False, value = None) + state_manager.set_item('output_video_resolution', pack_resolution(output_video_resolution)) + state_manager.set_item('output_video_fps', detect_video_fps(state_manager.get_item('target_path'))) + return gradio.Slider(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = True), gradio.Dropdown(visible = True), gradio.Dropdown(visible = True), gradio.Slider(visible = True), gradio.Dropdown(value = state_manager.get_item('output_video_resolution'), choices = output_video_resolutions, visible = True), gradio.Slider(value = state_manager.get_item('output_video_fps'), visible = True) + return gradio.Slider(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False), gradio.Dropdown(visible = False), gradio.Slider(visible = False) -def update_output_path(output_path : str) -> None: - facefusion.globals.output_path = output_path - - -def update_output_image_quality(output_image_quality : int) -> None: - facefusion.globals.output_image_quality = output_image_quality +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_resolution(output_image_resolution : str) -> None: - facefusion.globals.output_image_resolution = output_image_resolution + state_manager.set_item('output_image_resolution', output_image_resolution) -def update_output_video_encoder(output_video_encoder: OutputVideoEncoder) -> None: - facefusion.globals.output_video_encoder = output_video_encoder +def update_output_audio_encoder(output_audio_encoder : OutputAudioEncoder) -> None: + state_manager.set_item('output_audio_encoder', output_audio_encoder) + + +def update_output_video_encoder(output_video_encoder : OutputVideoEncoder) -> None: + state_manager.set_item('output_video_encoder', output_video_encoder) def update_output_video_preset(output_video_preset : OutputVideoPreset) -> None: - facefusion.globals.output_video_preset = output_video_preset + state_manager.set_item('output_video_preset', output_video_preset) -def update_output_video_quality(output_video_quality : int) -> None: - facefusion.globals.output_video_quality = output_video_quality +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_resolution(output_video_resolution : str) -> None: - facefusion.globals.output_video_resolution = output_video_resolution + state_manager.set_item('output_video_resolution', output_video_resolution) def update_output_video_fps(output_video_fps : Fps) -> None: - facefusion.globals.output_video_fps = output_video_fps + state_manager.set_item('output_video_fps', output_video_fps) diff --git a/facefusion/uis/components/preview.py b/facefusion/uis/components/preview.py index 72108cf4..6b3eb0a7 100755 --- a/facefusion/uis/components/preview.py +++ b/facefusion/uis/components/preview.py @@ -1,22 +1,23 @@ -from typing import Any, Dict, Optional from time import sleep +from typing import Optional + import cv2 import gradio import numpy -import facefusion.globals -from facefusion import logger, wording -from facefusion.audio import get_audio_frame, create_empty_audio_frame +from facefusion import logger, process_manager, state_manager, wording +from facefusion.audio import create_empty_audio_frame, get_audio_frame from facefusion.common_helper import get_first -from facefusion.core import conditional_append_reference_faces -from facefusion.face_analyser import get_average_face, clear_face_analyser -from facefusion.face_store import clear_static_faces, get_reference_faces, clear_reference_faces -from facefusion.typing import Face, FaceSet, AudioFrame, VisionFrame -from facefusion.vision import get_video_frame, count_video_frame_total, normalize_frame_color, resize_frame_resolution, read_static_image, read_static_images -from facefusion.filesystem import is_image, is_video, filter_audio_paths from facefusion.content_analyser import analyse_frame -from facefusion.processors.frame.core import load_frame_processor_module +from facefusion.core import conditional_append_reference_faces +from facefusion.face_analyser import get_average_face, get_many_faces +from facefusion.face_store import clear_reference_faces, clear_static_faces, get_reference_faces +from facefusion.filesystem import filter_audio_paths, is_image, is_video +from facefusion.processors.core import get_processors_modules +from facefusion.typing import AudioFrame, Face, FaceSet, VisionFrame from facefusion.uis.core import get_ui_component, get_ui_components, register_ui_component +from facefusion.uis.typing import ComponentOptions +from facefusion.vision import count_video_frame_total, detect_frame_orientation, get_video_frame, normalize_frame_color, read_static_image, read_static_images, resize_frame_resolution PREVIEW_IMAGE : Optional[gradio.Image] = None PREVIEW_FRAME_SLIDER : Optional[gradio.Slider] = None @@ -26,12 +27,11 @@ def render() -> None: global PREVIEW_IMAGE global PREVIEW_FRAME_SLIDER - preview_image_args : Dict[str, Any] =\ + preview_image_options : ComponentOptions =\ { - 'label': wording.get('uis.preview_image'), - 'interactive': False + 'label': wording.get('uis.preview_image') } - preview_frame_slider_args : Dict[str, Any] =\ + preview_frame_slider_options : ComponentOptions =\ { 'label': wording.get('uis.preview_frame_slider'), 'step': 1, @@ -40,34 +40,42 @@ def render() -> None: 'visible': False } conditional_append_reference_faces() - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - source_frames = read_static_images(facefusion.globals.source_paths) - source_face = get_average_face(source_frames) - source_audio_path = get_first(filter_audio_paths(facefusion.globals.source_paths)) + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + source_frames = read_static_images(state_manager.get_item('source_paths')) + source_faces = get_many_faces(source_frames) + source_face = get_average_face(source_faces) + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) source_audio_frame = create_empty_audio_frame() - if source_audio_path and facefusion.globals.output_video_fps and facefusion.globals.reference_frame_number: - temp_audio_frame = get_audio_frame(source_audio_path, facefusion.globals.output_video_fps, facefusion.globals.reference_frame_number) + + if source_audio_path and state_manager.get_item('output_video_fps') and state_manager.get_item('reference_frame_number'): + temp_audio_frame = get_audio_frame(source_audio_path, state_manager.get_item('output_video_fps'), state_manager.get_item('reference_frame_number')) if numpy.any(temp_audio_frame): source_audio_frame = temp_audio_frame - if is_image(facefusion.globals.target_path): - target_vision_frame = read_static_image(facefusion.globals.target_path) + + if is_image(state_manager.get_item('target_path')): + target_vision_frame = read_static_image(state_manager.get_item('target_path')) preview_vision_frame = process_preview_frame(reference_faces, source_face, source_audio_frame, target_vision_frame) - preview_image_args['value'] = normalize_frame_color(preview_vision_frame) - if is_video(facefusion.globals.target_path): - temp_vision_frame = get_video_frame(facefusion.globals.target_path, facefusion.globals.reference_frame_number) + preview_image_options['value'] = normalize_frame_color(preview_vision_frame) + 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 = get_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) preview_vision_frame = process_preview_frame(reference_faces, source_face, source_audio_frame, temp_vision_frame) - preview_image_args['value'] = normalize_frame_color(preview_vision_frame) - preview_image_args['visible'] = True - preview_frame_slider_args['value'] = facefusion.globals.reference_frame_number - preview_frame_slider_args['maximum'] = count_video_frame_total(facefusion.globals.target_path) - preview_frame_slider_args['visible'] = True - PREVIEW_IMAGE = gradio.Image(**preview_image_args) - PREVIEW_FRAME_SLIDER = gradio.Slider(**preview_frame_slider_args) + preview_image_options['value'] = normalize_frame_color(preview_vision_frame) + preview_image_options['elem_classes'] = [ 'image-preview', 'is-' + detect_frame_orientation(preview_vision_frame) ] + preview_image_options['visible'] = True + 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_IMAGE = gradio.Image(**preview_image_options) + PREVIEW_FRAME_SLIDER = gradio.Slider(**preview_frame_slider_options) register_ui_component('preview_frame_slider', PREVIEW_FRAME_SLIDER) def listen() -> None: - PREVIEW_FRAME_SLIDER.release(update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE) + PREVIEW_FRAME_SLIDER.release(update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE, show_progress = 'hidden') + PREVIEW_FRAME_SLIDER.change(slide_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE, show_progress = 'hidden') + reference_face_position_gallery = get_ui_component('reference_face_position_gallery') if reference_face_position_gallery: reference_face_position_gallery.select(update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE) @@ -94,23 +102,34 @@ def listen() -> None: [ 'face_debugger_items_checkbox_group', 'frame_colorizer_size_dropdown', - 'face_selector_mode_dropdown', 'face_mask_types_checkbox_group', - 'face_mask_region_checkbox_group', - 'face_analyser_order_dropdown', - 'face_analyser_age_dropdown', - 'face_analyser_gender_dropdown' + 'face_mask_region_checkbox_group' ]): ui_component.change(update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE) for ui_component in get_ui_components( [ + 'age_modifier_direction_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', 'frame_colorizer_blend_slider', 'frame_enhancer_blend_slider', - 'trim_frame_start_slider', - 'trim_frame_end_slider', 'reference_face_distance_slider', + 'face_selector_age_range_slider', 'face_mask_blur_slider', 'face_mask_padding_top_slider', 'face_mask_padding_bottom_slider', @@ -122,14 +141,24 @@ def listen() -> None: for ui_component in get_ui_components( [ - 'frame_processors_checkbox_group', + 'age_modifier_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_size_dropdown', + 'face_detector_angles_checkbox_group', + 'face_landmarker_model_dropdown' ]): ui_component.change(clear_and_update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE) @@ -142,66 +171,75 @@ def listen() -> None: def clear_and_update_preview_image(frame_number : int = 0) -> gradio.Image: - clear_face_analyser() clear_reference_faces() clear_static_faces() return update_preview_image(frame_number) -def update_preview_image(frame_number : int = 0) -> gradio.Image: - for frame_processor in facefusion.globals.frame_processors: - frame_processor_module = load_frame_processor_module(frame_processor) - while not frame_processor_module.post_check(): - logger.disable() - sleep(0.5) - logger.enable() - conditional_append_reference_faces() - reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None - source_frames = read_static_images(facefusion.globals.source_paths) - source_face = get_average_face(source_frames) - source_audio_path = get_first(filter_audio_paths(facefusion.globals.source_paths)) - source_audio_frame = create_empty_audio_frame() - if source_audio_path and facefusion.globals.output_video_fps and facefusion.globals.reference_frame_number: - reference_audio_frame_number = facefusion.globals.reference_frame_number - if facefusion.globals.trim_frame_start: - reference_audio_frame_number -= facefusion.globals.trim_frame_start - temp_audio_frame = get_audio_frame(source_audio_path, facefusion.globals.output_video_fps, reference_audio_frame_number) - if numpy.any(temp_audio_frame): - source_audio_frame = temp_audio_frame - if is_image(facefusion.globals.target_path): - target_vision_frame = read_static_image(facefusion.globals.target_path) - preview_vision_frame = process_preview_frame(reference_faces, source_face, source_audio_frame, target_vision_frame) - preview_vision_frame = normalize_frame_color(preview_vision_frame) - return gradio.Image(value = preview_vision_frame) - if is_video(facefusion.globals.target_path): - temp_vision_frame = get_video_frame(facefusion.globals.target_path, frame_number) - preview_vision_frame = process_preview_frame(reference_faces, source_face, source_audio_frame, temp_vision_frame) - preview_vision_frame = normalize_frame_color(preview_vision_frame) +def slide_preview_image(frame_number : int = 0) -> gradio.Image: + if is_video(state_manager.get_item('target_path')): + preview_vision_frame = normalize_frame_color(get_video_frame(state_manager.get_item('target_path'), frame_number)) + preview_vision_frame = resize_frame_resolution(preview_vision_frame, (1024, 1024)) return gradio.Image(value = preview_vision_frame) return gradio.Image(value = None) +def update_preview_image(frame_number : int = 0) -> gradio.Image: + while process_manager.is_checking(): + sleep(0.5) + conditional_append_reference_faces() + reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None + source_frames = read_static_images(state_manager.get_item('source_paths')) + source_faces = get_many_faces(source_frames) + source_face = get_average_face(source_faces) + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) + source_audio_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_audio_frame = get_audio_frame(source_audio_path, state_manager.get_item('output_video_fps'), reference_audio_frame_number) + if numpy.any(temp_audio_frame): + source_audio_frame = temp_audio_frame + + if is_image(state_manager.get_item('target_path')): + target_vision_frame = read_static_image(state_manager.get_item('target_path')) + preview_vision_frame = process_preview_frame(reference_faces, source_face, source_audio_frame, target_vision_frame) + preview_vision_frame = normalize_frame_color(preview_vision_frame) + 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')): + temp_vision_frame = get_video_frame(state_manager.get_item('target_path'), frame_number) + preview_vision_frame = process_preview_frame(reference_faces, source_face, source_audio_frame, temp_vision_frame) + preview_vision_frame = normalize_frame_color(preview_vision_frame) + 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 update_preview_frame_slider() -> gradio.Slider: - if is_video(facefusion.globals.target_path): - video_frame_total = count_video_frame_total(facefusion.globals.target_path) + 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 = None, maximum = None, visible = False) + return gradio.Slider(value = 0, visible = False) def process_preview_frame(reference_faces : FaceSet, source_face : Face, source_audio_frame : AudioFrame, target_vision_frame : VisionFrame) -> VisionFrame: - target_vision_frame = resize_frame_resolution(target_vision_frame, (640, 640)) + target_vision_frame = resize_frame_resolution(target_vision_frame, (1024, 1024)) + source_vision_frame = target_vision_frame.copy() if analyse_frame(target_vision_frame): return cv2.GaussianBlur(target_vision_frame, (99, 99), 0) - for frame_processor in facefusion.globals.frame_processors: - frame_processor_module = load_frame_processor_module(frame_processor) + + for processor_module in get_processors_modules(state_manager.get_item('processors')): logger.disable() - if frame_processor_module.pre_process('preview'): - logger.enable() - target_vision_frame = frame_processor_module.process_frame( + if processor_module.pre_process('preview'): + target_vision_frame = processor_module.process_frame( { 'reference_faces': reference_faces, 'source_face': source_face, 'source_audio_frame': source_audio_frame, + 'source_vision_frame': source_vision_frame, 'target_vision_frame': target_vision_frame }) + logger.enable() return target_vision_frame diff --git a/facefusion/uis/components/processors.py b/facefusion/uis/components/processors.py new file mode 100644 index 00000000..567f3e11 --- /dev/null +++ b/facefusion/uis/components/processors.py @@ -0,0 +1,40 @@ +from typing import List, Optional + +import gradio + +from facefusion import state_manager, wording +from facefusion.filesystem import list_directory +from facefusion.processors.core import clear_processors_modules, 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 = wording.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: + clear_processors_modules(state_manager.get_item('processors')) + state_manager.set_item('processors', processors) + + for processor_module in get_processors_modules(state_manager.get_item('processors')): + if not processor_module.pre_check(): + return gradio.CheckboxGroup() + 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 = list_directory('facefusion/processors/modules') + return sorted(available_processors, key = lambda processor : processors.index(processor) if processor in processors else len(processors)) diff --git a/facefusion/uis/components/source.py b/facefusion/uis/components/source.py index 4e79d89d..4f9c6f24 100644 --- a/facefusion/uis/components/source.py +++ b/facefusion/uis/components/source.py @@ -1,12 +1,12 @@ -from typing import Optional, List, Tuple +from typing import List, Optional, Tuple + import gradio -import facefusion.globals -from facefusion import wording -from facefusion.uis.typing import File +from facefusion import state_manager, wording from facefusion.common_helper import get_first -from facefusion.filesystem import has_audio, has_image, filter_audio_paths, filter_image_paths +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.typing import File SOURCE_FILE : Optional[gradio.File] = None SOURCE_AUDIO : Optional[gradio.Audio] = None @@ -18,22 +18,19 @@ def render() -> None: global SOURCE_AUDIO global SOURCE_IMAGE - has_source_audio = has_audio(facefusion.globals.source_paths) - has_source_image = has_image(facefusion.globals.source_paths) + 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( file_count = 'multiple', file_types = [ - '.mp3', - '.wav', - '.png', - '.jpg', - '.webp' + 'audio', + 'image' ], label = wording.get('uis.source_file'), - value = facefusion.globals.source_paths if has_source_audio or has_source_image else None + value = state_manager.get_item('source_paths') if has_source_audio or has_source_image else None ) - source_file_names = [ source_file_value['name'] for source_file_value in SOURCE_FILE.value ] if SOURCE_FILE.value 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( @@ -61,7 +58,7 @@ def update(files : List[File]) -> Tuple[gradio.Audio, gradio.Image]: 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)) - facefusion.globals.source_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) - facefusion.globals.source_paths = None + 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 index 4b133d8a..63a0d70d 100644 --- a/facefusion/uis/components/target.py +++ b/facefusion/uis/components/target.py @@ -1,12 +1,12 @@ -from typing import Tuple, Optional +from typing import Optional, Tuple + import gradio -import facefusion.globals -from facefusion import wording -from facefusion.face_store import clear_static_faces, clear_reference_faces -from facefusion.uis.typing import File +from facefusion import state_manager, wording +from facefusion.face_store import clear_reference_faces, clear_static_faces from facefusion.filesystem import get_file_size, is_image, is_video from facefusion.uis.core import register_ui_component +from facefusion.uis.typing import ComponentOptions, File from facefusion.vision import get_video_frame, normalize_frame_color FILE_SIZE_LIMIT = 512 * 1024 * 1024 @@ -21,44 +21,41 @@ def render() -> None: global TARGET_IMAGE global TARGET_VIDEO - is_target_image = is_image(facefusion.globals.target_path) - is_target_video = is_video(facefusion.globals.target_path) + 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 = wording.get('uis.target_file'), file_count = 'single', file_types = [ - '.png', - '.jpg', - '.webp', - '.webm', - '.mp4' + 'image', + 'video' ], - value = facefusion.globals.target_path if is_target_image or is_target_video else None + value = state_manager.get_item('target_path') if is_target_image or is_target_video else None ) - target_image_args =\ + target_image_options : ComponentOptions =\ { 'show_label': False, 'visible': False } - target_video_args =\ + target_video_options : ComponentOptions =\ { 'show_label': False, 'visible': False } if is_target_image: - target_image_args['value'] = TARGET_FILE.value['name'] - target_image_args['visible'] = True + target_image_options['value'] = TARGET_FILE.value.get('path') + target_image_options['visible'] = True if is_target_video: - if get_file_size(facefusion.globals.target_path) > FILE_SIZE_LIMIT: - preview_vision_frame = normalize_frame_color(get_video_frame(facefusion.globals.target_path)) - target_image_args['value'] = preview_vision_frame - target_image_args['visible'] = True + if get_file_size(state_manager.get_item('target_path')) > FILE_SIZE_LIMIT: + preview_vision_frame = normalize_frame_color(get_video_frame(state_manager.get_item('target_path'))) + target_image_options['value'] = preview_vision_frame + target_image_options['visible'] = True else: - target_video_args['value'] = TARGET_FILE.value['name'] - target_video_args['visible'] = True - TARGET_IMAGE = gradio.Image(**target_image_args) - TARGET_VIDEO = gradio.Video(**target_video_args) + 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) @@ -71,13 +68,13 @@ def update(file : File) -> Tuple[gradio.Image, gradio.Video]: clear_reference_faces() clear_static_faces() if file and is_image(file.name): - facefusion.globals.target_path = 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): - facefusion.globals.target_path = file.name + state_manager.set_item('target_path', file.name) if get_file_size(file.name) > FILE_SIZE_LIMIT: preview_vision_frame = normalize_frame_color(get_video_frame(file.name)) return gradio.Image(value = preview_vision_frame, visible = True), gradio.Video(value = None, visible = False) return gradio.Image(value = None, visible = False), gradio.Video(value = file.name, visible = True) - facefusion.globals.target_path = None + 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 index a601653f..c1a98389 100644 --- a/facefusion/uis/components/temp_frame.py +++ b/facefusion/uis/components/temp_frame.py @@ -1,11 +1,11 @@ from typing import Optional + import gradio -import facefusion.globals import facefusion.choices -from facefusion import wording -from facefusion.typing import TempFrameFormat +from facefusion import state_manager, wording from facefusion.filesystem import is_video +from facefusion.typing import TempFrameFormat from facefusion.uis.core import get_ui_component TEMP_FRAME_FORMAT_DROPDOWN : Optional[gradio.Dropdown] = None @@ -17,13 +17,14 @@ def render() -> None: TEMP_FRAME_FORMAT_DROPDOWN = gradio.Dropdown( label = wording.get('uis.temp_frame_format_dropdown'), choices = facefusion.choices.temp_frame_formats, - value = facefusion.globals.temp_frame_format, - visible = is_video(facefusion.globals.target_path) + 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 [ 'upload', 'change', 'clear' ]: @@ -31,11 +32,11 @@ def listen() -> None: def remote_update() -> gradio.Dropdown: - if is_video(facefusion.globals.target_path): + 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: - facefusion.globals.temp_frame_format = temp_frame_format + 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 00000000..bbe2a35b --- /dev/null +++ b/facefusion/uis/components/terminal.py @@ -0,0 +1,82 @@ +import io +import logging +import math +import os +from typing import Optional + +import gradio +from tqdm import tqdm + +from facefusion import logger, state_manager, wording +from facefusion.choices import log_level_set +from facefusion.typing 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 = wording.get('uis.log_level_dropdown'), + choices = log_level_set.keys(), + value = state_manager.get_item('log_level') + ) + TERMINAL_TEXTBOX = gradio.Textbox( + label = wording.get('uis.terminal_textbox'), + value = read_logs, + lines = 8, + max_lines = 8, + every = 0.5, + show_copy_button = True + ) + + +def listen() -> None: + global LOG_LEVEL_DROPDOWN + + 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 + wording.get('colon') + ' ' + str(percentage) + '% (' + str(self.n) + '/' + str(self.total) + ')' + if not self.disable and self.desc and self.unit: + return self.desc + wording.get('colon') + ' ' + str(self.n) + ' ' + self.unit + return None + + +def read_logs() -> str: + LOG_BUFFER.seek(0) + logs = LOG_BUFFER.read().rstrip() + return logs diff --git a/facefusion/uis/components/trim_frame.py b/facefusion/uis/components/trim_frame.py index c5c6c40d..c264f2a9 100644 --- a/facefusion/uis/components/trim_frame.py +++ b/facefusion/uis/components/trim_frame.py @@ -1,79 +1,62 @@ -from typing import Any, Dict, Tuple, Optional -import gradio +from typing import Optional, Tuple -import facefusion.globals -from facefusion import wording +from gradio_rangeslider import RangeSlider + +from facefusion import state_manager, wording from facefusion.face_store import clear_static_faces -from facefusion.vision import count_video_frame_total from facefusion.filesystem import is_video -from facefusion.uis.core import get_ui_components, register_ui_component +from facefusion.uis.core import get_ui_components +from facefusion.uis.typing import ComponentOptions +from facefusion.vision import count_video_frame_total -TRIM_FRAME_START_SLIDER : Optional[gradio.Slider] = None -TRIM_FRAME_END_SLIDER : Optional[gradio.Slider] = None +TRIM_FRAME_RANGE_SLIDER : Optional[RangeSlider] = None def render() -> None: - global TRIM_FRAME_START_SLIDER - global TRIM_FRAME_END_SLIDER + global TRIM_FRAME_RANGE_SLIDER - trim_frame_start_slider_args : Dict[str, Any] =\ + trim_frame_range_slider_options : ComponentOptions =\ { - 'label': wording.get('uis.trim_frame_start_slider'), - 'step': 1, + 'label': wording.get('uis.trim_frame_slider'), 'minimum': 0, - 'maximum': 100, + 'step': 1, 'visible': False } - trim_frame_end_slider_args : Dict[str, Any] =\ - { - 'label': wording.get('uis.trim_frame_end_slider'), - 'step': 1, - 'minimum': 0, - 'maximum': 100, - 'visible': False - } - if is_video(facefusion.globals.target_path): - video_frame_total = count_video_frame_total(facefusion.globals.target_path) - trim_frame_start_slider_args['value'] = facefusion.globals.trim_frame_start or 0 - trim_frame_start_slider_args['maximum'] = video_frame_total - trim_frame_start_slider_args['visible'] = True - trim_frame_end_slider_args['value'] = facefusion.globals.trim_frame_end or video_frame_total - trim_frame_end_slider_args['maximum'] = video_frame_total - trim_frame_end_slider_args['visible'] = True - with gradio.Row(): - TRIM_FRAME_START_SLIDER = gradio.Slider(**trim_frame_start_slider_args) - TRIM_FRAME_END_SLIDER = gradio.Slider(**trim_frame_end_slider_args) - register_ui_component('trim_frame_start_slider', TRIM_FRAME_START_SLIDER) - register_ui_component('trim_frame_end_slider', TRIM_FRAME_END_SLIDER) + 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_START_SLIDER.release(update_trim_frame_start, inputs = TRIM_FRAME_START_SLIDER) - TRIM_FRAME_END_SLIDER.release(update_trim_frame_end, inputs = TRIM_FRAME_END_SLIDER) + 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 [ 'upload', 'change', 'clear' ]: - getattr(ui_component, method)(remote_update, outputs = [ TRIM_FRAME_START_SLIDER, TRIM_FRAME_END_SLIDER ]) + getattr(ui_component, method)(remote_update, outputs = [ TRIM_FRAME_RANGE_SLIDER ]) -def remote_update() -> Tuple[gradio.Slider, gradio.Slider]: - if is_video(facefusion.globals.target_path): - video_frame_total = count_video_frame_total(facefusion.globals.target_path) - facefusion.globals.trim_frame_start = None - facefusion.globals.trim_frame_end = None - return gradio.Slider(value = 0, maximum = video_frame_total, visible = True), gradio.Slider(value = video_frame_total, maximum = video_frame_total, visible = True) - return gradio.Slider(value = None, maximum = None, visible = False), gradio.Slider(value = None, maximum = None, visible = False) +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_start(trim_frame_start : int) -> None: +def update_trim_frame(trim_frame : Tuple[float, float]) -> None: clear_static_faces() - facefusion.globals.trim_frame_start = trim_frame_start if trim_frame_start > 0 else None - - -def update_trim_frame_end(trim_frame_end : int) -> None: - clear_static_faces() - video_frame_total = count_video_frame_total(facefusion.globals.target_path) - facefusion.globals.trim_frame_end = trim_frame_end if trim_frame_end < video_frame_total else None + 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 00000000..47711a38 --- /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, wording +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 = wording.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/webcam.py b/facefusion/uis/components/webcam.py index 8571d9e1..cc3fcebf 100644 --- a/facefusion/uis/components/webcam.py +++ b/facefusion/uis/components/webcam.py @@ -1,26 +1,25 @@ -from typing import Optional, Generator, Deque import os import subprocess +from collections import deque +from concurrent.futures import ThreadPoolExecutor +from typing import Deque, Generator, Optional + import cv2 import gradio -from time import sleep -from concurrent.futures import ThreadPoolExecutor -from collections import deque from tqdm import tqdm -import facefusion.globals -from facefusion import logger, wording +from facefusion import logger, state_manager, wording from facefusion.audio import create_empty_audio_frame from facefusion.common_helper import is_windows from facefusion.content_analyser import analyse_stream -from facefusion.filesystem import filter_image_paths -from facefusion.typing import VisionFrame, Face, Fps -from facefusion.face_analyser import get_average_face -from facefusion.processors.frame.core import get_frame_processors_modules, load_frame_processor_module +from facefusion.face_analyser import get_average_face, get_many_faces from facefusion.ffmpeg import open_ffmpeg -from facefusion.vision import normalize_frame_color, read_static_images, unpack_resolution +from facefusion.filesystem import filter_image_paths +from facefusion.processors.core import get_processors_modules +from facefusion.typing import Face, Fps, VisionFrame +from facefusion.uis.core import get_ui_component from facefusion.uis.typing import StreamMode, WebcamMode -from facefusion.uis.core import get_ui_component, get_ui_components +from facefusion.vision import normalize_frame_color, read_static_images, unpack_resolution WEBCAM_CAPTURE : Optional[cv2.VideoCapture] = None WEBCAM_IMAGE : Optional[gradio.Image] = None @@ -69,32 +68,21 @@ def render() -> None: def listen() -> None: - start_event = None 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_mode_radio and webcam_resolution_dropdown and webcam_fps_slider: start_event = WEBCAM_START_BUTTON.click(start, inputs = [ webcam_mode_radio, webcam_resolution_dropdown, webcam_fps_slider ], outputs = WEBCAM_IMAGE) - WEBCAM_STOP_BUTTON.click(stop, cancels = start_event) - - for ui_component in get_ui_components( - [ - 'frame_processors_checkbox_group', - 'face_swapper_model_dropdown', - 'face_enhancer_model_dropdown', - 'frame_enhancer_model_dropdown', - 'lip_syncer_model_dropdown', - 'source_image' - ]): - ui_component.change(update, cancels = start_event) + WEBCAM_STOP_BUTTON.click(stop, cancels = start_event) def start(webcam_mode : WebcamMode, webcam_resolution : str, webcam_fps : Fps) -> Generator[VisionFrame, None, None]: - facefusion.globals.face_selector_mode = 'one' - facefusion.globals.face_analyser_order = 'large-small' - source_image_paths = filter_image_paths(facefusion.globals.source_paths) + state_manager.set_item('face_selector_mode', 'one') + source_image_paths = filter_image_paths(state_manager.get_item('source_paths')) source_frames = read_static_images(source_image_paths) - source_face = get_average_face(source_frames) + source_faces = get_many_faces(source_frames) + source_face = get_average_face(source_faces) stream = None if webcam_mode in [ 'udp', 'v4l2' ]: @@ -118,34 +106,33 @@ def start(webcam_mode : WebcamMode, webcam_resolution : str, webcam_fps : Fps) - def multi_process_capture(source_face : Face, webcam_capture : cv2.VideoCapture, webcam_fps : Fps) -> Generator[VisionFrame, None, None]: - with tqdm(desc = wording.get('processing'), unit = 'frame', ascii = ' =', disable = facefusion.globals.log_level in [ 'warn', 'error' ]) as progress: - with ThreadPoolExecutor(max_workers = facefusion.globals.execution_thread_count) as executor: + deque_capture_frames: Deque[VisionFrame] = deque() + with tqdm(desc = wording.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'), + 'execution_thread_count': state_manager.get_item('execution_thread_count') + }) + with ThreadPoolExecutor(max_workers = state_manager.get_item('execution_thread_count')) as executor: futures = [] - deque_capture_frames : Deque[VisionFrame] = deque() + while webcam_capture and webcam_capture.isOpened(): _, capture_frame = webcam_capture.read() if analyse_stream(capture_frame, webcam_fps): return future = executor.submit(process_stream_frame, source_face, capture_frame) futures.append(future) + for future_done in [ future for future in futures if future.done() ]: capture_frame = future_done.result() deque_capture_frames.append(capture_frame) futures.remove(future_done) + while deque_capture_frames: progress.update() yield deque_capture_frames.popleft() -def update() -> None: - for frame_processor in facefusion.globals.frame_processors: - frame_processor_module = load_frame_processor_module(frame_processor) - while not frame_processor_module.post_check(): - logger.disable() - sleep(0.5) - logger.enable() - - def stop() -> gradio.Image: clear_webcam_capture() return gradio.Image(value = None) @@ -153,16 +140,16 @@ def stop() -> gradio.Image: def process_stream_frame(source_face : Face, target_vision_frame : VisionFrame) -> VisionFrame: source_audio_frame = create_empty_audio_frame() - for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): + for processor_module in get_processors_modules(state_manager.get_item('processors')): logger.disable() - if frame_processor_module.pre_process('stream'): - logger.enable() - target_vision_frame = frame_processor_module.process_frame( + if processor_module.pre_process('stream'): + target_vision_frame = processor_module.process_frame( { 'source_face': source_face, 'source_audio_frame': source_audio_frame, 'target_vision_frame': target_vision_frame }) + logger.enable() return target_vision_frame @@ -176,5 +163,5 @@ def open_stream(stream_mode : StreamMode, stream_resolution : str, stream_fps : if device_name: commands.extend([ '-f', 'v4l2', '/dev/' + device_name ]) except FileNotFoundError: - logger.error(wording.get('stream_not_loaded').format(stream_mode = stream_mode), __name__.upper()) + logger.error(wording.get('stream_not_loaded').format(stream_mode = stream_mode), __name__) return open_ffmpeg(commands) diff --git a/facefusion/uis/components/webcam_options.py b/facefusion/uis/components/webcam_options.py index ea707b05..cbe7390c 100644 --- a/facefusion/uis/components/webcam_options.py +++ b/facefusion/uis/components/webcam_options.py @@ -1,4 +1,5 @@ from typing import Optional + import gradio from facefusion import wording diff --git a/facefusion/uis/core.py b/facefusion/uis/core.py index a664c714..e3b338f4 100644 --- a/facefusion/uis/core.py +++ b/facefusion/uis/core.py @@ -1,18 +1,22 @@ -from typing import Dict, Optional, Any, List -from types import ModuleType -import os import importlib -import sys -import gradio +import os +import warnings +from types import ModuleType +from typing import Any, Dict, List, Optional -import facefusion.globals -from facefusion.uis import overrides -from facefusion import metadata, logger, wording -from facefusion.uis.typing import Component, ComponentName +import gradio +from gradio.themes import Size + +from facefusion import logger, metadata, state_manager, wording +from facefusion.exit_helper import hard_exit from facefusion.filesystem import resolve_relative_path +from facefusion.uis import overrides +from facefusion.uis.typing import Component, ComponentName os.environ['GRADIO_ANALYTICS_ENABLED'] = '0' +warnings.filterwarnings('ignore', category = UserWarning, module = 'gradio') + gradio.processing_utils.encode_array_to_base64 = overrides.encode_array_to_base64 gradio.processing_utils.encode_pil_to_base64 = overrides.encode_pil_to_base64 @@ -35,12 +39,12 @@ def load_ui_layout_module(ui_layout : str) -> Any: if not hasattr(ui_layout_module, method_name): raise NotImplementedError except ModuleNotFoundError as exception: - logger.error(wording.get('ui_layout_not_loaded').format(ui_layout = ui_layout), __name__.upper()) - logger.debug(exception.msg, __name__.upper()) - sys.exit(1) + logger.error(wording.get('ui_layout_not_loaded').format(ui_layout = ui_layout), __name__) + logger.debug(exception.msg, __name__) + hard_exit(1) except NotImplementedError: - logger.error(wording.get('ui_layout_not_implemented').format(ui_layout = ui_layout), __name__.upper()) - sys.exit(1) + logger.error(wording.get('ui_layout_not_implemented').format(ui_layout = ui_layout), __name__) + hard_exit(1) return ui_layout_module @@ -75,9 +79,9 @@ def register_ui_component(component_name : ComponentName, component: Component) def launch() -> None: - ui_layouts_total = len(facefusion.globals.ui_layouts) - with gradio.Blocks(theme = get_theme(), css = get_css(), title = metadata.get('name') + ' ' + metadata.get('version')) as ui: - for ui_layout in facefusion.globals.ui_layouts: + 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_layout_module.pre_render(): if ui_layouts_total > 1: @@ -88,7 +92,7 @@ def launch() -> None: ui_layout_module.render() ui_layout_module.listen() - for ui_layout in facefusion.globals.ui_layouts: + for ui_layout in state_manager.get_item('ui_layouts'): ui_layout_module = load_ui_layout_module(ui_layout) ui_layout_module.run(ui) @@ -97,24 +101,34 @@ def get_theme() -> gradio.Theme: return gradio.themes.Base( primary_hue = gradio.themes.colors.red, secondary_hue = gradio.themes.colors.neutral, + 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( background_fill_primary = '*neutral_100', block_background_fill = 'white', block_border_width = '0', - block_label_background_fill = '*primary_100', - block_label_background_fill_dark = '*primary_600', + block_label_background_fill = '*neutral_100', + block_label_background_fill_dark = '*neutral_700', block_label_border_width = 'none', block_label_margin = '0.5rem', block_label_radius = '*radius_md', - block_label_text_color = '*primary_500', + 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 = '*primary_100', - block_title_background_fill_dark = '*primary_600', + block_title_background_fill = '*neutral_100', + block_title_background_fill_dark = '*neutral_700', block_title_padding = '*block_label_padding', block_title_radius = '*block_label_radius', - block_title_text_color = '*primary_500', + block_title_text_color = '*neutral_700', block_title_text_size = '*text_sm', block_title_text_weight = '600', block_padding = '0.5rem', @@ -151,6 +165,5 @@ def get_theme() -> gradio.Theme: def get_css() -> str: - fixes_css_path = resolve_relative_path('uis/assets/fixes.css') overrides_css_path = resolve_relative_path('uis/assets/overrides.css') - return open(fixes_css_path, 'r').read() + open(overrides_css_path, 'r').read() + return open(overrides_css_path, 'r').read() diff --git a/facefusion/uis/layouts/benchmark.py b/facefusion/uis/layouts/benchmark.py index 0a3324fe..7f501b81 100644 --- a/facefusion/uis/layouts/benchmark.py +++ b/facefusion/uis/layouts/benchmark.py @@ -1,24 +1,23 @@ -import multiprocessing import gradio -import facefusion.globals +from facefusion import state_manager from facefusion.download import conditional_download -from facefusion.uis.components import about, frame_processors, frame_processors_options, execution, execution_thread_count, execution_queue_count, memory, benchmark_options, benchmark +from facefusion.uis.components import about, age_modifier_options, benchmark, benchmark_options, execution, execution_queue_count, 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: - if not facefusion.globals.skip_download: + if not state_manager.get_item('skip_download'): conditional_download('.assets/examples', [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.mp3', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-360p.mp4', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-540p.mp4', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-720p.mp4', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-1080p.mp4', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-1440p.mp4', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-2160p.mp4' + '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', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-360p.mp4', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-540p.mp4', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-720p.mp4', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-1080p.mp4', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-1440p.mp4', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-2160p.mp4' ]) return True return False @@ -31,13 +30,29 @@ def pre_render() -> bool: def render() -> gradio.Blocks: with gradio.Blocks() as layout: with gradio.Row(): - with gradio.Column(scale = 2): + with gradio.Column(scale = 4): with gradio.Blocks(): about.render() with gradio.Blocks(): - frame_processors.render() + processors.render() with gradio.Blocks(): - frame_processors_options.render() + age_modifier_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() @@ -46,15 +61,23 @@ def render() -> gradio.Blocks: memory.render() with gradio.Blocks(): benchmark_options.render() - with gradio.Column(scale = 5): + with gradio.Column(scale = 11): with gradio.Blocks(): benchmark.render() return layout def listen() -> None: - frame_processors.listen() - frame_processors_options.listen() + processors.listen() + age_modifier_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() execution_queue_count.listen() @@ -63,5 +86,4 @@ def listen() -> None: def run(ui : gradio.Blocks) -> None: - concurrency_count = min(2, multiprocessing.cpu_count()) - ui.queue(concurrency_count = concurrency_count).launch(show_api = False, quiet = True, inbrowser = facefusion.globals.open_browser) + 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 index 67fff4fb..c5cf2d61 100755 --- a/facefusion/uis/layouts/default.py +++ b/facefusion/uis/layouts/default.py @@ -1,8 +1,7 @@ -import multiprocessing import gradio -import facefusion.globals -from facefusion.uis.components import about, frame_processors, frame_processors_options, execution, execution_thread_count, execution_queue_count, memory, temp_frame, output_options, common_options, source, target, output, preview, trim_frame, face_analyser, face_selector, face_masker +from facefusion import state_manager +from facefusion.uis.components import about, age_modifier_options, common_options, execution, execution_queue_count, 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, processors, source, target, temp_frame, terminal, trim_frame, ui_workflow def pre_check() -> bool: @@ -16,13 +15,29 @@ def pre_render() -> bool: def render() -> gradio.Blocks: with gradio.Blocks() as layout: with gradio.Row(): - with gradio.Column(scale = 2): + with gradio.Column(scale = 4): with gradio.Blocks(): about.render() with gradio.Blocks(): - frame_processors.render() + processors.render() with gradio.Blocks(): - frame_processors_options.render() + age_modifier_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() @@ -33,14 +48,21 @@ def render() -> gradio.Blocks: temp_frame.render() with gradio.Blocks(): output_options.render() - with gradio.Column(scale = 2): + with gradio.Column(scale = 4): with gradio.Blocks(): source.render() with gradio.Blocks(): target.render() with gradio.Blocks(): output.render() - with gradio.Column(scale = 3): + 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() with gradio.Blocks(): @@ -50,15 +72,25 @@ def render() -> gradio.Blocks: with gradio.Blocks(): face_masker.render() with gradio.Blocks(): - face_analyser.render() + face_detector.render() + with gradio.Blocks(): + face_landmarker.render() with gradio.Blocks(): common_options.render() return layout def listen() -> None: - frame_processors.listen() - frame_processors_options.listen() + processors.listen() + age_modifier_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() execution_queue_count.listen() @@ -68,14 +100,18 @@ def listen() -> None: source.listen() target.listen() output.listen() + instant_runner.listen() + job_runner.listen() + job_manager.listen() + terminal.listen() preview.listen() trim_frame.listen() face_selector.listen() face_masker.listen() - face_analyser.listen() + face_detector.listen() + face_landmarker.listen() common_options.listen() def run(ui : gradio.Blocks) -> None: - concurrency_count = min(8, multiprocessing.cpu_count()) - ui.queue(concurrency_count = concurrency_count).launch(show_api = False, quiet = True, inbrowser = facefusion.globals.open_browser) + 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 100755 index 00000000..740e31e8 --- /dev/null +++ b/facefusion/uis/layouts/jobs.py @@ -0,0 +1,35 @@ +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 pre_render() -> 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 index 424ab221..26f7f10f 100644 --- a/facefusion/uis/layouts/webcam.py +++ b/facefusion/uis/layouts/webcam.py @@ -1,8 +1,7 @@ -import multiprocessing import gradio -import facefusion.globals -from facefusion.uis.components import about, frame_processors, frame_processors_options, execution, execution_thread_count, webcam_options, source, webcam +from facefusion import state_manager +from facefusion.uis.components import about, age_modifier_options, execution, execution_thread_count, face_debugger_options, face_enhancer_options, face_swapper_options, frame_colorizer_options, frame_enhancer_options, lip_syncer_options, processors, source, webcam, webcam_options def pre_check() -> bool: @@ -16,13 +15,25 @@ def pre_render() -> bool: def render() -> gradio.Blocks: with gradio.Blocks() as layout: with gradio.Row(): - with gradio.Column(scale = 2): + with gradio.Column(scale = 4): with gradio.Blocks(): about.render() with gradio.Blocks(): - frame_processors.render() + processors.render() with gradio.Blocks(): - frame_processors_options.render() + age_modifier_options.render() + with gradio.Blocks(): + face_debugger_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() @@ -30,15 +41,21 @@ def render() -> gradio.Blocks: webcam_options.render() with gradio.Blocks(): source.render() - with gradio.Column(scale = 5): + with gradio.Column(scale = 11): with gradio.Blocks(): webcam.render() return layout def listen() -> None: - frame_processors.listen() - frame_processors_options.listen() + processors.listen() + age_modifier_options.listen() + face_debugger_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() source.listen() @@ -46,5 +63,4 @@ def listen() -> None: def run(ui : gradio.Blocks) -> None: - concurrency_count = min(2, multiprocessing.cpu_count()) - ui.queue(concurrency_count = concurrency_count).launch(show_api = False, quiet = True, inbrowser = facefusion.globals.open_browser) + 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 index 7f3c4707..1a1ee11d 100644 --- a/facefusion/uis/overrides.py +++ b/facefusion/uis/overrides.py @@ -1,11 +1,13 @@ +import base64 from typing import Any + import cv2 import numpy -import base64 +from numpy._typing import NDArray -def encode_array_to_base64(array : numpy.ndarray[Any, Any]) -> str: - buffer = cv2.imencode('.jpg', array[:, :, ::-1])[1] +def encode_array_to_base64(array : NDArray[Any]) -> str: + _, buffer = cv2.imencode('.jpg', array[:, :, ::-1]) return 'data:image/jpeg;base64,' + base64.b64encode(buffer.tobytes()).decode('utf-8') diff --git a/facefusion/uis/typing.py b/facefusion/uis/typing.py index 59d06f53..502b6a22 100644 --- a/facefusion/uis/typing.py +++ b/facefusion/uis/typing.py @@ -1,53 +1,80 @@ -from typing import Literal, Any, IO -import gradio +from typing import Any, Dict, IO, Literal File = IO[Any] -Component = gradio.File or gradio.Image or gradio.Video or gradio.Slider +Component = Any +ComponentOptions = Dict[str, Any] ComponentName = Literal\ [ + 'age_modifier_direction_slider', + 'age_modifier_model_dropdown', + 'expression_restorer_factor_slider', + 'expression_restorer_model_dropdown', + 'benchmark_cycles_slider', + 'benchmark_runs_checkbox_group', + 'face_debugger_items_checkbox_group', + 'face_detector_angles_checkbox_group', + 'face_detector_model_dropdown', + 'face_detector_score_slider', + 'face_detector_size_dropdown', + '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', + 'face_enhancer_blend_slider', + 'face_enhancer_model_dropdown', + 'face_landmarker_model_dropdown', + 'face_landmarker_score_slider', + '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_mask_region_checkbox_group', + 'face_mask_types_checkbox_group', + 'face_selector_gender_dropdown', + 'face_selector_race_dropdown', + 'face_selector_age_range_slider', + 'face_selector_mode_dropdown', + 'face_selector_order_dropdown', + 'face_swapper_model_dropdown', + 'face_swapper_pixel_boost_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', + 'output_image', + 'output_video', + 'output_video_fps_slider', + 'preview_frame_slider', + 'processors_checkbox_group', + 'reference_face_distance_slider', + 'reference_face_position_gallery', 'source_audio', 'source_image', 'target_image', 'target_video', - 'preview_frame_slider', - 'trim_frame_start_slider', - 'trim_frame_end_slider', - 'face_selector_mode_dropdown', - 'reference_face_position_gallery', - 'reference_face_distance_slider', - 'face_analyser_order_dropdown', - 'face_analyser_age_dropdown', - 'face_analyser_gender_dropdown', - 'face_detector_model_dropdown', - 'face_detector_size_dropdown', - 'face_detector_score_slider', - 'face_landmarker_score_slider', - 'face_mask_types_checkbox_group', - 'face_mask_blur_slider', - 'face_mask_padding_top_slider', - 'face_mask_padding_bottom_slider', - 'face_mask_padding_left_slider', - 'face_mask_padding_right_slider', - 'face_mask_region_checkbox_group', - 'frame_processors_checkbox_group', - 'face_debugger_items_checkbox_group', - 'face_enhancer_model_dropdown', - 'face_enhancer_blend_slider', - 'face_swapper_model_dropdown', - 'frame_colorizer_model_dropdown', - 'frame_colorizer_blend_slider', - 'frame_colorizer_size_dropdown', - 'frame_enhancer_model_dropdown', - 'frame_enhancer_blend_slider', - 'lip_syncer_model_dropdown', - 'output_path_textbox', - 'output_video_fps_slider', - 'benchmark_runs_checkbox_group', - 'benchmark_cycles_slider', + 'ui_workflow_dropdown', + 'webcam_fps_slider', 'webcam_mode_radio', - 'webcam_resolution_dropdown', - 'webcam_fps_slider' + 'webcam_resolution_dropdown' ] +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'] + WebcamMode = Literal['inline', 'udp', 'v4l2'] StreamMode = Literal['udp', 'v4l2'] diff --git a/facefusion/uis/ui_helper.py b/facefusion/uis/ui_helper.py new file mode 100644 index 00000000..91fa75da --- /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 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): + _, target_extension = os.path.splitext(target_path) + output_name = hashlib.sha1(str(state_manager.get_state()).encode('utf-8')).hexdigest()[:8] + return os.path.join(output_directory_path, output_name + target_extension) + return None diff --git a/facefusion/vision.py b/facefusion/vision.py index 692a73a1..d5c925db 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -1,13 +1,14 @@ -from typing import Optional, List, Tuple 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.typing import VisionFrame, Resolution, Fps from facefusion.choices import image_template_sizes, video_template_sizes +from facefusion.common_helper import is_windows from facefusion.filesystem import is_image, is_video, sanitize_path_for_windows +from facefusion.typing import Fps, Orientation, Resolution, VisionFrame @lru_cache(maxsize = 128) @@ -15,8 +16,9 @@ def read_static_image(image_path : str) -> Optional[VisionFrame]: return read_image(image_path) -def read_static_images(image_paths : List[str]) -> Optional[List[VisionFrame]]: +def read_static_images(image_paths : List[str]) -> List[VisionFrame]: frames = [] + if image_paths: for image_path in image_paths: frames.append(read_static_image(image_path)) @@ -176,6 +178,14 @@ def unpack_resolution(resolution : str) -> Resolution: 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 resize_frame_resolution(vision_frame : VisionFrame, max_resolution : Resolution) -> VisionFrame: height, width = vision_frame.shape[:2] max_width, max_height = max_resolution diff --git a/facefusion/voice_extractor.py b/facefusion/voice_extractor.py index 4404d5a6..60d8c4a6 100644 --- a/facefusion/voice_extractor.py +++ b/facefusion/voice_extractor.py @@ -1,56 +1,57 @@ -from typing import Any, Tuple -from time import sleep -import scipy +from typing import Tuple + import numpy -import onnxruntime +import scipy -import facefusion.globals -from facefusion import process_manager -from facefusion.thread_helper import thread_lock, thread_semaphore -from facefusion.typing import ModelSet, AudioChunk, Audio -from facefusion.execution import apply_execution_provider_options -from facefusion.filesystem import resolve_relative_path, is_file -from facefusion.download import conditional_download +from facefusion import inference_manager +from facefusion.download import conditional_download_hashes, conditional_download_sources +from facefusion.filesystem import resolve_relative_path +from facefusion.thread_helper import thread_semaphore +from facefusion.typing import Audio, AudioChunk, InferencePool, ModelOptions, ModelSet -VOICE_EXTRACTOR = None -MODELS : ModelSet =\ +MODEL_SET : ModelSet =\ { - 'voice_extractor': + 'kim_vocal_2': { - 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/voice_extractor.onnx', - 'path': resolve_relative_path('../.assets/models/voice_extractor.onnx') + 'hashes': + { + 'voice_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/kim_vocal_2.hash', + 'path': resolve_relative_path('../.assets/models/kim_vocal_2.hash') + } + }, + 'sources': + { + 'voice_extractor': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/kim_vocal_2.onnx', + 'path': resolve_relative_path('../.assets/models/kim_vocal_2.onnx') + } + } } } -def get_voice_extractor() -> Any: - global VOICE_EXTRACTOR - - with thread_lock(): - while process_manager.is_checking(): - sleep(0.5) - if VOICE_EXTRACTOR is None: - model_path = MODELS.get('voice_extractor').get('path') - VOICE_EXTRACTOR = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(facefusion.globals.execution_device_id, facefusion.globals.execution_providers)) - return VOICE_EXTRACTOR +def get_inference_pool() -> InferencePool: + model_sources = get_model_options().get('sources') + return inference_manager.get_inference_pool(__name__, model_sources) -def clear_voice_extractor() -> None: - global VOICE_EXTRACTOR +def clear_inference_pool() -> None: + inference_manager.clear_inference_pool(__name__) - VOICE_EXTRACTOR = None + +def get_model_options() -> ModelOptions: + return MODEL_SET.get('kim_vocal_2') def pre_check() -> bool: download_directory_path = resolve_relative_path('../.assets/models') - model_url = MODELS.get('voice_extractor').get('url') - model_path = MODELS.get('voice_extractor').get('path') + model_hashes = get_model_options().get('hashes') + model_sources = get_model_options().get('sources') - if not facefusion.globals.skip_download: - process_manager.check() - conditional_download(download_directory_path, [ model_url ]) - process_manager.end() - return is_file(model_path) + return conditional_download_hashes(download_directory_path, model_hashes) and conditional_download_sources(download_directory_path, model_sources) def batch_extract_voice(audio : Audio, chunk_size : int, step_size : int) -> Audio: @@ -61,23 +62,32 @@ def batch_extract_voice(audio : Audio, chunk_size : int, step_size : int) -> Aud end = min(start + chunk_size, audio.shape[0]) temp_audio[start:end, ...] += extract_voice(audio[start:end, ...]) temp_chunk[start:end, ...] += 1 + audio = temp_audio / temp_chunk return audio def extract_voice(temp_audio_chunk : AudioChunk) -> AudioChunk: - voice_extractor = get_voice_extractor() - chunk_size = 1024 * (voice_extractor.get_inputs()[0].shape[3] - 1) + voice_extractor = get_inference_pool().get('voice_extractor') + chunk_size = (voice_extractor.get_inputs()[0].shape[3] - 1) * 1024 trim_size = 3840 temp_audio_chunk, pad_size = prepare_audio_chunk(temp_audio_chunk.T, chunk_size, trim_size) temp_audio_chunk = decompose_audio_chunk(temp_audio_chunk, trim_size) + temp_audio_chunk = forward(temp_audio_chunk) + temp_audio_chunk = compose_audio_chunk(temp_audio_chunk, trim_size) + temp_audio_chunk = normalize_audio_chunk(temp_audio_chunk, chunk_size, trim_size, pad_size) + return temp_audio_chunk + + +def forward(temp_audio_chunk : AudioChunk) -> AudioChunk: + voice_extractor = get_inference_pool().get('voice_extractor') + with thread_semaphore(): temp_audio_chunk = voice_extractor.run(None, { - voice_extractor.get_inputs()[0].name: temp_audio_chunk + 'input': temp_audio_chunk })[0] - temp_audio_chunk = compose_audio_chunk(temp_audio_chunk, trim_size) - temp_audio_chunk = normalize_audio_chunk(temp_audio_chunk, chunk_size, trim_size, pad_size) + return temp_audio_chunk @@ -91,6 +101,7 @@ def prepare_audio_chunk(temp_audio_chunk : AudioChunk, chunk_size : int, trim_si for index in range(0, audio_chunk_size, 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, pad_size @@ -99,12 +110,14 @@ def prepare_audio_chunk(temp_audio_chunk : AudioChunk, chunk_size : int, trim_si def decompose_audio_chunk(temp_audio_chunk : AudioChunk, trim_size : int) -> AudioChunk: frame_size = 7680 frame_overlap = 6656 - voice_extractor_shape = get_voice_extractor().get_inputs()[0].shape + frame_total = 3072 + bin_total = 256 + channel_total = 4 window = scipy.signal.windows.hann(frame_size) temp_audio_chunk = scipy.signal.stft(temp_audio_chunk, nperseg = frame_size, noverlap = 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, trim_size + 1, voice_extractor_shape[3]).reshape(-1, voice_extractor_shape[1], trim_size + 1, voice_extractor_shape[3]) - temp_audio_chunk = temp_audio_chunk[:, :, :voice_extractor_shape[2]] + temp_audio_chunk = temp_audio_chunk.reshape(-1, 2, 2, trim_size + 1, bin_total).reshape(-1, channel_total, trim_size + 1, bin_total) + temp_audio_chunk = temp_audio_chunk[:, :, :frame_total] temp_audio_chunk /= numpy.sqrt(1.0 / window.sum() ** 2) return temp_audio_chunk @@ -112,10 +125,11 @@ def decompose_audio_chunk(temp_audio_chunk : AudioChunk, trim_size : int) -> Aud def compose_audio_chunk(temp_audio_chunk : AudioChunk, trim_size : int) -> AudioChunk: frame_size = 7680 frame_overlap = 6656 - voice_extractor_shape = get_voice_extractor().get_inputs()[0].shape + frame_total = 3072 + bin_total = 256 window = scipy.signal.windows.hann(frame_size) - temp_audio_chunk = numpy.pad(temp_audio_chunk, ((0, 0), (0, 0), (0, trim_size + 1 - voice_extractor_shape[2]), (0, 0))) - temp_audio_chunk = temp_audio_chunk.reshape(-1, 2, trim_size + 1, voice_extractor_shape[3]).transpose((0, 2, 3, 1)) + temp_audio_chunk = numpy.pad(temp_audio_chunk, ((0, 0), (0, 0), (0, trim_size + 1 - frame_total), (0, 0))) + temp_audio_chunk = temp_audio_chunk.reshape(-1, 2, trim_size + 1, 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 = frame_size, noverlap = frame_overlap, window = window)[1] temp_audio_chunk *= numpy.sqrt(1.0 / window.sum() ** 2) diff --git a/facefusion/wording.py b/facefusion/wording.py index 1401e44c..7b5a6b34 100755 --- a/facefusion/wording.py +++ b/facefusion/wording.py @@ -4,6 +4,7 @@ WORDING : Dict[str, Any] =\ { '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', @@ -31,19 +32,54 @@ WORDING : Dict[str, Any] =\ 'processing_image_failed': 'Processing to image failed', 'processing_video_succeed': 'Processing to video succeed in {seconds} seconds', 'processing_video_failed': 'Processing to video failed', - 'model_download_not_done': 'Download of the model is not done', - 'model_file_not_present': 'File of the model is not present', - 'select_image_source': 'Select a image for source path', - 'select_audio_source': 'Select a audio for source path', - 'select_video_target': 'Select a video for target path', - 'select_image_or_video_target': 'Select a image or video for target path', - 'select_file_or_directory_output': 'Select a file or directory for output path', + 'choose_image_source': 'Choose a image for the source', + 'choose_audio_source': 'Choose a audio for the source', + 'choose_video_target': 'Choose a video for the target', + 'choose_image_or_video_target': 'Choose a 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', - 'frame_processor_not_loaded': 'Frame processor {frame_processor} could not be loaded', - 'frame_processor_not_implemented': 'Frame processor {frame_processor} not implemented correctly', + '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', + '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_succeed': 'Processing of job {job_id} succeed', + 'processing_jobs_succeed': 'Processing of all job succeed', + '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_succeed': 'Validating hash for {hash_file_name} succeed', + 'validating_hash_failed': 'Validating hash for {hash_file_name} failed', + 'validating_source_succeed': 'Validating source for {source_file_name} succeed', + 'validating_source_failed': 'Validating source for {source_file_name} failed', + 'deleting_corrupt_source': 'Deleting corrupt source for {source_file_name}', + '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': ':', @@ -52,36 +88,28 @@ WORDING : Dict[str, Any] =\ 'help': { # installer - 'install_dependency': 'select the variant of {dependency} to install', + 'install_dependency': 'choose the variant of {dependency} to install', 'skip_conda': 'skip the conda environment check', # general - 'config': 'choose the config file to override defaults', - 'source': 'choose single or multiple source images or audios', - 'target': 'choose single target image or video', - 'output': 'specify the output file or directory', - # misc - 'force_download': 'force automate downloads and exit', - 'skip_download': 'omit automate downloads and remote lookups', - 'headless': 'run the program without a user interface', - 'log_level': 'adjust the message severity displayed in the terminal', - # execution - 'execution_device_id': 'specify the device used for processing', - 'execution_providers': 'accelerate the model inference using different providers (choices: {choices}, ...)', - 'execution_thread_count': 'specify the amount of parallel threads while processing', - 'execution_queue_count': 'specify the amount of frames each thread is processing', - # memory - 'video_memory_strategy': 'balance fast frame processing and low VRAM usage', - 'system_memory_limit': 'limit the available RAM that can be used while processing', + 'config_path': 'choose the config file to override defaults', + 'source_paths': 'choose single or multiple source images or audios', + 'target_path': 'choose single target image or video', + 'output_path': 'specify the output image or video within a directory', + 'jobs_path': 'specify the directory to store jobs', # face analyser - 'face_analyser_order': 'specify the order in which the face analyser detects faces', - 'face_analyser_age': 'filter the detected faces based on their age', - 'face_analyser_gender': 'filter the detected faces based on their gender', - 'face_detector_model': 'choose the model responsible for detecting the face', + 'face_detector_model': 'choose the model responsible for detecting the faces', 'face_detector_size': 'specify the size of the frame provided to the face detector', + 'face_detector_angles': 'specify the angles to rotate the frame before detecting faces', 'face_detector_score': 'filter the detected faces base on the confidence score', - 'face_landmarker_score': 'filter the detected landmarks base on the confidence score', + 'face_landmarker_model': 'choose the model responsible for detecting the face landmarks', + 'face_landmarker_score': 'filter the detected face landmarks base on the confidence score', # face selector 'face_selector_mode': 'use reference based tracking or simple matching', + 'face_selector_order': 'specify the order of the detected faces', + 'face_selector_gender': 'filter the detected faces based on their gender', + 'face_selector_age_start': 'filter the detected faces based on their age start', + 'face_selector_age_end': 'filter the detected faces based on their age end', + '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 desired similarity between the reference face and target face', 'reference_frame_number': 'specify the frame used to create the reference face', @@ -98,18 +126,39 @@ WORDING : Dict[str, Any] =\ # output creation 'output_image_quality': 'specify the image quality which translates to the compression factor', 'output_image_resolution': 'specify the image output resolution based on the target image', - 'output_video_encoder': 'specify the encoder use for the video compression', + 'output_audio_encoder': 'specify the encoder used for the audio output', + 'output_video_encoder': 'specify the encoder used for the video output', 'output_video_preset': 'balance fast video processing and video file size', 'output_video_quality': 'specify the video quality which translates to the compression factor', 'output_video_resolution': 'specify the video output resolution based on the target video', 'output_video_fps': 'specify the video output fps based on the target video', 'skip_audio': 'omit the audio from the target video', - # frame processors - 'frame_processors': 'load a single or multiple frame processors. (choices: {choices}, ...)', - 'face_debugger_items': 'load a single or multiple frame processors (choices: {choices})', + # processors + 'processors': 'load a single or multiple processors. (choices: {choices}, ...)', + 'age_modifier_model': 'choose the model responsible for aging the face', + 'age_modifier_direction': 'specify the direction in which the age should be modified', + 'expression_restorer_model': 'choose the model responsible for restoring the expression', + 'expression_restorer_factor': 'restore factor of expression from target face', + 'face_debugger_items': 'load a single or multiple processors (choices: {choices})', + 'face_editor_model': 'choose the model responsible for editing the face', + 'face_editor_eyebrow_direction': 'specify the eyebrow direction', + 'face_editor_eye_gaze_horizontal': 'specify the horizontal eye gaze', + 'face_editor_eye_gaze_vertical': 'specify the vertical eye gaze', + 'face_editor_eye_open_ratio': 'specify the ratio of eye opening', + 'face_editor_lip_open_ratio': 'specify the ratio of lip opening', + 'face_editor_mouth_grim': 'specify the mouth grim amount', + 'face_editor_mouth_pout': 'specify the mouth pout amount', + 'face_editor_mouth_purse': 'specify the mouth purse amount', + 'face_editor_mouth_smile': 'specify the mouth smile amount', + 'face_editor_mouth_position_horizontal': 'specify the mouth position horizontal amount', + 'face_editor_mouth_position_vertical': 'specify the mouth position vertical amount', + 'face_editor_head_pitch': 'specify the head pitch amount', + 'face_editor_head_yaw': 'specify the head yaw amount', + 'face_editor_head_roll': 'specify the head roll amount', 'face_enhancer_model': 'choose the model responsible for enhancing the face', 'face_enhancer_blend': 'blend the enhanced into the previous face', 'face_swapper_model': 'choose the model responsible for swapping the face', + 'face_swapper_pixel_boost': 'choose the pixel boost resolution for the face swapper', 'frame_colorizer_model': 'choose the model responsible for colorizing the frame', 'frame_colorizer_blend': 'blend the colorized into the previous frame', 'frame_colorizer_size': 'specify the size of the frame provided to the frame colorizer', @@ -118,94 +167,144 @@ WORDING : Dict[str, Any] =\ 'lip_syncer_model': 'choose the model responsible for syncing the lips', # uis 'open_browser': 'open the browser once the program is ready', - 'ui_layouts': 'launch a single or multiple UI layouts (choices: {choices}, ...)' + 'ui_layouts': 'launch a single or multiple UI layouts (choices: {choices}, ...)', + 'ui_workflow': 'choose the ui workflow', + # execution + 'execution_device_id': 'specify the device used for processing', + 'execution_providers': 'accelerate the model inference using different providers (choices: {choices}, ...)', + 'execution_thread_count': 'specify the amount of parallel threads while processing', + 'execution_queue_count': 'specify the amount of frames each thread is processing', + # memory + 'video_memory_strategy': 'balance fast processing and low VRAM usage', + 'system_memory_limit': 'limit the available RAM that can be used while processing', + # misc + 'skip_download': 'omit downloads and remote lookups', + 'log_level': 'adjust the message severity displayed in the terminal', + # run + 'run': 'run the program', + 'headless_run': 'run the program in headless mode', + 'force_download': 'force automate downloads and exit', + # jobs + 'job_id': 'specify the job id', + 'step_index': 'specify the step index', + # job manager + '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_list': 'list jobs by status', + '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 runner + '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': + { + 'become_a_member': 'become a member', + 'join_our_community': 'join our community', + 'read_the_documentation': 'read the documentation' }, 'uis': { - # general - 'start_button': 'START', - 'stop_button': 'STOP', - 'clear_button': 'CLEAR', - # about - 'donate_button': 'DONATE', - # benchmark - 'benchmark_results_dataframe': 'BENCHMARK RESULTS', - # benchmark options - 'benchmark_runs_checkbox_group': 'BENCHMARK RUNS', + 'age_modifier_direction_slider': 'AGE MODIFIER DIRECTION', + 'age_modifier_model_dropdown': 'AGE MODIFIER MODEL', + 'apply_button': 'APPLY', 'benchmark_cycles_slider': 'BENCHMARK CYCLES', - # common options + 'benchmark_runs_checkbox_group': 'BENCHMARK RUNS', + 'clear_button': 'CLEAR', 'common_options_checkbox_group': 'OPTIONS', - # execution 'execution_providers_checkbox_group': 'EXECUTION PROVIDERS', - # execution queue count 'execution_queue_count_slider': 'EXECUTION QUEUE COUNT', - # execution thread count 'execution_thread_count_slider': 'EXECUTION THREAD COUNT', - # face analyser - 'face_analyser_order_dropdown': 'FACE ANALYSER ORDER', - 'face_analyser_age_dropdown': 'FACE ANALYSER AGE', - 'face_analyser_gender_dropdown': 'FACE ANALYSER GENDER', + 'expression_restorer_factor_slider': 'EXPRESSION RESTORER FACTOR', + 'expression_restorer_model_dropdown': 'EXPRESSION RESTORER MODEL', + 'face_debugger_items_checkbox_group': 'FACE DEBUGGER ITEMS', + 'face_detector_angles_checkbox_group': 'FACE DETECTOR ANGLES', 'face_detector_model_dropdown': 'FACE DETECTOR MODEL', - 'face_detector_size_dropdown': 'FACE DETECTOR SIZE', 'face_detector_score_slider': 'FACE DETECTOR SCORE', + 'face_detector_size_dropdown': 'FACE DETECTOR SIZE', + 'face_editor_model_dropdown': 'FACE EDITOR MODEL', + 'face_editor_eye_gaze_horizontal_slider': 'FACE EDITOR EYE GAZE HORIZONTAL', + 'face_editor_eye_gaze_vertical_slider': 'FACE EDITOR EYE GAZE VERTICAL', + 'face_editor_eye_open_ratio_slider': 'FACE EDITOR EYE OPEN RATIO', + 'face_editor_eyebrow_direction_slider': 'FACE EDITOR EYEBROW DIRECTION', + 'face_editor_lip_open_ratio_slider': 'FACE EDITOR LIP OPEN RATIO', + 'face_editor_mouth_grim_slider': 'FACE EDITOR MOUTH GRIM', + 'face_editor_mouth_pout_slider': 'FACE EDITOR MOUTH POUT', + 'face_editor_mouth_purse_slider': 'FACE EDITOR MOUTH PURSE', + 'face_editor_mouth_smile_slider': 'FACE EDITOR MOUTH SMILE', + 'face_editor_mouth_position_horizontal_slider': 'FACE EDITOR MOUTH POSITION HORIZONTAL', + 'face_editor_mouth_position_vertical_slider': 'FACE EDITOR MOUTH POSITION VERTICAL', + 'face_editor_head_pitch_slider': 'FACE EDITOR HEAD PITCH', + 'face_editor_head_yaw_slider': 'FACE EDITOR HEAD YAW', + 'face_editor_head_roll_slider': 'FACE EDITOR HEAD ROLL', + 'face_enhancer_blend_slider': 'FACE ENHANCER BLEND', + 'face_enhancer_model_dropdown': 'FACE ENHANCER MODEL', + 'face_landmarker_model_dropdown': 'FACE LANDMARKER MODEL', 'face_landmarker_score_slider': 'FACE LANDMARKER SCORE', - # face masker - 'face_mask_types_checkbox_group': 'FACE MASK TYPES', 'face_mask_blur_slider': 'FACE MASK BLUR', - 'face_mask_padding_top_slider': 'FACE MASK PADDING TOP', - 'face_mask_padding_right_slider': 'FACE MASK PADDING RIGHT', '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_region_checkbox_group': 'FACE MASK REGIONS', - # face selector + 'face_mask_types_checkbox_group': 'FACE MASK TYPES', + 'face_selector_gender_dropdown': 'FACE SELECTOR GENDER', + 'face_selector_race_dropdown': 'FACE SELECTOR RACE', + 'face_selector_age_range_slider': 'FACE SELECTOR AGE', 'face_selector_mode_dropdown': 'FACE SELECTOR MODE', - 'reference_face_gallery': 'REFERENCE FACE', - 'reference_face_distance_slider': 'REFERENCE FACE DISTANCE', - # frame processors - 'frame_processors_checkbox_group': 'FRAME PROCESSORS', - # frame processors options - 'face_debugger_items_checkbox_group': 'FACE DEBUGGER ITEMS', - 'face_enhancer_model_dropdown': 'FACE ENHANCER MODEL', - 'face_enhancer_blend_slider': 'FACE ENHANCER BLEND', + 'face_selector_order_dropdown': 'FACE SELECTOR ORDER', 'face_swapper_model_dropdown': 'FACE SWAPPER MODEL', - 'frame_colorizer_model_dropdown': 'FRAME COLORIZER MODEL', + 'face_swapper_pixel_boost_dropdown': 'FACE SWAPPER PIXEL BOOST', 'frame_colorizer_blend_slider': 'FRAME COLORIZER BLEND', + 'frame_colorizer_model_dropdown': 'FRAME COLORIZER MODEL', 'frame_colorizer_size_dropdown': 'FRAME COLORIZER SIZE', - 'frame_enhancer_model_dropdown': 'FRAME ENHANCER MODEL', 'frame_enhancer_blend_slider': 'FRAME ENHANCER BLEND', + 'frame_enhancer_model_dropdown': 'FRAME ENHANCER 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', 'lip_syncer_model_dropdown': 'LIP SYNCER MODEL', - # memory - 'video_memory_strategy_dropdown': 'VIDEO MEMORY STRATEGY', - 'system_memory_limit_slider': 'SYSTEM MEMORY LIMIT', - # output + 'log_level_dropdown': 'LOG LEVEL', + 'output_audio_encoder_dropdown': 'OUTPUT AUDIO ENCODER', 'output_image_or_video': 'OUTPUT', - # output options - 'output_path_textbox': 'OUTPUT PATH', 'output_image_quality_slider': 'OUTPUT IMAGE QUALITY', 'output_image_resolution_dropdown': 'OUTPUT IMAGE RESOLUTION', + '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_resolution_dropdown': 'OUTPUT VIDEO RESOLUTION', - 'output_video_fps_slider': 'OUTPUT VIDEO FPS', - # preview - 'preview_image': 'PREVIEW', 'preview_frame_slider': 'PREVIEW FRAME', - # source + 'preview_image': 'PREVIEW', + 'processors_checkbox_group': 'PROCESSORS', + 'reference_face_distance_slider': 'REFERENCE FACE DISTANCE', + 'reference_face_gallery': 'REFERENCE FACE', + 'refresh_button': 'REFRESH', 'source_file': 'SOURCE', - # target + 'start_button': 'START', + 'stop_button': 'STOP', + 'system_memory_limit_slider': 'SYSTEM MEMORY LIMIT', 'target_file': 'TARGET', - # temp frame + 'terminal_textbox': 'TERMINAL', 'temp_frame_format_dropdown': 'TEMP FRAME FORMAT', - # trim frame - 'trim_frame_start_slider': 'TRIM FRAME START', - 'trim_frame_end_slider': 'TRIM FRAME END', - # webcam + '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 options 'webcam_mode_radio': 'WEBCAM MODE', 'webcam_resolution_dropdown': 'WEBCAM RESOLUTION', - 'webcam_fps_slider': 'WEBCAM FPS' } } @@ -213,8 +312,8 @@ WORDING : Dict[str, Any] =\ def get(key : str) -> Optional[str]: if '.' in key: section, name = key.split('.') - if section in WORDING and name in WORDING[section]: - return WORDING[section][name] + if section in WORDING and name in WORDING.get(section): + return WORDING.get(section).get(name) if key in WORDING: - return WORDING[key] + return WORDING.get(key) return None diff --git a/install.py b/install.py index 6feca23e..000f1e72 100755 --- a/install.py +++ b/install.py @@ -1,10 +1,8 @@ #!/usr/bin/env python3 import os -import subprocess -os.environ['PIP_BREAK_SYSTEM_PACKAGES'] = '1' -subprocess.call([ 'pip', 'install', 'inquirer', '-q' ]) +os.environ['SYSTEM_VERSION_COMPAT'] = '0' from facefusion import installer diff --git a/requirements.txt b/requirements.txt index 8222e415..0f1f443d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,10 @@ filetype==1.2.0 -gradio==3.50.2 -numpy==1.26.4 -onnx==1.16.0 -onnxruntime==1.17.3 -opencv-python==4.9.0.80 -psutil==5.9.8 -tqdm==4.66.4 -scipy==1.13.0 +gradio==4.44.0 +gradio_rangeslider==0.0.6 +numpy==2.1.0 +onnx==1.16.2 +onnxruntime==1.19.2 +opencv-python==4.10.0.84 +psutil==6.0.0 +tqdm==4.66.5 +scipy==1.14.1 diff --git a/tests/helper.py b/tests/helper.py new file mode 100644 index 00000000..45d810bb --- /dev/null +++ b/tests/helper.py @@ -0,0 +1,44 @@ +import os + +from facefusion.filesystem import create_directory, is_directory, is_file, remove_directory +from facefusion.temp_helper import get_base_directory_path +from facefusion.typing 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(get_base_directory_path(), '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(get_base_directory_path(), '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(get_base_directory_path(), '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 index 765acfb8..66039f1e 100644 --- a/tests/test_audio.py +++ b/tests/test_audio.py @@ -1,26 +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.mp3' + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.mp3' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/source.mp3', '.assets/examples/source.wav' ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), get_test_example_file('source.wav') ]) def test_get_audio_frame() -> None: - assert get_audio_frame('.assets/examples/source.mp3', 25) is not None - assert get_audio_frame('.assets/examples/source.wav', 25) is not None + assert get_audio_frame(get_test_example_file('source.mp3'), 25) is not None + assert get_audio_frame(get_test_example_file('source.wav'), 25) is not None assert get_audio_frame('invalid', 25) is None def test_read_static_audio() -> None: - assert len(read_static_audio('.assets/examples/source.mp3', 25)) == 280 - assert len(read_static_audio('.assets/examples/source.wav', 25)) == 280 + 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 00000000..d4f28143 --- /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', '-j', 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', '-j', 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_expression_restorer.py b/tests/test_cli_expression_restorer.py new file mode 100644 index 00000000..03dd3064 --- /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', '-j', 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', '-j', 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 index 51167af7..e41fafb9 100644 --- a/tests/test_cli_face_debugger.py +++ b/tests/test_cli_face_debugger.py @@ -1,31 +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + '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', '.assets/examples/target-240p.mp4', '-vframes', '1', '.assets/examples/target-240p.jpg' ]) + 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, 'run.py', '--frame-processors', 'face_debugger', '-t', '.assets/examples/target-240p.jpg', '-o', '.assets/examples/test_debug_face_to_image.jpg', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'image succeed' in run.stdout.decode() + 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, 'run.py', '--frame-processors', 'face_debugger', '-t', '.assets/examples/target-240p.mp4', '-o', '.assets/examples/test_debug_face_to_video.mp4', '--trim-frame-end', '10', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'video succeed' in run.stdout.decode() + 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 00000000..633ff65f --- /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', '-j', 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', '-j', 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 index ccedf7d5..cc01d8dc 100644 --- a/tests/test_cli_face_enhancer.py +++ b/tests/test_cli_face_enhancer.py @@ -1,32 +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + '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', '.assets/examples/target-240p.mp4', '-vframes', '1', '.assets/examples/target-240p.jpg' ]) + 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, 'run.py', '--frame-processors', 'face_enhancer', '-t', '.assets/examples/target-240p.jpg', '-o', '.assets/examples/test_enhance_face_to_image.jpg', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'image succeed' in run.stdout.decode() + 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, 'run.py', '--frame-processors', 'face_enhancer', '-t', '.assets/examples/target-240p.mp4', '-o', '.assets/examples/test_enhance_face_to_video.mp4', '--trim-frame-end', '10', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) - - assert run.returncode == 0 - assert 'video succeed' in run.stdout.decode() + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 index 399870ab..bf0c0983 100644 --- a/tests/test_cli_face_swapper.py +++ b/tests/test_cli_face_swapper.py @@ -1,31 +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + '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', '.assets/examples/target-240p.mp4', '-vframes', '1', '.assets/examples/target-240p.jpg' ]) + 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, 'run.py', '--frame-processors', 'face_swapper', '-s', '.assets/examples/source.jpg', '-t', '.assets/examples/target-240p.jpg', '-o', '.assets/examples/test_swap_face_to_image.jpg', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'image succeed' in run.stdout.decode() + 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, 'run.py', '--frame-processors', 'face_swapper', '-s', '.assets/examples/source.jpg', '-t', '.assets/examples/target-240p.mp4', '-o', '.assets/examples/test_swap_face_to_video.mp4', '--trim-frame-end', '10', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'video succeed' in run.stdout.decode() + 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 index 97ff08ec..9796651e 100644 --- a/tests/test_cli_frame_colorizer.py +++ b/tests/test_cli_frame_colorizer.py @@ -1,32 +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + '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', '.assets/examples/target-240p.mp4', '-vframes', '1', '-vf', 'hue=s=0', '.assets/examples/target-240p-0sat.jpg' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'hue=s=0', '.assets/examples/target-240p-0sat.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, 'run.py', '--frame-processors', 'frame_colorizer', '-t', '.assets/examples/target-240p-0sat.jpg', '-o', '.assets/examples/test_colorize_frame_to_image.jpg', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'image succeed' in run.stdout.decode() + 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, 'run.py', '--frame-processors', 'frame_colorizer', '-t', '.assets/examples/target-240p-0sat.mp4', '-o', '.assets/examples/test_colorize_frame_to_video.mp4', '--trim-frame-end', '10', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'video succeed' in run.stdout.decode() + 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 index 89399f46..0530bdae 100644 --- a/tests/test_cli_frame_enhancer.py +++ b/tests/test_cli_frame_enhancer.py @@ -1,31 +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + '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', '.assets/examples/target-240p.mp4', '-vframes', '1', '.assets/examples/target-240p.jpg' ]) + 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, 'run.py', '--frame-processors', 'frame_enhancer', '-t', '.assets/examples/target-240p.jpg', '-o', '.assets/examples/test_enhance_frame_to_image.jpg', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'image succeed' in run.stdout.decode() + 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, 'run.py', '--frame-processors', 'frame_enhancer', '-t', '.assets/examples/target-240p.mp4', '-o', '.assets/examples/test_enhance_frame_to_video.mp4', '--trim-frame-end', '10', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'video succeed' in run.stdout.decode() + 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 00000000..7abddcc5 --- /dev/null +++ b/tests/test_cli_job_manager.py @@ -0,0 +1,204 @@ +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()) + + +def test_job_create() -> None: + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-create', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-submit', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-submit-all-1', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-submit-all-2', '-j', 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', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + + 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-delete', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-delete', 'test-job-delete', '-j', 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-delete-all-1', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-delete-all-2', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-delete-all', '-j', get_test_jobs_directory() ] + + 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', '-j', 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', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-add-step', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-remix-step', '-j', 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', '-j', 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', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-insert-step', '-j', 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', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-remove-step', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-remove-step', '-j', 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', '-j', 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', '-j', 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 00000000..906ef243 --- /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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-run', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run', '-j', 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-submit', 'test-job-run', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-run', 'test-job-run', '-j', 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-run-all-1', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-run-all-2', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-1', '-j', 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', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-submit-all', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-run-all', '-j', get_test_jobs_directory() ] + + 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-retry', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry', '-j', 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', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + + assert subprocess.run(commands).returncode == 1 + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-retry-all-1', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-retry-all-2', '-j', get_test_jobs_directory() ] + subprocess.run(commands) + + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-1', '-j', 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', '-j', 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', '-j', 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', '-j', get_test_jobs_directory() ] + + 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', '-j', get_test_jobs_directory() ] + + 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 index 089ff41a..bd8d078d 100644 --- a/tests/test_cli_lip_syncer.py +++ b/tests/test_cli_lip_syncer.py @@ -1,32 +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.mp3', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + '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', '.assets/examples/target-240p.mp4', '-vframes', '1', '.assets/examples/target-240p.jpg' ]) + 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, 'run.py', '--frame-processors', 'lip_syncer', '-s', '.assets/examples/source.mp3', '-t', '.assets/examples/target-240p.jpg', '-o', '.assets/examples/test_sync_lip_to_image.jpg', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'image succeed' in run.stdout.decode() + 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, 'run.py', '--frame-processors', 'lip_syncer', '-s', '.assets/examples/source.mp3', '-t', '.assets/examples/target-240p.mp4', '-o', '.assets/examples/test_sync_lip_to_video.mp4', '--trim-frame-end', '10', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) + commands = [ sys.executable, 'facefusion.py', 'headless-run', '-j', 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 run.returncode == 0 - assert 'video succeed' in run.stdout.decode() + assert subprocess.run(commands).returncode == 0 + assert is_test_output_file('test_sync_lip_to_video.mp4') is True diff --git a/tests/test_common_helper.py b/tests/test_common_helper.py index 6cef2c3c..ed7dcfdb 100644 --- a/tests/test_common_helper.py +++ b/tests/test_common_helper.py @@ -1,8 +1,12 @@ -from facefusion.common_helper import create_metavar, create_int_range, create_float_range +from facefusion.common_helper import calc_float_step, calc_int_step, create_float_metavar, create_float_range, create_int_metavar, create_int_range -def test_create_metavar() -> None: - assert create_metavar([ 1, 2, 3, 4, 5 ]) == '[1-5]' +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: @@ -13,3 +17,11 @@ def test_create_int_range() -> None: 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 calc_int_step([ 0, 1 ]) == 1 + + +def test_calc_float_step() -> None: + assert calc_float_step([ 0.1, 0.2 ]) == 0.1 diff --git a/tests/test_config.py b/tests/test_config.py index 8c830fd6..88550e51 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,4 +1,5 @@ from configparser import ConfigParser + import pytest from facefusion import config diff --git a/tests/test_date_helper.py b/tests/test_date_helper.py new file mode 100644 index 00000000..7ec714f8 --- /dev/null +++ b/tests/test_date_helper.py @@ -0,0 +1,15 @@ +from datetime import datetime, timedelta + +from facefusion.date_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, 5)) == '5 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(5, 10, 15)) == '5 days, 10 hours and 15 minutes ago' diff --git a/tests/test_download.py b/tests/test_download.py index 6df94b17..8ca1d368 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -1,23 +1,24 @@ import pytest from facefusion.download import conditional_download, get_download_size, is_download_done +from .helper import get_test_example_file, get_test_examples_directory @pytest.fixture(scope = 'module', autouse = True) def before_all() -> None: - conditional_download('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' ]) def test_get_download_size() -> None: - assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4') == 191675 - assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-360p.mp4') == 370732 + assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4') == 191675 + assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-360p.mp4') == 370732 assert get_download_size('invalid') == 0 def test_is_download_done() -> None: - assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4', '.assets/examples/target-240p.mp4') is True - assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4', 'invalid') is False + assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4', get_test_example_file('target-240p.mp4')) is True + assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4', 'invalid') is False assert is_download_done('invalid', 'invalid') is False diff --git a/tests/test_execution.py b/tests/test_execution.py index eaada164..1823f1e8 100644 --- a/tests/test_execution.py +++ b/tests/test_execution.py @@ -1,27 +1,24 @@ -from facefusion.execution import encode_execution_providers, decode_execution_providers, has_execution_provider, apply_execution_provider_options +from facefusion.execution import create_execution_providers, get_execution_provider_choices, has_execution_provider -def test_encode_execution_providers() -> None: - assert encode_execution_providers([ 'CPUExecutionProvider' ]) == [ 'cpu' ] - - -def test_decode_execution_providers() -> None: - assert decode_execution_providers([ 'cpu' ]) == [ 'CPUExecutionProvider' ] +def test_get_execution_provider_choices() -> None: + assert 'cpu' in get_execution_provider_choices() def test_has_execution_provider() -> None: - assert has_execution_provider('CPUExecutionProvider') is True - assert has_execution_provider('InvalidExecutionProvider') is False + assert has_execution_provider('cpu') is True + assert has_execution_provider('openvino') is False def test_multiple_execution_providers() -> None: execution_provider_with_options =\ [ - 'CPUExecutionProvider', ('CUDAExecutionProvider', { 'device_id': '1', 'cudnn_conv_algo_search': 'DEFAULT' - }) + }), + 'CPUExecutionProvider' ] - assert apply_execution_provider_options('1', [ 'CPUExecutionProvider', 'CUDAExecutionProvider' ]) == execution_provider_with_options + + assert create_execution_providers('1', [ 'cpu', 'cuda' ]) == execution_provider_with_options diff --git a/tests/test_face_analyser.py b/tests/test_face_analyser.py index 957dfc8d..7c351861 100644 --- a/tests/test_face_analyser.py +++ b/tests/test_face_analyser.py @@ -1,103 +1,109 @@ import subprocess + import pytest -import facefusion.globals +from facefusion import face_classifier, face_detector, face_landmarker, face_recognizer, state_manager from facefusion.download import conditional_download -from facefusion.face_analyser import pre_check, clear_face_analyser, get_one_face +from facefusion.face_analyser import get_many_faces, get_one_face from facefusion.typing import Face 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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg' + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/source.jpg', '-vf', 'crop=iw*0.8:ih*0.8', '.assets/examples/source-80crop.jpg' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/source.jpg', '-vf', 'crop=iw*0.7:ih*0.7', '.assets/examples/source-70crop.jpg' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/source.jpg', '-vf', 'crop=iw*0.6:ih*0.6', '.assets/examples/source-60crop.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_id', 0) + state_manager.init_item('execution_providers', [ 'cpu' ]) + 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: - facefusion.globals.face_detector_score = 0.5 - facefusion.globals.face_landmarker_score = 0.5 - facefusion.globals.face_recognizer_model = 'arcface_inswapper' - clear_face_analyser() + 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: - facefusion.globals.face_detector_model = 'retinaface' - facefusion.globals.face_detector_size = '320x320' + state_manager.init_item('face_detector_model', 'retinaface') + state_manager.init_item('face_detector_size', '320x320') + face_detector.pre_check() - pre_check() source_paths =\ [ - '.assets/examples/source.jpg', - '.assets/examples/source-80crop.jpg', - '.assets/examples/source-70crop.jpg', - '.assets/examples/source-60crop.jpg' + 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) - face = get_one_face(source_frame) + many_faces = get_many_faces([ source_frame ]) + face = get_one_face(many_faces) assert isinstance(face, Face) def test_get_one_face_with_scrfd() -> None: - facefusion.globals.face_detector_model = 'scrfd' - facefusion.globals.face_detector_size = '640x640' + state_manager.init_item('face_detector_model', 'scrfd') + state_manager.init_item('face_detector_size', '640x640') + face_detector.pre_check() - pre_check() source_paths =\ [ - '.assets/examples/source.jpg', - '.assets/examples/source-80crop.jpg', - '.assets/examples/source-70crop.jpg', - '.assets/examples/source-60crop.jpg' + 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) - face = get_one_face(source_frame) + many_faces = get_many_faces([ source_frame ]) + face = get_one_face(many_faces) assert isinstance(face, Face) def test_get_one_face_with_yoloface() -> None: - facefusion.globals.face_detector_model = 'yoloface' - facefusion.globals.face_detector_size = '640x640' + state_manager.init_item('face_detector_model', 'yoloface') + state_manager.init_item('face_detector_size', '640x640') + face_detector.pre_check() - pre_check() source_paths =\ [ - '.assets/examples/source.jpg', - '.assets/examples/source-80crop.jpg', - '.assets/examples/source-70crop.jpg', - '.assets/examples/source-60crop.jpg' + 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) - face = get_one_face(source_frame) + many_faces = get_many_faces([ source_frame ]) + face = get_one_face(many_faces) assert isinstance(face, Face) -def test_get_one_face_with_yunet() -> None: - facefusion.globals.face_detector_model = 'yunet' - facefusion.globals.face_detector_size = '640x640' +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 ]) - pre_check() - source_paths =\ - [ - '.assets/examples/source.jpg', - '.assets/examples/source-80crop.jpg', - '.assets/examples/source-70crop.jpg', - '.assets/examples/source-60crop.jpg' - ] - for source_path in source_paths: - source_frame = read_static_image(source_path) - face = get_one_face(source_frame) - - assert isinstance(face, Face) + assert isinstance(many_faces[0], Face) + assert isinstance(many_faces[1], Face) + assert isinstance(many_faces[2], Face) diff --git a/tests/test_ffmpeg.py b/tests/test_ffmpeg.py index 5c4bd2d0..ef2e22cf 100644 --- a/tests/test_ffmpeg.py +++ b/tests/test_ffmpeg.py @@ -1,113 +1,127 @@ import glob import subprocess + import pytest -import facefusion.globals -from facefusion import process_manager -from facefusion.filesystem import get_temp_directory_path, create_temp, clear_temp +from facefusion import process_manager, state_manager from facefusion.download import conditional_download -from facefusion.ffmpeg import extract_frames, read_audio_buffer +from facefusion.ffmpeg import concat_video, extract_frames, read_audio_buffer +from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_directory_path +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.mp3', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + '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', '.assets/examples/source.mp3', '.assets/examples/source.wav' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=25', '.assets/examples/target-240p-25fps.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=30', '.assets/examples/target-240p-30fps.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=60', '.assets/examples/target-240p-60fps.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') ]) + state_manager.init_item('temp_frame_format', 'jpg') + state_manager.init_item('output_audio_encoder', 'aac') @pytest.fixture(scope = 'function', autouse = True) def before_each() -> None: - facefusion.globals.trim_frame_start = None - facefusion.globals.trim_frame_end = None - facefusion.globals.temp_frame_format = 'jpg' + state_manager.clear_item('trim_frame_start') + state_manager.clear_item('trim_frame_end') + prepare_test_output_directory() def test_extract_frames() -> None: target_paths =\ [ - '.assets/examples/target-240p-25fps.mp4', - '.assets/examples/target-240p-30fps.mp4', - '.assets/examples/target-240p-60fps.mp4' + get_test_example_file('target-240p-25fps.mp4'), + get_test_example_file('target-240p-30fps.mp4'), + get_test_example_file('target-240p-60fps.mp4') ] for target_path in target_paths: temp_directory_path = get_temp_directory_path(target_path) - create_temp(target_path) + create_temp_directory(target_path) assert extract_frames(target_path, '452x240', 30.0) is True assert len(glob.glob1(temp_directory_path, '*.jpg')) == 324 - clear_temp(target_path) + clear_temp_directory(target_path) def test_extract_frames_with_trim_start() -> None: - facefusion.globals.trim_frame_start = 224 - data_provider =\ + state_manager.init_item('trim_frame_start', 224) + providers =\ [ - ('.assets/examples/target-240p-25fps.mp4', 55), - ('.assets/examples/target-240p-30fps.mp4', 100), - ('.assets/examples/target-240p-60fps.mp4', 212) + (get_test_example_file('target-240p-25fps.mp4'), 55), + (get_test_example_file('target-240p-30fps.mp4'), 100), + (get_test_example_file('target-240p-60fps.mp4'), 212) ] - for target_path, frame_total in data_provider: + for target_path, frame_total in providers: temp_directory_path = get_temp_directory_path(target_path) - create_temp(target_path) + create_temp_directory(target_path) assert extract_frames(target_path, '452x240', 30.0) is True assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total - clear_temp(target_path) + clear_temp_directory(target_path) def test_extract_frames_with_trim_start_and_trim_end() -> None: - facefusion.globals.trim_frame_start = 124 - facefusion.globals.trim_frame_end = 224 - data_provider =\ + state_manager.init_item('trim_frame_start', 124) + state_manager.init_item('trim_frame_end', 224) + providers =\ [ - ('.assets/examples/target-240p-25fps.mp4', 120), - ('.assets/examples/target-240p-30fps.mp4', 100), - ('.assets/examples/target-240p-60fps.mp4', 50) + (get_test_example_file('target-240p-25fps.mp4'), 120), + (get_test_example_file('target-240p-30fps.mp4'), 100), + (get_test_example_file('target-240p-60fps.mp4'), 50) ] - for target_path, frame_total in data_provider: + for target_path, frame_total in providers: temp_directory_path = get_temp_directory_path(target_path) - create_temp(target_path) + create_temp_directory(target_path) assert extract_frames(target_path, '452x240', 30.0) is True assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total - clear_temp(target_path) + clear_temp_directory(target_path) def test_extract_frames_with_trim_end() -> None: - facefusion.globals.trim_frame_end = 100 - data_provider =\ + state_manager.init_item('trim_frame_end', 100) + providers =\ [ - ('.assets/examples/target-240p-25fps.mp4', 120), - ('.assets/examples/target-240p-30fps.mp4', 100), - ('.assets/examples/target-240p-60fps.mp4', 50) + (get_test_example_file('target-240p-25fps.mp4'), 120), + (get_test_example_file('target-240p-30fps.mp4'), 100), + (get_test_example_file('target-240p-60fps.mp4'), 50) ] - for target_path, frame_total in data_provider: + for target_path, frame_total in providers: temp_directory_path = get_temp_directory_path(target_path) - create_temp(target_path) + create_temp_directory(target_path) assert extract_frames(target_path, '426x240', 30.0) is True assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total - clear_temp(target_path) + clear_temp_directory(target_path) + + +def test_concat_video() -> None: + output_path = get_test_output_file('test-concat-video.mp4') + temp_output_paths =\ + [ + get_test_example_file('target-240p.mp4'), + get_test_example_file('target-240p.mp4') + ] + + assert concat_video(output_path, temp_output_paths) is True def test_read_audio_buffer() -> None: - assert isinstance(read_audio_buffer('.assets/examples/source.mp3', 1, 1), bytes) - assert isinstance(read_audio_buffer('.assets/examples/source.wav', 1, 1), bytes) - assert read_audio_buffer('.assets/examples/invalid.mp3', 1, 1) is None + assert isinstance(read_audio_buffer(get_test_example_file('source.mp3'), 1, 1), bytes) + assert isinstance(read_audio_buffer(get_test_example_file('source.wav'), 1, 1), bytes) + assert read_audio_buffer(get_test_example_file('invalid.mp3'), 1, 1) is None diff --git a/tests/test_filesystem.py b/tests/test_filesystem.py index 143dbc94..fedac41c 100644 --- a/tests/test_filesystem.py +++ b/tests/test_filesystem.py @@ -1,91 +1,119 @@ -import shutil +import os.path + import pytest from facefusion.common_helper import is_windows from facefusion.download import conditional_download -from facefusion.filesystem import get_file_size, is_file, is_directory, is_audio, has_audio, is_image, has_image, is_video, filter_audio_paths, filter_image_paths, list_directory, sanitize_path_for_windows +from facefusion.filesystem import copy_file, create_directory, filter_audio_paths, filter_image_paths, get_file_size, has_audio, has_image, in_directory, is_audio, is_directory, is_file, is_image, is_video, list_directory, remove_directory, same_file_extension, sanitize_path_for_windows +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('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.mp3', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + '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' ]) - shutil.copyfile('.assets/examples/source.jpg', '.assets/examples/söurce.jpg') + copy_file(get_test_example_file('source.jpg'), get_test_example_file('söurce.jpg')) def test_get_file_size() -> None: - assert get_file_size('.assets/examples/source.jpg') > 0 + assert get_file_size(get_test_example_file('source.jpg')) > 0 assert get_file_size('invalid') == 0 +def test_same_file_extension() -> None: + assert same_file_extension([ 'target.jpg', 'output.jpg' ]) is True + assert same_file_extension([ 'target.jpg', 'output.mp4' ]) is False + + def test_is_file() -> None: - assert is_file('.assets/examples/source.jpg') is True - assert is_file('.assets/examples') is False + 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_directory() -> None: - assert is_directory('.assets/examples') is True - assert is_directory('.assets/examples/source.jpg') is False + 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 + + def test_is_audio() -> None: - assert is_audio('.assets/examples/source.mp3') is True - assert is_audio('.assets/examples/source.jpg') is False + 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([ '.assets/examples/source.mp3' ]) is True - assert has_audio([ '.assets/examples/source.mp3', '.assets/examples/source.jpg' ]) is True - assert has_audio([ '.assets/examples/source.jpg', '.assets/examples/source.jpg' ]) is False + 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('.assets/examples/source.jpg') is True - assert is_image('.assets/examples/target-240p.mp4') is False + 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([ '.assets/examples/source.jpg' ]) is True - assert has_image([ '.assets/examples/source.jpg', '.assets/examples/source.mp3' ]) is True - assert has_image([ '.assets/examples/source.mp3', '.assets/examples/source.mp3' ]) is False + 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('.assets/examples/target-240p.mp4') is True - assert is_video('.assets/examples/source.jpg') is False + 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_filter_audio_paths() -> None: - assert filter_audio_paths([ '.assets/examples/source.jpg', '.assets/examples/source.mp3' ]) == [ '.assets/examples/source.mp3' ] - assert filter_audio_paths([ '.assets/examples/source.jpg', '.assets/examples/source.jpg' ]) == [] + 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([ '.assets/examples/source.jpg', '.assets/examples/source.mp3' ]) == [ '.assets/examples/source.jpg' ] - assert filter_image_paths([ '.assets/examples/source.mp3', '.assets/examples/source.mp3' ]) == [] + 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_sanitize_path_for_windows() -> None: + if is_windows(): + assert sanitize_path_for_windows(get_test_example_file('söurce.jpg')).endswith('SURCE~1.JPG') + assert sanitize_path_for_windows('invalid') is None + + +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_list_directory() -> None: - assert list_directory('.assets/examples') - assert list_directory('.assets/examples/source.jpg') is None + assert list_directory(get_test_examples_directory()) + assert list_directory(get_test_example_file('source.jpg')) is None assert list_directory('invalid') is None -@pytest.mark.skip() -def test_sanitize_path_for_windows() -> None: - if is_windows(): - assert sanitize_path_for_windows('.assets/examples/söurce.jpg') == 'ASSETS~1/examples/SURCE~1.JPG' - assert sanitize_path_for_windows('invalid') is None +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 diff --git a/tests/test_inference_pool.py b/tests/test_inference_pool.py new file mode 100644 index 00000000..70352bb8 --- /dev/null +++ b/tests/test_inference_pool.py @@ -0,0 +1,29 @@ +from unittest.mock import patch + +import pytest +from onnxruntime import InferenceSession + +from facefusion import content_analyser, state_manager +from facefusion.inference_manager import INFERENCE_POOLS, get_inference_pool + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + content_analyser.pre_check() + state_manager.init_item('execution_device_id', 0) + state_manager.init_item('execution_providers', [ 'cpu' ]) + + +def test_get_inference_pool() -> None: + model_sources = content_analyser.get_model_options().get('sources') + + with patch('facefusion.inference_manager.detect_app_context', return_value = 'cli'): + get_inference_pool('test', model_sources) + + assert isinstance(INFERENCE_POOLS.get('cli').get('test.cpu').get('content_analyser'), InferenceSession) + + with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'): + get_inference_pool('test', model_sources) + + assert isinstance(INFERENCE_POOLS.get('ui').get('test.cpu').get('content_analyser'), InferenceSession) + diff --git a/tests/test_job_helper.py b/tests/test_job_helper.py new file mode 100644 index 00000000..08fe6f86 --- /dev/null +++ b/tests/test_job_helper.py @@ -0,0 +1,8 @@ +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') diff --git a/tests/test_job_list.py b/tests/test_job_list.py new file mode 100644 index 00000000..732a199f --- /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 00000000..3ee6c0d1 --- /dev/null +++ b/tests/test_job_manager.py @@ -0,0 +1,373 @@ +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, 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' + } + + assert submit_jobs() is False + + create_job('job-test-submit-jobs-1') + create_job('job-test-submit-jobs-2') + + assert submit_jobs() is False + + add_step('job-test-submit-jobs-1', args_1) + add_step('job-test-submit-jobs-2', args_2) + + assert submit_jobs() is True + assert submit_jobs() 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: + assert delete_jobs() is False + + create_job('job-test-delete-jobs-1') + create_job('job-test-delete-jobs-2') + + assert delete_jobs() is True + + +@pytest.mark.skip() +def test_find_jobs() -> None: + pass + + +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 00000000..84e86444 --- /dev/null +++ b/tests/test_job_runner.py @@ -0,0 +1,228 @@ +import subprocess + +import pytest + +from facefusion import state_manager +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, submit_job, submit_jobs +from facefusion.jobs.job_runner import collect_output_set, finalize_steps, run_job, run_jobs, run_steps +from facefusion.typing 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') ]) + state_manager.init_item('output_audio_encoder', 'aac') + + +@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-1.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-1.jpg') + } + + assert run_jobs(process_step) 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) is False + + submit_jobs() + + assert run_jobs(process_step) is True + + +@pytest.mark.skip() +def test_retry_job() -> None: + pass + + +@pytest.mark.skip() +def test_retry_jobs() -> None: + pass + + +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-1.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-1.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-1-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-1.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-1.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-1.jpg'): + [ + get_test_output_file('output-1-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 00000000..c1d8a387 --- /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_normalizer.py b/tests/test_normalizer.py index 3273c58d..0673f64f 100644 --- a/tests/test_normalizer.py +++ b/tests/test_normalizer.py @@ -1,18 +1,4 @@ -from facefusion.common_helper import is_linux, is_macos -from facefusion.normalizer import normalize_output_path, normalize_padding, normalize_fps - - -def test_normalize_output_path() -> None: - if is_linux() or is_macos(): - assert normalize_output_path('.assets/examples/target-240p.mp4', '.assets/examples/target-240p.mp4') == '.assets/examples/target-240p.mp4' - assert normalize_output_path('.assets/examples/target-240p.mp4', '.assets/examples').startswith('.assets/examples/target-240p') - assert normalize_output_path('.assets/examples/target-240p.mp4', '.assets/examples').endswith('.mp4') - assert normalize_output_path('.assets/examples/target-240p.mp4', '.assets/examples/output.mp4') == '.assets/examples/output.mp4' - assert normalize_output_path('.assets/examples/target-240p.mp4', '.assets/examples/invalid') is None - assert normalize_output_path('.assets/examples/target-240p.mp4', '.assets/invalid/output.mp4') is None - assert normalize_output_path('.assets/examples/target-240p.mp4', 'invalid') is None - assert normalize_output_path('.assets/examples/target-240p.mp4', None) is None - assert normalize_output_path(None, '.assets/examples/output.mp4') is None +from facefusion.normalizer import normalize_fps, normalize_padding def test_normalize_padding() -> None: diff --git a/tests/test_process_manager.py b/tests/test_process_manager.py index 1fbe74bb..85e64645 100644 --- a/tests/test_process_manager.py +++ b/tests/test_process_manager.py @@ -1,4 +1,4 @@ -from facefusion.process_manager import set_process_state, is_processing, is_stopping, is_pending, start, stop, end +from facefusion.process_manager import end, is_pending, is_processing, is_stopping, set_process_state, start, stop def test_start() -> None: diff --git a/tests/test_program_helper.py b/tests/test_program_helper.py new file mode 100644 index 00000000..92b64fb2 --- /dev/null +++ b/tests/test_program_helper.py @@ -0,0 +1,40 @@ +from argparse import ArgumentParser + +import pytest + +from facefusion.program_helper import find_argument_group, validate_actions + + +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, 'invalid') is None + + +@pytest.mark.skip() +def test_validate_args() -> None: + pass + + +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_temp_helper.py b/tests/test_temp_helper.py new file mode 100644 index 00000000..48aad129 --- /dev/null +++ b/tests/test_temp_helper.py @@ -0,0 +1,33 @@ +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_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_vision.py b/tests/test_vision.py index 6cf48647..7cb69860 100644 --- a/tests/test_vision.py +++ b/tests/test_vision.py @@ -1,41 +1,43 @@ import subprocess + import pytest from facefusion.download import conditional_download -from facefusion.vision import detect_image_resolution, restrict_image_resolution, create_image_resolutions, get_video_frame, count_video_frame_total, detect_video_fps, restrict_video_fps, detect_video_resolution, restrict_video_resolution, create_video_resolutions, normalize_resolution, pack_resolution, unpack_resolution +from facefusion.vision import count_video_frame_total, create_image_resolutions, create_video_resolutions, detect_image_resolution, detect_video_fps, detect_video_resolution, get_video_frame, normalize_resolution, pack_resolution, restrict_image_resolution, restrict_video_fps, restrict_video_resolution, unpack_resolution +from .helper import get_test_example_file, get_test_examples_directory @pytest.fixture(scope = 'module', autouse = True) def before_all() -> None: - conditional_download('.assets/examples', + conditional_download(get_test_examples_directory(), [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-1080p.mp4' + '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', '.assets/examples/target-240p.mp4', '-vframes', '1', '.assets/examples/target-240p.jpg' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-1080p.mp4', '-vframes', '1', '.assets/examples/target-1080p.jpg' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vframes', '1', '-vf', 'transpose=0', '.assets/examples/target-240p-90deg.jpg' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-1080p.mp4', '-vframes', '1', '-vf', 'transpose=0', '.assets/examples/target-1080p-90deg.jpg' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=25', '.assets/examples/target-240p-25fps.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=30', '.assets/examples/target-240p-30fps.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=60', '.assets/examples/target-240p-60fps.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'transpose=0', '.assets/examples/target-240p-90deg.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-1080p.mp4', '-vf', 'transpose=0', '.assets/examples/target-1080p-90deg.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-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', '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') ]) def test_detect_image_resolution() -> None: - assert detect_image_resolution('.assets/examples/target-240p.jpg') == (426, 226) - assert detect_image_resolution('.assets/examples/target-240p-90deg.jpg') == (226, 426) - assert detect_image_resolution('.assets/examples/target-1080p.jpg') == (2048, 1080) - assert detect_image_resolution('.assets/examples/target-1080p-90deg.jpg') == (1080, 2048) + 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('.assets/examples/target-1080p.jpg', (426, 226)) == (426, 226) - assert restrict_image_resolution('.assets/examples/target-1080p.jpg', (2048, 1080)) == (2048, 1080) - assert restrict_image_resolution('.assets/examples/target-1080p.jpg', (4096, 2160)) == (2048, 1080) + 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_create_image_resolutions() -> None: @@ -47,42 +49,42 @@ def test_create_image_resolutions() -> None: def test_get_video_frame() -> None: - assert get_video_frame('.assets/examples/target-240p-25fps.mp4') is not None + assert get_video_frame(get_test_example_file('target-240p-25fps.mp4')) is not None assert get_video_frame('invalid') is None def test_count_video_frame_total() -> None: - assert count_video_frame_total('.assets/examples/target-240p-25fps.mp4') == 270 - assert count_video_frame_total('.assets/examples/target-240p-30fps.mp4') == 324 - assert count_video_frame_total('.assets/examples/target-240p-60fps.mp4') == 648 + 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_detect_video_fps() -> None: - assert detect_video_fps('.assets/examples/target-240p-25fps.mp4') == 25.0 - assert detect_video_fps('.assets/examples/target-240p-30fps.mp4') == 30.0 - assert detect_video_fps('.assets/examples/target-240p-60fps.mp4') == 60.0 + 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('.assets/examples/target-1080p.mp4', 20.0) == 20.0 - assert restrict_video_fps('.assets/examples/target-1080p.mp4', 25.0) == 25.0 - assert restrict_video_fps('.assets/examples/target-1080p.mp4', 60.0) == 25.0 + 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_resolution() -> None: - assert detect_video_resolution('.assets/examples/target-240p.mp4') == (426, 226) - assert detect_video_resolution('.assets/examples/target-240p-90deg.mp4') == (226, 426) - assert detect_video_resolution('.assets/examples/target-1080p.mp4') == (2048, 1080) - assert detect_video_resolution('.assets/examples/target-1080p-90deg.mp4') == (1080, 2048) + 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('.assets/examples/target-1080p.mp4', (426, 226)) == (426, 226) - assert restrict_video_resolution('.assets/examples/target-1080p.mp4', (2048, 1080)) == (2048, 1080) - assert restrict_video_resolution('.assets/examples/target-1080p.mp4', (4096, 2160)) == (2048, 1080) + 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_create_video_resolutions() -> None: diff --git a/tests/test_wording.py b/tests/test_wording.py index 1deaa773..5d987f9e 100644 --- a/tests/test_wording.py +++ b/tests/test_wording.py @@ -3,5 +3,5 @@ from facefusion import wording def test_get() -> None: assert wording.get('python_not_supported') - assert wording.get('help.source') + assert wording.get('help.source_paths') assert wording.get('invalid') is None