diff --git a/facefusion.ini b/facefusion.ini index e0bd2044..539b83f6 100644 --- a/facefusion.ini +++ b/facefusion.ini @@ -32,6 +32,8 @@ reference_face_distance = reference_frame_number = [face_masker] +face_occluder_model = +face_parser_model = face_mask_types = face_mask_blur = face_mask_padding = diff --git a/facefusion/args.py b/facefusion/args.py index 8ceab01c..f8ee798f 100644 --- a/facefusion/args.py +++ b/facefusion/args.py @@ -71,6 +71,8 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('reference_face_distance', args.get('reference_face_distance')) apply_state_item('reference_frame_number', args.get('reference_frame_number')) # face masker + apply_state_item('face_occluder_model', args.get('face_occluder_model')) + apply_state_item('face_parser_model', args.get('face_parser_model')) apply_state_item('face_mask_types', args.get('face_mask_types')) apply_state_item('face_mask_blur', args.get('face_mask_blur')) apply_state_item('face_mask_padding', normalize_padding(args.get('face_mask_padding'))) diff --git a/facefusion/choices.py b/facefusion/choices.py index 6e91560a..bcd239d1 100755 --- a/facefusion/choices.py +++ b/facefusion/choices.py @@ -2,7 +2,7 @@ import logging from typing import List, Sequence from facefusion.common_helper import create_float_range, create_int_range -from facefusion.typing import Angle, DownloadProvider, DownloadProviderSet, DownloadScope, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskRegion, FaceMaskType, FaceSelectorMode, FaceSelectorOrder, Gender, JobStatus, LogLevel, LogLevelSet, OutputAudioEncoder, OutputVideoEncoder, OutputVideoPreset, Race, Score, TempFrameFormat, UiWorkflow, VideoMemoryStrategy +from facefusion.typing import Angle, DownloadProvider, DownloadProviderSet, DownloadScope, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskRegion, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorMode, FaceSelectorOrder, Gender, JobStatus, LogLevel, LogLevelSet, OutputAudioEncoder, OutputVideoEncoder, OutputVideoPreset, Race, Score, TempFrameFormat, UiWorkflow, VideoMemoryStrategy face_detector_set : FaceDetectorSet =\ { @@ -15,8 +15,10 @@ face_detector_models : List[FaceDetectorModel] = list(face_detector_set.keys()) face_landmarker_models : List[FaceLandmarkerModel] = [ 'many', '2dfan4', 'peppa_wutz' ] face_selector_modes : List[FaceSelectorMode] = [ 'many', 'one', 'reference' ] face_selector_orders : List[FaceSelectorOrder] = [ 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best' ] -face_selector_genders : List[Gender] = ['female', 'male'] -face_selector_races : List[Race] = ['white', 'black', 'latino', 'asian', 'indian', 'arabic'] +face_selector_genders : List[Gender] = [ 'female', 'male' ] +face_selector_races : List[Race] = [ 'white', 'black', 'latino', 'asian', 'indian', 'arabic' ] +face_occluder_models : List[FaceOccluderModel] = [ 'xseg_groggy_5' ] +face_parser_models : List[FaceParserModel] = [ 'bisenet_resnet_18', 'bisenet_resnet_34' ] 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' ] diff --git a/facefusion/face_masker.py b/facefusion/face_masker.py index 52d3198b..007db421 100755 --- a/facefusion/face_masker.py +++ b/facefusion/face_masker.py @@ -5,7 +5,7 @@ import cv2 import numpy from cv2.typing import Size -from facefusion import inference_manager +from facefusion import inference_manager, state_manager from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url from facefusion.filesystem import resolve_relative_path from facefusion.thread_helper import conditional_thread_semaphore @@ -30,7 +30,7 @@ FACE_MASK_REGIONS : Dict[FaceMaskRegion, int] =\ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: return\ { - 'face_occluder': + 'xseg_groggy_5': { 'hashes': { @@ -50,7 +50,27 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: }, 'size': (256, 256) }, - 'face_parser': + 'bisenet_resnet_18': + { + 'hashes': + { + 'face_parser': + { + 'url': resolve_download_url('models-3.1.0', 'bisenet_resnet_18.hash'), + 'path': resolve_relative_path('../.assets/models/bisenet_resnet_18.hash') + } + }, + 'sources': + { + 'face_parser': + { + 'url': resolve_download_url('models-3.1.0', 'bisenet_resnet_18.onnx'), + 'path': resolve_relative_path('../.assets/models/bisenet_resnet_18.onnx') + } + }, + 'size': (512, 512) + }, + 'bisenet_resnet_34': { 'hashes': { @@ -83,17 +103,19 @@ def clear_inference_pool() -> None: def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]: + model_hashes = {} + model_sources = {} model_set = create_static_model_set('full') - 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') - } + + if state_manager.get_item('face_occluder_model') == 'xseg_groggy_5': + model_hashes['xseg_groggy_5'] = model_set.get('xseg_groggy_5').get('hashes').get('face_occluder') + model_sources['xseg_groggy_5'] = model_set.get('xseg_groggy_5').get('sources').get('face_occluder') + if state_manager.get_item('face_parser_model') == 'bisenet_resnet_18': + model_hashes['bisenet_resnet_18'] = model_set.get('bisenet_resnet_18').get('hashes').get('face_parser') + model_sources['bisenet_resnet_18'] = model_set.get('bisenet_resnet_18').get('sources').get('face_parser') + if state_manager.get_item('face_parser_model') == 'bisenet_resnet_34': + model_hashes['bisenet_resnet_34'] = model_set.get('bisenet_resnet_34').get('hashes').get('face_parser') + model_sources['bisenet_resnet_34'] = model_set.get('bisenet_resnet_34').get('sources').get('face_parser') return model_hashes, model_sources @@ -118,7 +140,8 @@ def create_static_box_mask(crop_size : Size, face_mask_blur : float, face_mask_p def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask: - model_size = create_static_model_set('full').get('face_occluder').get('size') + face_occluder_model = state_manager.get_item('face_occluder_model') + model_size = create_static_model_set('full').get(face_occluder_model).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) @@ -130,7 +153,8 @@ def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask: def create_region_mask(crop_vision_frame : VisionFrame, face_mask_regions : List[FaceMaskRegion]) -> Mask: - model_size = create_static_model_set('full').get('face_parser').get('size') + face_parser_model = state_manager.get_item('face_parser_model') + model_size = create_static_model_set('full').get(face_parser_model).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)) @@ -154,7 +178,8 @@ def create_mouth_mask(face_landmark_68 : FaceLandmark68) -> Mask: def forward_occlude_face(prepare_vision_frame : VisionFrame) -> Mask: - face_occluder = get_inference_pool().get('face_occluder') + face_occluder_model = state_manager.get_item('face_occluder_model') + face_occluder = get_inference_pool().get(face_occluder_model) with conditional_thread_semaphore(): occlusion_mask : Mask = face_occluder.run(None, @@ -166,7 +191,8 @@ def forward_occlude_face(prepare_vision_frame : VisionFrame) -> Mask: def forward_parse_face(prepare_vision_frame : VisionFrame) -> Mask: - face_parser = get_inference_pool().get('face_parser') + face_parser_model = state_manager.get_item('face_parser_model') + face_parser = get_inference_pool().get(face_parser_model) with conditional_thread_semaphore(): region_mask : Mask = face_parser.run(None, diff --git a/facefusion/program.py b/facefusion/program.py index 9af0fe09..8a07df67 100755 --- a/facefusion/program.py +++ b/facefusion/program.py @@ -132,6 +132,8 @@ def create_face_selector_program() -> ArgumentParser: def create_face_masker_program() -> ArgumentParser: program = ArgumentParser(add_help = False) group_face_masker = program.add_argument_group('face masker') + group_face_masker.add_argument('--face-occluder-model', help = wording.get('help.face_occluder_model'), default = config.get_str_value('face_detector.face_occluder_model', 'xseg_groggy_5'), choices = facefusion.choices.face_occluder_models) + group_face_masker.add_argument('--face-parser-model', help = wording.get('help.face_parser_model'), default = config.get_str_value('face_detector.face_parser_model', 'bisenet_resnet_34'), choices = facefusion.choices.face_parser_models) 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 = '+') diff --git a/facefusion/typing.py b/facefusion/typing.py index 6785debe..90f388af 100755 --- a/facefusion/typing.py +++ b/facefusion/typing.py @@ -101,6 +101,8 @@ FaceLandmarkerModel = Literal['many', '2dfan4', 'peppa_wutz'] FaceDetectorSet = Dict[FaceDetectorModel, List[str]] FaceSelectorMode = Literal['many', 'one', 'reference'] FaceSelectorOrder = Literal['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best'] +FaceOccluderModel = Literal['xseg_groggy_5'] +FaceParserModel = Literal['bisenet_resnet_18', 'bisenet_resnet_34'] 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['bmp', 'jpg', 'png'] @@ -231,6 +233,8 @@ StateKey = Literal\ 'reference_face_position', 'reference_face_distance', 'reference_frame_number', + 'face_occluder_model', + 'face_parser_model', 'face_mask_types', 'face_mask_blur', 'face_mask_padding', @@ -292,6 +296,8 @@ State = TypedDict('State', 'reference_face_position' : int, 'reference_face_distance' : float, 'reference_frame_number' : int, + 'face_occluder_model' : FaceOccluderModel, + 'face_parser_model' : FaceParserModel, 'face_mask_types' : List[FaceMaskType], 'face_mask_blur' : float, 'face_mask_padding' : Padding, diff --git a/facefusion/uis/components/face_masker.py b/facefusion/uis/components/face_masker.py index 46571c5b..03c35457 100755 --- a/facefusion/uis/components/face_masker.py +++ b/facefusion/uis/components/face_masker.py @@ -3,11 +3,13 @@ from typing import List, Optional, Tuple import gradio import facefusion.choices -from facefusion import state_manager, wording +from facefusion import face_masker, state_manager, wording from facefusion.common_helper import calc_float_step, calc_int_step -from facefusion.typing import FaceMaskRegion, FaceMaskType +from facefusion.typing import FaceMaskRegion, FaceMaskType, FaceOccluderModel, FaceParserModel from facefusion.uis.core import register_ui_component +FACE_OCCLUDER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None +FACE_PARSER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None FACE_MASK_TYPES_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None FACE_MASK_REGIONS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None FACE_MASK_BLUR_SLIDER : Optional[gradio.Slider] = None @@ -18,6 +20,8 @@ FACE_MASK_PADDING_LEFT_SLIDER : Optional[gradio.Slider] = None def render() -> None: + global FACE_OCCLUDER_MODEL_DROPDOWN + global FACE_PARSER_MODEL_DROPDOWN global FACE_MASK_TYPES_CHECKBOX_GROUP global FACE_MASK_REGIONS_CHECKBOX_GROUP global FACE_MASK_BLUR_SLIDER @@ -26,8 +30,20 @@ def render() -> None: global FACE_MASK_PADDING_BOTTOM_SLIDER global FACE_MASK_PADDING_LEFT_SLIDER + has_box_mask = 'box' in state_manager.get_item('face_mask_types') has_region_mask = 'region' in state_manager.get_item('face_mask_types') + with gradio.Row(): + FACE_OCCLUDER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.face_occluder_model_dropdown'), + choices = facefusion.choices.face_occluder_models, + value = state_manager.get_item('face_occluder_model') + ) + FACE_PARSER_MODEL_DROPDOWN = gradio.Dropdown( + label = wording.get('uis.face_parser_model_dropdown'), + choices = facefusion.choices.face_parser_models, + value = state_manager.get_item('face_parser_model') + ) FACE_MASK_TYPES_CHECKBOX_GROUP = gradio.CheckboxGroup( label = wording.get('uis.face_mask_types_checkbox_group'), choices = facefusion.choices.face_mask_types, @@ -82,6 +98,8 @@ def render() -> None: value = state_manager.get_item('face_mask_padding')[3], visible = has_box_mask ) + register_ui_component('face_occluder_model_dropdown', FACE_OCCLUDER_MODEL_DROPDOWN) + register_ui_component('face_parser_model_dropdown', FACE_PARSER_MODEL_DROPDOWN) register_ui_component('face_mask_types_checkbox_group', FACE_MASK_TYPES_CHECKBOX_GROUP) register_ui_component('face_mask_regions_checkbox_group', FACE_MASK_REGIONS_CHECKBOX_GROUP) register_ui_component('face_mask_blur_slider', FACE_MASK_BLUR_SLIDER) @@ -92,6 +110,8 @@ def render() -> None: def listen() -> None: + FACE_OCCLUDER_MODEL_DROPDOWN.change(update_face_occluder_model, inputs = FACE_OCCLUDER_MODEL_DROPDOWN) + FACE_PARSER_MODEL_DROPDOWN.change(update_face_parser_model, inputs = FACE_PARSER_MODEL_DROPDOWN) FACE_MASK_TYPES_CHECKBOX_GROUP.change(update_face_mask_types, inputs = FACE_MASK_TYPES_CHECKBOX_GROUP, outputs = [ FACE_MASK_TYPES_CHECKBOX_GROUP, FACE_MASK_REGIONS_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_REGIONS_CHECKBOX_GROUP.change(update_face_mask_regions, inputs = FACE_MASK_REGIONS_CHECKBOX_GROUP, outputs = FACE_MASK_REGIONS_CHECKBOX_GROUP) FACE_MASK_BLUR_SLIDER.release(update_face_mask_blur, inputs = FACE_MASK_BLUR_SLIDER) @@ -100,6 +120,24 @@ def listen() -> None: face_mask_padding_slider.release(update_face_mask_padding, inputs = face_mask_padding_sliders) +def update_face_occluder_model(face_occluder_model : FaceOccluderModel) -> gradio.Dropdown: + face_masker.clear_inference_pool() + state_manager.set_item('face_occluder_model', face_occluder_model) + + if face_masker.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('face_occluder_model')) + return gradio.Dropdown() + + +def update_face_parser_model(face_parser_model : FaceParserModel) -> gradio.Dropdown: + face_masker.clear_inference_pool() + state_manager.set_item('face_parser_model', face_parser_model) + + if face_masker.pre_check(): + return gradio.Dropdown(value = state_manager.get_item('face_parser_model')) + return gradio.Dropdown() + + def update_face_mask_types(face_mask_types : List[FaceMaskType]) -> Tuple[gradio.CheckboxGroup, gradio.CheckboxGroup, gradio.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) diff --git a/facefusion/uis/components/preview.py b/facefusion/uis/components/preview.py index 75cf9ed9..a38338f7 100755 --- a/facefusion/uis/components/preview.py +++ b/facefusion/uis/components/preview.py @@ -162,7 +162,9 @@ def listen() -> None: 'face_detector_model_dropdown', 'face_detector_size_dropdown', 'face_detector_angles_checkbox_group', - 'face_landmarker_model_dropdown' + 'face_landmarker_model_dropdown', + 'face_occluder_model_dropdown', + 'face_parser_model_dropdown' ]): ui_component.change(clear_and_update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE) diff --git a/facefusion/uis/typing.py b/facefusion/uis/typing.py index f7e894cc..78c6f718 100644 --- a/facefusion/uis/typing.py +++ b/facefusion/uis/typing.py @@ -52,6 +52,8 @@ ComponentName = Literal\ 'face_selector_race_dropdown', 'face_swapper_model_dropdown', 'face_swapper_pixel_boost_dropdown', + 'face_occluder_model_dropdown', + 'face_parser_model_dropdown', 'frame_colorizer_blend_slider', 'frame_colorizer_model_dropdown', 'frame_colorizer_size_dropdown', diff --git a/facefusion/wording.py b/facefusion/wording.py index 07823ba8..5d5a9f01 100755 --- a/facefusion/wording.py +++ b/facefusion/wording.py @@ -126,6 +126,8 @@ WORDING : Dict[str, Any] =\ 'reference_face_distance': 'specify the similarity between the reference face and target face', 'reference_frame_number': 'specify the frame used to create the reference face', # face masker + 'face_occluder_model': 'choose the model responsible for occluding the face', + 'face_parser_model': 'choose the model responsible for parsing the face', 'face_mask_types': 'mix and match different face mask types (choices: {choices})', 'face_mask_blur': 'specify the degree of blur applied to the box mask', 'face_mask_padding': 'apply top, right, bottom and left padding to the box mask', @@ -285,6 +287,8 @@ WORDING : Dict[str, Any] =\ 'face_selector_race_dropdown': 'FACE SELECTOR RACE', 'face_swapper_model_dropdown': 'FACE SWAPPER MODEL', 'face_swapper_pixel_boost_dropdown': 'FACE SWAPPER PIXEL BOOST', + 'face_occluder_model_dropdown': 'FACE OCCLUDER MODEL', + 'face_parser_model_dropdown': 'FACE PARSER MODEL', 'frame_colorizer_blend_slider': 'FRAME COLORIZER BLEND', 'frame_colorizer_model_dropdown': 'FRAME COLORIZER MODEL', 'frame_colorizer_size_dropdown': 'FRAME COLORIZER SIZE',