Limit download scopes to force-download command

This commit is contained in:
henryruhs 2024-11-26 15:58:03 +01:00
parent 4de4d18dda
commit 92d897057d
19 changed files with 38 additions and 51 deletions

View File

@ -43,7 +43,7 @@ download_provider_set : DownloadProviderSet =\
'github': 'https://github.com/facefusion/facefusion-assets/releases/download/{base_name}/{file_name}',
'huggingface': 'https://huggingface.co/facefusion/{base_name}/resolve/main/{file_name}'
}
download_scopes : List[DownloadScope] = [ 'lite', 'medium', 'full' ]
download_scopes : List[DownloadScope] = [ 'lite', 'full' ]
log_level_set : LogLevelSet =\
{

View File

@ -54,8 +54,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get('open_nsfw')
return create_static_model_set('full').get('open_nsfw')
def pre_check() -> bool:

View File

@ -3,7 +3,7 @@ from typing import List, Tuple
import numpy
from facefusion import inference_manager, state_manager
from facefusion import inference_manager
from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url
from facefusion.face_helper import warp_face_by_face_landmark_5
from facefusion.filesystem import resolve_relative_path
@ -51,8 +51,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get('fairface')
return create_static_model_set('full').get('fairface')
def pre_check() -> bool:

View File

@ -91,8 +91,7 @@ def clear_inference_pool() -> None:
def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]:
model_hashes = {}
model_sources = {}
download_scope = state_manager.get_item('download_scope')
model_set = create_static_model_set(download_scope)
model_set = create_static_model_set('full')
if state_manager.get_item('face_detector_model') in [ 'many', 'retinaface' ]:
model_hashes['retinaface'] = model_set.get('retinaface').get('hashes').get('retinaface')

View File

