Introduce face masker models

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 =

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')))

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_genders : List[Gender] = [ 'female', 'male' ]
-face_selector_races : List[Race] = ['white', 'black', 'latino', 'asian', 'indian', 'arabic']
+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' ]

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 =\
+
-	{
+	if state_manager.get_item('face_occluder_model') == 'xseg_groggy_5':
-		'face_occluder': model_set.get('face_occluder').get('hashes').get('face_occluder'),
+		model_hashes['xseg_groggy_5'] = model_set.get('xseg_groggy_5').get('hashes').get('face_occluder')
-		'face_parser': model_set.get('face_parser').get('hashes').get('face_parser')
+		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_sources =\
+		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')
-		'face_occluder': model_set.get('face_occluder').get('sources').get('face_occluder'),
+	if state_manager.get_item('face_parser_model') == 'bisenet_resnet_34':
-		'face_parser': model_set.get('face_parser').get('sources').get('face_parser')
+		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,

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 = '+')

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,

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)

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)
 

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',

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',