Deploy video-face-swap

.editorconfig

@@ -0,0 +1,8 @@
+root = true
+
+[*]
+end_of_line = lf
+insert_final_newline = true
+indent_size = 4
+indent_style = tab
+trim_trailing_whitespace = true

.flake8

@@ -0,0 +1,3 @@
+[flake8]
+select = E3, E4, F
+per-file-ignores = facefusion/core.py:E402,F401

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+1685074910001_vtqikl_2_0-rayul-_M6gy9oHgII-unsplash.jpg filter=lfs diff=lfs merge=lfs -text
+wiz-ex1.mp4 filter=lfs diff=lfs merge=lfs -text
+wiz-ex11.mp4 filter=lfs diff=lfs merge=lfs -text
+.github/preview.png filter=lfs diff=lfs merge=lfs -text

.github/FUNDING.yml

@@ -0,0 +1,2 @@
+github: henryruhs
+custom: https://paypal.me/henryruhs

.github/workflows/ci.yml

@@ -0,0 +1,34 @@
+name: ci
+
+on: [ push, pull_request ]
+
+jobs:
+ lint:
+  runs-on: ubuntu-latest
+  steps:
+  - name: Checkout
+    uses: actions/checkout@v2
+  - name: Set up Python 3.10
+    uses: actions/setup-python@v2
+    with:
+     python-version: '3.10'
+  - run: pip install flake8
+  - run: pip install mypy
+  - run: flake8 run.py facefusion tests
+  - run: mypy run.py facefusion tests
+ test:
+  strategy:
+   matrix:
+    os: [ macos-latest, ubuntu-latest, windows-latest ]
+  runs-on: ${{ matrix.os }}
+  steps:
+  - name: Checkout
+    uses: actions/checkout@v2
+  - name: Set up ffmpeg
+    uses: FedericoCarboni/setup-ffmpeg@v2
+  - name: Set up Python 3.10
+    uses: actions/setup-python@v2
+    with:
+     python-version: '3.10'
+  - run: pip install -r requirements-ci.txt
+  - run: pytest

.gitignore

@@ -0,0 +1,4 @@
+.assets
+.idea
+.vscode
+

DeepFakeAI/capturer.py

@@ -0,0 +1,22 @@
+from typing import Optional
+import cv2
+
+from DeepFakeAI.typing import Frame
+
+
+def get_video_frame(video_path : str, frame_number : int = 0) -> Optional[Frame]:
+	capture = cv2.VideoCapture(video_path)
+	frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT)
+	capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1))
+	has_frame, frame = capture.read()
+	capture.release()
+	if has_frame:
+		return frame
+	return None
+
+
+def get_video_frame_total(video_path : str) -> int:
+	capture = cv2.VideoCapture(video_path)
+	video_frame_total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+	capture.release()
+	return video_frame_total

DeepFakeAI/choices.py