@ -90,8 +90,7 @@ def clear_inference_pool() -> None:
def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]:
download_scope = state_manager.get_item('download_scope')
model_set = create_static_model_set(download_scope)
model_set = create_static_model_set('full')
model_hashes =\
{
'fan_68_5': model_set.get('fan_68_5').get('hashes').get('fan_68_5')
@ -133,8 +132,7 @@ def detect_face_landmarks(vision_frame : VisionFrame, bounding_box : BoundingBox
def detect_with_2dfan4(temp_vision_frame: VisionFrame, bounding_box: BoundingBox, face_angle: Angle) -> Tuple[FaceLandmark68, Score]:
download_scope = state_manager.get_item('download_scope')
model_size = create_static_model_set(download_scope).get('2dfan4').get('size')
model_size = create_static_model_set('full').get('2dfan4').get('size')
scale = 195 / numpy.subtract(bounding_box[2:], bounding_box[:2]).max().clip(1, None)
translation = (model_size[0] - numpy.add(bounding_box[2:], bounding_box[:2]) * scale) * 0.5
rotated_matrix, rotated_size = create_rotated_matrix_and_size(face_angle, model_size)
@ -153,8 +151,7 @@ def detect_with_2dfan4(temp_vision_frame: VisionFrame, bounding_box: BoundingBox
def detect_with_peppa_wutz(temp_vision_frame : VisionFrame, bounding_box : BoundingBox, face_angle : Angle) -> Tuple[FaceLandmark68, Score]:
download_scope = state_manager.get_item('download_scope')
model_size = create_static_model_set(download_scope).get('peppa_wutz').get('size')
model_size = create_static_model_set('full').get('peppa_wutz').get('size')
scale = 195 / numpy.subtract(bounding_box[2:], bounding_box[:2]).max().clip(1, None)
translation = (model_size[0] - numpy.add(bounding_box[2:], bounding_box[:2]) * scale) * 0.5
rotated_matrix, rotated_size = create_rotated_matrix_and_size(face_angle, model_size)

View File

@ -5,7 +5,7 @@ import cv2
import numpy
from cv2.typing import Size
from facefusion import inference_manager, state_manager
from facefusion import inference_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
@ -83,8 +83,7 @@ def clear_inference_pool() -> None:
def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]:
download_scope = state_manager.get_item('download_scope')
model_set = create_static_model_set(download_scope)
model_set = create_static_model_set('full')
model_hashes =\
{
'face_occluder': model_set.get('face_occluder').get('hashes').get('face_occluder'),
@ -119,8 +118,7 @@ def create_static_box_mask(crop_size : Size, face_mask_blur : float, face_mask_p
def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask:
download_scope = state_manager.get_item('download_scope')
model_size = create_static_model_set(download_scope).get('face_occluder').get('size')
model_size = create_static_model_set('full').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)
@ -132,8 +130,7 @@ def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask:
def create_region_mask(crop_vision_frame : VisionFrame, face_mask_regions : List[FaceMaskRegion]) -> Mask:
download_scope = state_manager.get_item('download_scope')
model_size = create_static_model_set(download_scope).get('face_parser').get('size')
model_size = create_static_model_set('full').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))

View File

@ -3,7 +3,7 @@ from typing import Tuple
import numpy
from facefusion import inference_manager, state_manager
from facefusion import inference_manager
from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url
from facefusion.face_helper import warp_face_by_face_landmark_5
from facefusion.filesystem import resolve_relative_path
@ -49,8 +49,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get('arcface')
return create_static_model_set('full').get('arcface')
def pre_check() -> bool:

View File

@ -76,8 +76,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
age_modifier_model = state_manager.get_item('age_modifier_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(age_modifier_model)
return create_static_model_set('full').get(age_modifier_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -104,7 +104,7 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet:
('edel', 'sidney_sweeney_224', (224, 224)),
('edel', 'winona_ryder_224', (224, 224))
])
if download_scope == 'medium' or download_scope == 'full':
if download_scope in [ 'lite', 'full' ]:
model_config.extend(
[
('iperov', 'alexandra_daddario_224', (224, 224)),
@ -230,8 +230,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
deep_swapper_model = state_manager.get_item('deep_swapper_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(deep_swapper_model)
return create_static_model_set('full').get(deep_swapper_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -87,8 +87,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
expression_restorer_model = state_manager.get_item('expression_restorer_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(expression_restorer_model)
return create_static_model_set('full').get(expression_restorer_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -117,8 +117,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
face_editor_model = state_manager.get_item('face_editor_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(face_editor_model)
return create_static_model_set('full').get(face_editor_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -234,8 +234,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
face_enhancer_model = state_manager.get_item('face_enhancer_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(face_enhancer_model)
return create_static_model_set('full').get(face_enhancer_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -348,8 +348,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
face_swapper_model = state_manager.get_item('face_swapper_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(face_swapper_model)
return create_static_model_set('full').get(face_swapper_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -140,8 +140,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
frame_colorizer_model = state_manager.get_item('frame_colorizer_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(frame_colorizer_model)
return create_static_model_set('full').get(frame_colorizer_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -397,8 +397,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
frame_enhancer_model = state_manager.get_item('frame_enhancer_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(frame_enhancer_model)
return create_static_model_set('full').get(frame_enhancer_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -86,8 +86,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
lip_syncer_model = state_manager.get_item('lip_syncer_model')
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get(lip_syncer_model)
return create_static_model_set('full').get(lip_syncer_model)
def register_args(program : ArgumentParser) -> None:

View File

@ -199,12 +199,19 @@ def create_execution_program() -> ArgumentParser:
return program
def create_download_program() -> ArgumentParser:
def create_download_providers_program() -> ArgumentParser:
program = ArgumentParser(add_help = False)
group_download = program.add_argument_group('download')
group_download.add_argument('--download-providers', help = wording.get('help.download_providers').format(choices = ', '.join(list(facefusion.choices.download_provider_set.keys()))), default = config.get_str_list('download.download_providers', 'github'), choices = list(facefusion.choices.download_provider_set.keys()), nargs = '+', metavar = 'DOWNLOAD_PROVIDERS')
job_store.register_job_keys([ 'download_providers' ])
return program
def create_download_scope_program() -> ArgumentParser:
program = ArgumentParser(add_help = False)
group_download = program.add_argument_group('download')
group_download.add_argument('--download-scope', help = wording.get('help.download_scope'), default = config.get_str_value('download.download_scope', 'lite'), choices = list(facefusion.choices.download_scopes))
job_store.register_job_keys([ 'download_providers', 'download_scope' ])
job_store.register_job_keys([ 'download_scope' ])
return program
@ -249,7 +256,7 @@ def collect_step_program() -> ArgumentParser:
def collect_job_program() -> ArgumentParser:
return ArgumentParser(parents= [ create_execution_program(), create_download_program(), create_memory_program(), create_log_level_program() ], add_help = False)
return ArgumentParser(parents= [ create_execution_program(), create_download_providers_program(), create_memory_program(), create_log_level_program() ], add_help = False)
def create_program() -> ArgumentParser:
@ -261,7 +268,7 @@ def create_program() -> ArgumentParser:
sub_program.add_parser('run', help = wording.get('help.run'), parents = [ create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), create_source_paths_program(), create_target_path_program(), create_output_path_program(), collect_step_program(), create_uis_program(), collect_job_program() ], formatter_class = create_help_formatter_large)
sub_program.add_parser('headless-run', help = wording.get('help.headless_run'), parents = [ create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), create_source_paths_program(), create_target_path_program(), create_output_path_program(), collect_step_program(), collect_job_program() ], formatter_class = create_help_formatter_large)
sub_program.add_parser('batch-run', help = wording.get('help.batch_run'), parents = [ create_config_path_program(), create_temp_path_program(), create_jobs_path_program(), create_source_pattern_program(), create_target_pattern_program(), create_output_pattern_program(), collect_step_program(), collect_job_program() ], formatter_class = create_help_formatter_large)
sub_program.add_parser('force-download', help = wording.get('help.force_download'), parents = [ create_download_program(), create_log_level_program() ], formatter_class = create_help_formatter_large)
sub_program.add_parser('force-download', help = wording.get('help.force_download'), parents = [ create_download_providers_program(), create_download_scope_program(), create_log_level_program() ], formatter_class = create_help_formatter_large)
# job manager
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-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)

View File

@ -119,7 +119,7 @@ ExecutionProviderSet = Dict[ExecutionProviderKey, ExecutionProviderValue]
DownloadProviderKey = Literal['github', 'huggingface']
DownloadProviderSet = Dict[DownloadProviderKey, str]
DownloadScope = Literal['lite', 'medium', 'full']
DownloadScope = Literal['lite', 'full']
Download = TypedDict('Download',
{
'url' : str,

View File

@ -4,7 +4,7 @@ from typing import Tuple
import numpy
import scipy
from facefusion import inference_manager, state_manager
from facefusion import inference_manager
from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url
from facefusion.filesystem import resolve_relative_path
from facefusion.thread_helper import thread_semaphore
@ -47,8 +47,7 @@ def clear_inference_pool() -> None:
def get_model_options() -> ModelOptions:
download_scope = state_manager.get_item('download_scope')
return create_static_model_set(download_scope).get('kim_vocal_2')
return create_static_model_set('full').get('kim_vocal_2')
def pre_check() -> bool: