Initial commit

facefusion/processors/frame/core.py

@@ -0,0 +1,113 @@
+import os
+import sys
+import importlib
+import psutil
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from queue import Queue
+from types import ModuleType
+from typing import Any, List, Callable
+from tqdm import tqdm
+
+import facefusion.globals
+from facefusion import wording
+
+FRAME_PROCESSORS_MODULES : List[ModuleType] = []
+FRAME_PROCESSORS_METHODS =\
+[
+	'get_frame_processor',
+	'clear_frame_processor',
+	'pre_check',
+	'pre_process',
+	'process_frame',
+	'process_frames',
+	'process_image',
+	'process_video',
+	'post_process'
+]
+
+
+def load_frame_processor_module(frame_processor : str) -> Any:
+	try:
+		frame_processor_module = importlib.import_module('facefusion.processors.frame.modules.' + frame_processor)
+		for method_name in FRAME_PROCESSORS_METHODS:
+			if not hasattr(frame_processor_module, method_name):
+				raise NotImplementedError
+	except ModuleNotFoundError:
+		sys.exit(wording.get('frame_processor_not_loaded').format(frame_processor = frame_processor))
+	except NotImplementedError:
+		sys.exit(wording.get('frame_processor_not_implemented').format(frame_processor = frame_processor))
+	return frame_processor_module
+
+
+def get_frame_processors_modules(frame_processors : List[str]) -> List[ModuleType]:
+	global FRAME_PROCESSORS_MODULES
+
+	if not FRAME_PROCESSORS_MODULES:
+		for frame_processor in frame_processors:
+			frame_processor_module = load_frame_processor_module(frame_processor)
+			FRAME_PROCESSORS_MODULES.append(frame_processor_module)
+	return FRAME_PROCESSORS_MODULES
+
+
+def clear_frame_processors_modules() -> None:
+	global FRAME_PROCESSORS_MODULES
+
+	for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors):
+		frame_processor_module.clear_frame_processor()
+	FRAME_PROCESSORS_MODULES = []
+
+
+def multi_process_frame(source_path : str, temp_frame_paths : List[str], process_frames: Callable[[str, List[str], Any], None], update: Callable[[], None]) -> None:
+	with ThreadPoolExecutor(max_workers = facefusion.globals.execution_thread_count) as executor:
+		futures = []
+		queue = create_queue(temp_frame_paths)
+		queue_per_future = max(len(temp_frame_paths) // facefusion.globals.execution_thread_count * facefusion.globals.execution_queue_count, 1)
+		while not queue.empty():
+			future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
+			futures.append(future)
+		for future in as_completed(futures):
+			future.result()
+
+
+def create_queue(temp_frame_paths : List[str]) -> Queue[str]:
+	queue: Queue[str] = Queue()
+	for frame_path in temp_frame_paths:
+		queue.put(frame_path)
+	return queue
+
+
+def pick_queue(queue : Queue[str], queue_per_future : int) -> List[str]:
+	queues = []
+	for _ in range(queue_per_future):
+		if not queue.empty():
+			queues.append(queue.get())
+	return queues
+
+
+def process_video(source_path : str, frame_paths : List[str], process_frames : Callable[[str, List[str], Any], None]) -> None:
+	progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
+	total = len(frame_paths)
+	with tqdm(total = total, desc = wording.get('processing'), unit = 'frame', dynamic_ncols = True, bar_format = progress_bar_format) as progress:
+		multi_process_frame(source_path, frame_paths, process_frames, lambda: update_progress(progress))
+
+
+def update_progress(progress : Any = None) -> None:
+	process = psutil.Process(os.getpid())
+	memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
+	progress.set_postfix(
+	{
+		'memory_usage': '{:.2f}'.format(memory_usage).zfill(5) + 'GB',
+		'execution_providers': facefusion.globals.execution_providers,
+		'execution_thread_count': facefusion.globals.execution_thread_count,
+		'execution_queue_count': facefusion.globals.execution_queue_count
+	})
+	progress.refresh()
+	progress.update(1)
+
+
+def get_device() -> str:
+	if 'CUDAExecutionProvider' in facefusion.globals.execution_providers:
+		return 'cuda'
+	if 'CoreMLExecutionProvider' in facefusion.globals.execution_providers:
+		return 'mps'
+	return 'cpu'

facefusion/processors/frame/modules/face_enhancer.py

@@ -0,0 +1,100 @@
+from typing import Any, List, Callable
+import cv2
+import threading
+from gfpgan.utils import GFPGANer
+
+import facefusion.globals
+import facefusion.processors.frame.core as frame_processors
+from facefusion import wording
+from facefusion.core import update_status
+from facefusion.face_analyser import get_many_faces
+from facefusion.typing import Frame, Face
+from facefusion.utilities import conditional_download, resolve_relative_path, is_image, is_video
+
+FRAME_PROCESSOR = None
+THREAD_SEMAPHORE = threading.Semaphore()
+THREAD_LOCK = threading.Lock()
+NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_ENHANCER'
+
+
+def get_frame_processor() -> Any:
+	global FRAME_PROCESSOR
+
+	with THREAD_LOCK:
+		if FRAME_PROCESSOR is None:
+			model_path = resolve_relative_path('../.assets/models/GFPGANv1.4.pth')
+			FRAME_PROCESSOR = GFPGANer(
+				model_path = model_path,
+				upscale = 1,
+				device = frame_processors.get_device()
+			)
+	return FRAME_PROCESSOR
+
+
+def clear_frame_processor() -> None:
+	global FRAME_PROCESSOR
+
+	FRAME_PROCESSOR = None
+
+
+def pre_check() -> bool:
+	download_directory_path = resolve_relative_path('../.assets/models')
+	conditional_download(download_directory_path, ['https://huggingface.co/facefusion/models/resolve/main/GFPGANv1.4.pth'])
+	return True
+
+
+def pre_process() -> bool:
+	if not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path):
+		update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME)
+		return False
+	return True
+
+
+def post_process() -> None:
+	clear_frame_processor()
+
+
+def enhance_face(target_face : Face, temp_frame : Frame) -> Frame:
+	start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
+	padding_x = int((end_x - start_x) * 0.5)
+	padding_y = int((end_y - start_y) * 0.5)
+	start_x = max(0, start_x - padding_x)
+	start_y = max(0, start_y - padding_y)
+	end_x = max(0, end_x + padding_x)
+	end_y = max(0, end_y + padding_y)
+	crop_frame = temp_frame[start_y:end_y, start_x:end_x]
+	if crop_frame.size:
+		with THREAD_SEMAPHORE:
+			_, _, crop_frame = get_frame_processor().enhance(
+				crop_frame,
+				paste_back = True
+			)
+		temp_frame[start_y:end_y, start_x:end_x] = crop_frame
+	return temp_frame
+
+
+def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) -> Frame:
+	many_faces = get_many_faces(temp_frame)
+	if many_faces:
+		for target_face in many_faces:
+			temp_frame = enhance_face(target_face, temp_frame)
+	return temp_frame
+
+
+def process_frames(source_path : str, temp_frame_paths : List[str], update: Callable[[], None]) -> None:
+	for temp_frame_path in temp_frame_paths:
+		temp_frame = cv2.imread(temp_frame_path)
+		result_frame = process_frame(None, None, temp_frame)
+		cv2.imwrite(temp_frame_path, result_frame)
+		if update:
+			update()
+
+
+def process_image(source_path : str, target_path : str, output_path : str) -> None:
+	target_frame = cv2.imread(target_path)
+	result_frame = process_frame(None, None, target_frame)
+	cv2.imwrite(output_path, result_frame)
+
+
+def process_video(source_path : str, temp_frame_paths : List[str]) -> None:
+	facefusion.processors.frame.core.process_video(None, temp_frame_paths, process_frames)

facefusion/processors/frame/modules/face_swapper.py

@@ -0,0 +1,105 @@
+from typing import Any, List, Callable
+import cv2
+import insightface
+import threading
+
+import facefusion.globals
+import facefusion.processors.frame.core as frame_processors
+from facefusion import wording
+from facefusion.core import update_status
+from facefusion.face_analyser import get_one_face, get_many_faces, find_similar_faces
+from facefusion.face_reference import get_face_reference, set_face_reference
+from facefusion.typing import Face, Frame
+from facefusion.utilities import conditional_download, resolve_relative_path, is_image, is_video
+
+FRAME_PROCESSOR = None
+THREAD_LOCK = threading.Lock()
+NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_SWAPPER'
+
+
+def get_frame_processor() -> Any:
+	global FRAME_PROCESSOR
+
+	with THREAD_LOCK:
+		if FRAME_PROCESSOR is None:
+			model_path = resolve_relative_path('../.assets/models/inswapper_128.onnx')
+			FRAME_PROCESSOR = insightface.model_zoo.get_model(model_path, providers = facefusion.globals.execution_providers)
+	return FRAME_PROCESSOR
+
+
+def clear_frame_processor() -> None:
+	global FRAME_PROCESSOR
+
+	FRAME_PROCESSOR = None
+
+
+def pre_check() -> bool:
+	download_directory_path = resolve_relative_path('../.assets/models')
+	conditional_download(download_directory_path, ['https://huggingface.co/facefusion/models/resolve/main/inswapper_128.onnx'])
+	return True
+
+
+def pre_process() -> bool:
+	if not is_image(facefusion.globals.source_path):
+		update_status(wording.get('select_image_source') + wording.get('exclamation_mark'), NAME)
+		return False
+	elif not get_one_face(cv2.imread(facefusion.globals.source_path)):
+		update_status(wording.get('no_source_face_detected') + wording.get('exclamation_mark'), NAME)
+		return False
+	if not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path):
+		update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME)
+		return False
+	return True
+
+
+def post_process() -> None:
+	clear_frame_processor()
+
+
+def swap_face(source_face : Face, target_face : Face, temp_frame : Frame) -> Frame:
+	return get_frame_processor().get(temp_frame, target_face, source_face, paste_back = True)
+
+
+def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) -> Frame:
+	if 'reference' in facefusion.globals.face_recognition:
+		similar_faces = find_similar_faces(temp_frame, reference_face, facefusion.globals.reference_face_distance)
+		if similar_faces:
+			for similar_face in similar_faces:
+				temp_frame = swap_face(source_face, similar_face, temp_frame)
+	if 'many' in facefusion.globals.face_recognition:
+		many_faces = get_many_faces(temp_frame)
+		if many_faces:
+			for target_face in many_faces:
+				temp_frame = swap_face(source_face, target_face, temp_frame)
+	return temp_frame
+
+
+def process_frames(source_path : str, temp_frame_paths : List[str], update: Callable[[], None]) -> None:
+	source_face = get_one_face(cv2.imread(source_path))
+	reference_face = get_face_reference() if 'reference' in facefusion.globals.face_recognition else None
+	for temp_frame_path in temp_frame_paths:
+		temp_frame = cv2.imread(temp_frame_path)
+		result_frame = process_frame(source_face, reference_face, temp_frame)
+		cv2.imwrite(temp_frame_path, result_frame)
+		if update:
+			update()
+
+
+def process_image(source_path : str, target_path : str, output_path : str) -> None:
+	source_face = get_one_face(cv2.imread(source_path))
+	target_frame = cv2.imread(target_path)
+	reference_face = get_one_face(target_frame, facefusion.globals.reference_face_position) if 'reference' in facefusion.globals.face_recognition else None
+	result_frame = process_frame(source_face, reference_face, target_frame)
+	cv2.imwrite(output_path, result_frame)
+
+
+def process_video(source_path : str, temp_frame_paths : List[str]) -> None:
+	conditional_set_face_reference(temp_frame_paths)
+	frame_processors.process_video(source_path, temp_frame_paths, process_frames)
+
+
+def conditional_set_face_reference(temp_frame_paths : List[str]) -> None:
+	if 'reference' in facefusion.globals.face_recognition and not get_face_reference():
+		reference_frame = cv2.imread(temp_frame_paths[facefusion.globals.reference_frame_number])
+		reference_face = get_one_face(reference_frame, facefusion.globals.reference_face_position)
+		set_face_reference(reference_face)

facefusion/processors/frame/modules/frame_enhancer.py

@@ -0,0 +1,88 @@
+from typing import Any, List, Callable
+import cv2
+import threading
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from realesrgan import RealESRGANer
+
+import facefusion.processors.frame.core as frame_processors
+from facefusion.typing import Frame, Face
+from facefusion.utilities import conditional_download, resolve_relative_path
+
+FRAME_PROCESSOR = None
+THREAD_SEMAPHORE = threading.Semaphore()
+THREAD_LOCK = threading.Lock()
+NAME = 'FACEFUSION.FRAME_PROCESSOR.FRAME_ENHANCER'
+
+
+def get_frame_processor() -> Any:
+	global FRAME_PROCESSOR
+
+	with THREAD_LOCK:
+		if FRAME_PROCESSOR is None:
+			model_path = resolve_relative_path('../.assets/models/RealESRGAN_x4plus.pth')
+			FRAME_PROCESSOR = RealESRGANer(
+				model_path = model_path,
+				model = RRDBNet(
+					num_in_ch = 3,
+					num_out_ch = 3,
+					num_feat = 64,
+					num_block = 23,
+					num_grow_ch = 32,
+					scale = 4
+				),
+				device = frame_processors.get_device(),
+				tile = 512,
+				tile_pad = 32,
+				pre_pad = 0,
+				scale = 4
+			)
+	return FRAME_PROCESSOR
+
+
+def clear_frame_processor() -> None:
+	global FRAME_PROCESSOR
+
+	FRAME_PROCESSOR = None
+
+
+def pre_check() -> bool:
+	download_directory_path = resolve_relative_path('../.assets/models')
+	conditional_download(download_directory_path, ['https://huggingface.co/facefusion/models/resolve/main/RealESRGAN_x4plus.pth'])
+	return True
+
+
+def pre_process() -> bool:
+	return True
+
+
+def post_process() -> None:
+	clear_frame_processor()
+
+
+def enhance_frame(temp_frame : Frame) -> Frame:
+	with THREAD_SEMAPHORE:
+		temp_frame, _ = get_frame_processor().enhance(temp_frame, outscale = 1)
+	return temp_frame
+
+
+def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) -> Frame:
+	return enhance_frame(temp_frame)
+
+
+def process_frames(source_path : str, temp_frame_paths : List[str], update: Callable[[], None]) -> None:
+	for temp_frame_path in temp_frame_paths:
+		temp_frame = cv2.imread(temp_frame_path)
+		result_frame = process_frame(None, None, temp_frame)
+		cv2.imwrite(temp_frame_path, result_frame)
+		if update:
+			update()
+
+
+def process_image(source_path : str, target_path : str, output_path : str) -> None:
+	target_frame = cv2.imread(target_path)
+	result = process_frame(None, None, target_frame)
+	cv2.imwrite(output_path, result)
+
+
+def process_video(source_path : str, temp_frame_paths : List[str]) -> None:
+	frame_processors.process_video(None, temp_frame_paths, process_frames)