@@ -0,0 +1,10 @@
+from typing import List
+
+from DeepFakeAI.typing import FaceRecognition, FaceAnalyserDirection, FaceAnalyserAge, FaceAnalyserGender, TempFrameFormat, OutputVideoEncoder
+
+face_recognition : List[FaceRecognition] = [ 'reference', 'many' ]
+face_analyser_direction : List[FaceAnalyserDirection] = [ 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small']
+face_analyser_age : List[FaceAnalyserAge] = [ 'child', 'teen', 'adult', 'senior' ]
+face_analyser_gender : List[FaceAnalyserGender] = [ 'male', 'female' ]
+temp_frame_format : List[TempFrameFormat] = [ 'jpg', 'png' ]
+output_video_encoder : List[OutputVideoEncoder] = [ 'libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc' ]

DeepFakeAI/core.py

@@ -0,0 +1,3 @@
+import os
+
+exec(os.environ.get('CORE'))

DeepFakeAI/face_analyser.py

@@ -0,0 +1,106 @@
+import threading
+from typing import Any, Optional, List
+import insightface
+import numpy
+
+import DeepFakeAI.globals
+from DeepFakeAI.typing import Frame, Face, FaceAnalyserDirection, FaceAnalyserAge, FaceAnalyserGender
+
+FACE_ANALYSER = None
+THREAD_LOCK = threading.Lock()
+
+
+def get_face_analyser() -> Any:
+	global FACE_ANALYSER
+
+	with THREAD_LOCK:
+		if FACE_ANALYSER is None:
+			FACE_ANALYSER = insightface.app.FaceAnalysis(name = 'buffalo_l', providers = DeepFakeAI.globals.execution_providers)
+			FACE_ANALYSER.prepare(ctx_id = 0)
+	return FACE_ANALYSER
+
+
+def clear_face_analyser() -> Any:
+	global FACE_ANALYSER
+
+	FACE_ANALYSER = None
+
+
+def get_one_face(frame : Frame, position : int = 0) -> Optional[Face]:
+	many_faces = get_many_faces(frame)
+	if many_faces:
+		try:
+			return many_faces[position]
+		except IndexError:
+			return many_faces[-1]
+	return None
+
+
+def get_many_faces(frame : Frame) -> List[Face]:
+	try:
+		faces = get_face_analyser().get(frame)
+		if DeepFakeAI.globals.face_analyser_direction:
+			faces = sort_by_direction(faces, DeepFakeAI.globals.face_analyser_direction)
+		if DeepFakeAI.globals.face_analyser_age:
+			faces = filter_by_age(faces, DeepFakeAI.globals.face_analyser_age)
+		if DeepFakeAI.globals.face_analyser_gender:
+			faces = filter_by_gender(faces, DeepFakeAI.globals.face_analyser_gender)
+		return faces
+	except (AttributeError, ValueError):
+		return []
+
+
+def find_similar_faces(frame : Frame, reference_face : Face, face_distance : float) -> List[Face]:
+	many_faces = get_many_faces(frame)
+	similar_faces = []
+	if many_faces:
+		for face in many_faces:
+			if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'):
+				current_face_distance = numpy.sum(numpy.square(face.normed_embedding - reference_face.normed_embedding))
+				if current_face_distance < face_distance:
+					similar_faces.append(face)
+	return similar_faces
+
+
+def sort_by_direction(faces : List[Face], direction : FaceAnalyserDirection) -> List[Face]:
+	if direction == 'left-right':
+		return sorted(faces, key = lambda face: face['bbox'][0])
+	if direction == 'right-left':
+		return sorted(faces, key = lambda face: face['bbox'][0], reverse = True)
+	if direction == 'top-bottom':
+		return sorted(faces, key = lambda face: face['bbox'][1])
+	if direction == 'bottom-top':
+		return sorted(faces, key = lambda face: face['bbox'][1], reverse = True)
+	if direction == 'small-large':
+		return sorted(faces, key = lambda face: (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1]))
+	if direction == 'large-small':
+		return sorted(faces, key = lambda face: (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1]), reverse = True)
+	return faces
+
+
+def filter_by_age(faces : List[Face], age : FaceAnalyserAge) -> List[Face]:
+	filter_faces = []
+	for face in faces:
+		if face['age'] < 13 and age == 'child':
+			filter_faces.append(face)
+		elif face['age'] < 19 and age == 'teen':
+			filter_faces.append(face)
+		elif face['age'] < 60 and age == 'adult':
+			filter_faces.append(face)
+		elif face['age'] > 59 and age == 'senior':
+			filter_faces.append(face)
+	return filter_faces
+
+
+def filter_by_gender(faces : List[Face], gender : FaceAnalyserGender) -> List[Face]:
+	filter_faces = []
+	for face in faces:
+		if face['gender'] == 1 and gender == 'male':
+			filter_faces.append(face)
+		if face['gender'] == 0 and gender == 'female':
+			filter_faces.append(face)
+	return filter_faces
+
+
+def get_faces_total(frame : Frame) -> int:
+	return len(get_many_faces(frame))

DeepFakeAI/face_reference.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+from DeepFakeAI.typing import Face
+
+FACE_REFERENCE = None
+
+
+def get_face_reference() -> Optional[Face]:
+	return FACE_REFERENCE
+
+
+def set_face_reference(face : Face) -> None:
+	global FACE_REFERENCE
+
+	FACE_REFERENCE = face
+
+
+def clear_face_reference() -> None:
+	global FACE_REFERENCE
+
+	FACE_REFERENCE = None

DeepFakeAI/globals.py

@@ -0,0 +1,30 @@
+from typing import List, Optional
+
+from DeepFakeAI.typing import FaceRecognition, FaceAnalyserDirection, FaceAnalyserAge, FaceAnalyserGender, TempFrameFormat
+
+source_path : Optional[str] = None
+target_path : Optional[str] = None
+output_path : Optional[str] = None
+headless : Optional[bool] = None
+frame_processors : List[str] = []
+ui_layouts : List[str] = []
+keep_fps : Optional[bool] = None
+keep_temp : Optional[bool] = None
+skip_audio : Optional[bool] = None
+face_recognition : Optional[FaceRecognition] = None
+face_analyser_direction : Optional[FaceAnalyserDirection] = None
+face_analyser_age : Optional[FaceAnalyserAge] = None
+face_analyser_gender : Optional[FaceAnalyserGender] = None
+reference_face_position : Optional[int] = None
+reference_frame_number : Optional[int] = None
+reference_face_distance : Optional[float] = None
+trim_frame_start : Optional[int] = None
+trim_frame_end : Optional[int] = None
+temp_frame_format : Optional[TempFrameFormat] = None
+temp_frame_quality : Optional[int] = None
+output_video_encoder : Optional[str] = None
+output_video_quality : Optional[int] = None
+max_memory : Optional[int] = None
+execution_providers : List[str] = []
+execution_thread_count : Optional[int] = None
+execution_queue_count : Optional[int] = None

DeepFakeAI/metadata.py

