Provide different XSeg models

facefusion/choices.py

@@ -17,7 +17,7 @@ face_selector_modes : List[FaceSelectorMode] = [ 'many', 'one', 'reference' ]
 face_selector_orders : List[FaceSelectorOrder] = [ 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best' ]
 face_selector_genders : List[Gender] = [ 'female', 'male' ]
 face_selector_races : List[Race] = [ 'white', 'black', 'latino', 'asian', 'indian', 'arabic' ]
-face_occluder_models : List[FaceOccluderModel] = [ 'xseg_groggy_5' ]
+face_occluder_models : List[FaceOccluderModel] = [ 'xseg_1', 'xseg_2' ]
 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' ]

facefusion/face_masker.py

@@ -30,22 +30,42 @@ FACE_MASK_REGIONS : Dict[FaceMaskRegion, int] =\
 def create_static_model_set(download_scope : DownloadScope) -> ModelSet:
 	return\
 	{
-		'xseg_groggy_5':
+		'xseg_1':
 		{
 			'hashes':
 			{
 				'face_occluder':
 				{
-					'url': resolve_download_url('models-3.1.0', 'xseg_groggy_5.hash'),
+					'url': resolve_download_url('models-3.1.0', 'xseg_1.hash'),
-					'path': resolve_relative_path('../.assets/models/xseg_groggy_5.hash')
+					'path': resolve_relative_path('../.assets/models/xseg_1.hash')
 				}
 			},
 			'sources':
 			{
 				'face_occluder':
 				{
-					'url': resolve_download_url('models-3.1.0', 'xseg_groggy_5.onnx'),
+					'url': resolve_download_url('models-3.1.0', 'xseg_1.onnx'),
-					'path': resolve_relative_path('../.assets/models/xseg_groggy_5.onnx')
+					'path': resolve_relative_path('../.assets/models/xseg_1.onnx')
+				}
+			},
+			'size': (256, 256)
+		},
+		'xseg_2':
+		{
+			'hashes':
+			{
+				'face_occluder':
+				{
+					'url': resolve_download_url('models-3.1.0', 'xseg_2.hash'),
+					'path': resolve_relative_path('../.assets/models/xseg_2.hash')
+				}
+			},
+			'sources':
+			{
+				'face_occluder':
+				{
+					'url': resolve_download_url('models-3.1.0', 'xseg_2.onnx'),
+					'path': resolve_relative_path('../.assets/models/xseg_2.onnx')
 				}
 			},
 			'size': (256, 256)
@@ -107,9 +127,13 @@ def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]:
 	model_sources = {}
 	model_set = create_static_model_set('full')
 
-	if state_manager.get_item('face_occluder_model') == 'xseg_groggy_5':
+	if state_manager.get_item('face_occluder_model') == 'xseg_1':
-		model_hashes['xseg_groggy_5'] = model_set.get('xseg_groggy_5').get('hashes').get('face_occluder')
+		model_hashes['xseg_1'] = model_set.get('xseg_1').get('hashes').get('face_occluder')
-		model_sources['xseg_groggy_5'] = model_set.get('xseg_groggy_5').get('sources').get('face_occluder')
+		model_sources['xseg_1'] = model_set.get('xseg_1').get('sources').get('face_occluder')
+
+	if state_manager.get_item('face_occluder_model') == 'xseg_2':
+		model_hashes['xseg_2'] = model_set.get('xseg_2').get('hashes').get('face_occluder')
+		model_sources['xseg_2'] = model_set.get('xseg_2').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')

facefusion/program.py

@@ -132,7 +132,7 @@ 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-occluder-model', help = wording.get('help.face_occluder_model'), default = config.get_str_value('face_detector.face_occluder_model', 'xseg_1'), 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))

facefusion/typing.py

@@ -101,7 +101,7 @@ 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']
+FaceOccluderModel = Literal['xseg_1', 'xseg_2']
 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']