Upload 4 files

processors/modules/frame_colorizer/choices.py

@@ -0,0 +1,10 @@
+from typing import List, Sequence, get_args
+
+from facefusion.common_helper import create_int_range
+from facefusion.processors.modules.frame_colorizer.types import FrameColorizerModel
+
+frame_colorizer_models : List[FrameColorizerModel] = list(get_args(FrameColorizerModel))
+
+frame_colorizer_sizes : List[str] = [ '192x192', '256x256', '384x384', '512x512' ]
+
+frame_colorizer_blend_range : Sequence[int] = create_int_range(0, 100, 1)

processors/modules/frame_colorizer/core.py

@@ -0,0 +1,313 @@
+from argparse import ArgumentParser
+from functools import lru_cache
+from types import ModuleType
+from typing import List
+
+import cv2
+import numpy
+
+import facefusion.choices
+import facefusion.jobs.job_manager
+import facefusion.jobs.job_store
+from facefusion import config, content_analyser, inference_manager, logger, state_manager, translator, video_manager
+from facefusion.common_helper import create_int_metavar, is_macos
+from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url
+from facefusion.execution import has_execution_provider
+from facefusion.filesystem import in_directory, is_image, is_video, resolve_relative_path, same_file_extension
+from facefusion.processors.modules.frame_colorizer import choices as frame_colorizer_choices
+from facefusion.processors.modules.frame_colorizer.types import FrameColorizerInputs
+from facefusion.processors.types import ProcessorOutputs
+from facefusion.program_helper import find_argument_group
+from facefusion.thread_helper import thread_semaphore
+from facefusion.types import ApplyStateItem, Args, DownloadScope, InferencePool, InferenceProvider, ModelOptions, ModelSet, ProcessMode, VisionFrame
+from facefusion.vision import blend_frame, read_static_image, read_static_video_frame, unpack_resolution
+
+
+@lru_cache()
+def create_static_model_set(download_scope : DownloadScope) -> ModelSet:
+	return\
+	{
+		'ddcolor':
+		{
+			'__metadata__':
+			{
+				'vendor': 'piddnad',
+				'license': 'Apache-2.0',
+				'year': 2023
+			},
+			'hashes':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'ddcolor.hash'),
+					'path': resolve_relative_path('../.assets/models/ddcolor.hash')
+				}
+			},
+			'sources':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'ddcolor.onnx'),
+					'path': resolve_relative_path('../.assets/models/ddcolor.onnx')
+				}
+			},
+			'type': 'ddcolor'
+		},
+		'ddcolor_artistic':
+		{
+			'__metadata__':
+			{
+				'vendor': 'piddnad',
+				'license': 'Apache-2.0',
+				'year': 2023
+			},
+			'hashes':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'ddcolor_artistic.hash'),
+					'path': resolve_relative_path('../.assets/models/ddcolor_artistic.hash')
+				}
+			},
+			'sources':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'ddcolor_artistic.onnx'),
+					'path': resolve_relative_path('../.assets/models/ddcolor_artistic.onnx')
+				}
+			},
+			'type': 'ddcolor'
+		},
+		'deoldify':
+		{
+			'__metadata__':
+			{
+				'vendor': 'jantic',
+				'license': 'MIT',
+				'year': 2022
+			},
+			'hashes':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'deoldify.hash'),
+					'path': resolve_relative_path('../.assets/models/deoldify.hash')
+				}
+			},
+			'sources':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'deoldify.onnx'),
+					'path': resolve_relative_path('../.assets/models/deoldify.onnx')
+				}
+			},
+			'type': 'deoldify'
+		},
+		'deoldify_artistic':
+		{
+			'__metadata__':
+			{
+				'vendor': 'jantic',
+				'license': 'MIT',
+				'year': 2022
+			},
+			'hashes':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'deoldify_artistic.hash'),
+					'path': resolve_relative_path('../.assets/models/deoldify_artistic.hash')
+				}
+			},
+			'sources':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'deoldify_artistic.onnx'),
+					'path': resolve_relative_path('../.assets/models/deoldify_artistic.onnx')
+				}
+			},
+			'type': 'deoldify'
+		},
+		'deoldify_stable':
+		{
+			'__metadata__':
+			{
+				'vendor': 'jantic',
+				'license': 'MIT',
+				'year': 2022
+			},
+			'hashes':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'deoldify_stable.hash'),
+					'path': resolve_relative_path('../.assets/models/deoldify_stable.hash')
+				}
+			},
+			'sources':
+			{
+				'frame_colorizer':
+				{
+					'url': resolve_download_url('models-3.0.0', 'deoldify_stable.onnx'),
+					'path': resolve_relative_path('../.assets/models/deoldify_stable.onnx')
+				}
+			},
+			'type': 'deoldify'
+		}
+	}
+
+
+def get_inference_pool() -> InferencePool:
+	model_names = [ state_manager.get_item('frame_colorizer_model') ]
+	model_source_set = get_model_options().get('sources')
+
+	return inference_manager.get_inference_pool(__name__, model_names, model_source_set)
+
+
+def clear_inference_pool() -> None:
+	model_names = [ state_manager.get_item('frame_colorizer_model') ]
+	inference_manager.clear_inference_pool(__name__, model_names)
+
+
+def override_inference_providers() -> List[InferenceProvider]:
+	if is_macos() and has_execution_provider('coreml'):
+		return [ facefusion.choices.execution_provider_set.get('cpu') ]
+
+	return []
+
+
+def get_model_options() -> ModelOptions:
+	model_name = state_manager.get_item('frame_colorizer_model')
+	return create_static_model_set('full').get(model_name)
+
+
+def register_args(program : ArgumentParser) -> None:
+	group_processors = find_argument_group(program, 'processors')
+	if group_processors:
+		group_processors.add_argument('--frame-colorizer-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'frame_colorizer_model', 'ddcolor'), choices = frame_colorizer_choices.frame_colorizer_models)
+		group_processors.add_argument('--frame-colorizer-size', help = translator.get('help.size', __package__), type = str, default = config.get_str_value('processors', 'frame_colorizer_size', '256x256'), choices = frame_colorizer_choices.frame_colorizer_sizes)
+		group_processors.add_argument('--frame-colorizer-blend', help = translator.get('help.blend', __package__), type = int, default = config.get_int_value('processors', 'frame_colorizer_blend', '100'), choices = frame_colorizer_choices.frame_colorizer_blend_range, metavar = create_int_metavar(frame_colorizer_choices.frame_colorizer_blend_range))
+		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 get_common_modules() -> List[ModuleType]:
+	return [ content_analyser ]
+
+
+def pre_check() -> bool:
+	model_hash_set = get_model_options().get('hashes')
+	model_source_set = get_model_options().get('sources')
+
+	for common_module in get_common_modules():
+		if not common_module.pre_check():
+			return False
+
+	return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set)
+
+
+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(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__)
+		return False
+	if mode == 'output' and not in_directory(state_manager.get_item('output_path')):
+		logger.error(translator.get('specify_image_or_video_output') + translator.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(translator.get('match_target_and_output_extension') + translator.get('exclamation_mark'), __name__)
+		return False
+	return True
+
+
+def post_process() -> None:
+	read_static_image.cache_clear()
+	read_static_video_frame.cache_clear()
+	video_manager.clear_video_pool()
+
+	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
+		clear_inference_pool()
+
+	if state_manager.get_item('video_memory_strategy') == 'strict':
+		for common_module in get_common_modules():
+			common_module.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_color_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_color_frame(temp_vision_frame : VisionFrame, color_vision_frame : VisionFrame) -> VisionFrame:
+	frame_colorizer_blend = 1 - (state_manager.get_item('frame_colorizer_blend') / 100)
+	temp_vision_frame = blend_frame(temp_vision_frame, color_vision_frame, 1 - frame_colorizer_blend)
+	return temp_vision_frame
+
+
+def process_frame(inputs : FrameColorizerInputs) -> ProcessorOutputs:
+	temp_vision_frame = inputs.get('temp_vision_frame')
+	temp_vision_mask = inputs.get('temp_vision_mask')
+	temp_vision_frame = colorize_frame(temp_vision_frame)
+	return temp_vision_frame, temp_vision_mask

processors/modules/frame_colorizer/locales.py

@@ -0,0 +1,20 @@
+from facefusion.types import Locales
+
+LOCALES : Locales =\
+{
+	'en':
+	{
+		'help':
+		{
+			'model': 'choose the model responsible for colorizing the frame',
+			'size': 'specify the frame size provided to the frame colorizer',
+			'blend': 'blend the colorized into the previous frame'
+		},
+		'uis':
+		{
+			'blend_slider': 'FRAME COLORIZER BLEND',
+			'model_dropdown': 'FRAME COLORIZER MODEL',
+			'size_dropdown': 'FRAME COLORIZER SIZE'
+		}
+	}
+}

processors/modules/frame_colorizer/types.py

@@ -0,0 +1,12 @@
+from typing import List, Literal, TypedDict
+
+from facefusion.types import Mask, VisionFrame
+
+FrameColorizerInputs = TypedDict('FrameColorizerInputs',
+{
+	'target_vision_frames' : List[VisionFrame],
+	'temp_vision_frame' : VisionFrame,
+	'temp_vision_mask' : Mask
+})
+
+FrameColorizerModel = Literal['ddcolor', 'ddcolor_artistic', 'deoldify', 'deoldify_artistic', 'deoldify_stable']