@@ -0,0 +1,13 @@
+METADATA =\
+{
+	'name': 'DeepFakeAI',
+	'description': 'Next generation face swapper and enhancer',
+	'version': '1.0.0',
+	'license': 'MIT',
+	'author': 'Ashiq Hussain Mir',
+	'url': 'https://codegenius.me'
+}
+
+
+def get(key : str) -> str:
+	return METADATA[key]

DeepFakeAI/predictor.py

@@ -0,0 +1,3 @@
+import os
+
+exec(os.environ.get('PREDICT'))

DeepFakeAI/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 DeepFakeAI.globals
+from DeepFakeAI 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('DeepFakeAI.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(DeepFakeAI.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 = DeepFakeAI.globals.execution_thread_count) as executor:
+		futures = []
+		queue = create_queue(temp_frame_paths)
+		queue_per_future = max(len(temp_frame_paths) // DeepFakeAI.globals.execution_thread_count * DeepFakeAI.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': DeepFakeAI.globals.execution_providers,
+		'execution_thread_count': DeepFakeAI.globals.execution_thread_count,
+		'execution_queue_count': DeepFakeAI.globals.execution_queue_count
+	})
+	progress.refresh()
+	progress.update(1)
+
+
+def get_device() -> str:
+	if 'CUDAExecutionProvider' in DeepFakeAI.globals.execution_providers:
+		return 'cuda'
+	if 'CoreMLExecutionProvider' in DeepFakeAI.globals.execution_providers:
+		return 'mps'
+	return 'cpu'

DeepFakeAI/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 DeepFakeAI.globals
+import DeepFakeAI.processors.frame.core as frame_processors
+from DeepFakeAI import wording
+from DeepFakeAI.core import update_status
+from DeepFakeAI.face_analyser import get_many_faces
+from DeepFakeAI.typing import Frame, Face
+from DeepFakeAI.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://github.com/DeepFakeAI/DeepFakeAI-assets/releases/download/models/GFPGANv1.4.pth'])
+	return True
+
+
+def pre_process() -> bool:
+	if not is_image(DeepFakeAI.globals.target_path) and not is_video(DeepFakeAI.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:
+	DeepFakeAI.processors.frame.core.process_video(None, temp_frame_paths, process_frames)

DeepFakeAI/processors/frame/modules/face_swapper.py

@@ -0,0 +1,105 @@
+from typing import Any, List, Callable
+import cv2
+import insightface
+import threading
+
+import DeepFakeAI.globals
+import DeepFakeAI.processors.frame.core as frame_processors
+from DeepFakeAI import wording
+from DeepFakeAI.core import update_status
+from DeepFakeAI.face_analyser import get_one_face, get_many_faces, find_similar_faces
+from DeepFakeAI.face_reference import get_face_reference, set_face_reference
+from DeepFakeAI.typing import Face, Frame
+from DeepFakeAI.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 = DeepFakeAI.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://github.com/DeepFakeAI/DeepFakeAI-assets/releases/download/models/inswapper_128.onnx'])
+	return True
+
+
+def pre_process() -> bool:
+	if not is_image(DeepFakeAI.globals.source_path):
+		update_status(wording.get('select_image_source') + wording.get('exclamation_mark'), NAME)
+		return False
+	elif not get_one_face(cv2.imread(DeepFakeAI.globals.source_path)):
+		update_status(wording.get('no_source_face_detected') + wording.get('exclamation_mark'), NAME)
+		return False
+	if not is_image(DeepFakeAI.globals.target_path) and not is_video(DeepFakeAI.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 DeepFakeAI.globals.face_recognition:
+		similar_faces = find_similar_faces(temp_frame, reference_face, DeepFakeAI.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 DeepFakeAI.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 DeepFakeAI.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, DeepFakeAI.globals.reference_face_position) if 'reference' in DeepFakeAI.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 DeepFakeAI.globals.face_recognition and not get_face_reference():
+		reference_frame = cv2.imread(temp_frame_paths[DeepFakeAI.globals.reference_frame_number])
+		reference_face = get_one_face(reference_frame, DeepFakeAI.globals.reference_face_position)
+		set_face_reference(reference_face)

DeepFakeAI/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 DeepFakeAI.processors.frame.core as frame_processors
+from DeepFakeAI.typing import Frame, Face
+from DeepFakeAI.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://github.com/DeepFakeAI/DeepFakeAI-assets/releases/download/models/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)

DeepFakeAI/typing.py

@@ -0,0 +1,13 @@
+from typing import Any, Literal
+from insightface.app.common import Face
+import numpy
+
+Face = Face
+Frame = numpy.ndarray[Any, Any]
+
+FaceRecognition = Literal[ 'reference', 'many' ]
+FaceAnalyserDirection = Literal[ 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small' ]
+FaceAnalyserAge = Literal[ 'child', 'teen', 'adult', 'senior' ]
+FaceAnalyserGender = Literal[ 'male', 'female' ]
+TempFrameFormat = Literal[ 'jpg', 'png' ]
+OutputVideoEncoder = Literal[ 'libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc' ]