import glob import mimetypes import os import platform import shutil import ssl import subprocess import urllib from pathlib import Path from typing import List, Optional, Any, Callable from tqdm import tqdm import time import onnxruntime # import tensorflow import sys # single thread doubles cuda performance - needs to be set before torch import if any(arg.startswith('--execution-provider') for arg in sys.argv): os.environ['OMP_NUM_THREADS'] = '1' # reduce tensorflow log level os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' import warnings import signal import importlib import psutil from concurrent.futures import ThreadPoolExecutor, as_completed from queue import Queue import cv2 import threading import insightface import numpy from insightface.app.common import Face # from gfpgan.utils import GFPGANer import streamlit as st Face = Face Frame = numpy.ndarray[Any, Any] FACE_ENHANCER = None FACE_ANALYSER = None FACE_SWAPPER = None THREAD_LOCK = threading.Lock() THREAD_SEMAPHORE = threading.Semaphore() output_video_encoder = 'libx264' working_dir = os.path.dirname(os.path.abspath(__file__)) source_path = os.path.join(working_dir,"source.png") target_path = os.path.join(working_dir,"target.mp4") source_name, _ = os.path.splitext(os.path.basename(source_path)) target_name, target_extension = os.path.splitext(os.path.basename(target_path)) output_path = os.path.join(working_dir, source_name + '-' + target_name + target_extension) face_swapper_path = working_dir face_enhancer_path = working_dir def pre_check() -> bool: if sys.version_info < (3, 9): print('Python version is not supported - please upgrade to 3.9 or higher.') return False if not shutil.which('ffmpeg'): print('ffmpeg is not installed.') return False return True def has_image_extension(image_path: str) -> bool: return image_path.lower().endswith(('png', 'jpg', 'jpeg', 'webp')) def is_image(image_path: str) -> bool: if image_path and os.path.isfile(image_path): mimetype, _ = mimetypes.guess_type(image_path) return bool(mimetype and mimetype.startswith('image/')) return False def is_video(video_path: str) -> bool: if video_path and os.path.isfile(video_path): mimetype, _ = mimetypes.guess_type(video_path) return bool(mimetype and mimetype.startswith('video/')) return False 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=execution_providers) FACE_ANALYSER.prepare(ctx_id=0) return FACE_ANALYSER def get_many_faces(frame: Frame) -> Optional[List[Face]]: try: return get_face_analyser().get(frame) except ValueError: return 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 pre_start() -> bool: if not is_image(source_path): st.write('Select an image for source path.') return False elif not get_one_face(cv2.imread(source_path)): st.write('No face in source path detected.') return False if not is_video(target_path): st.write('Select a video for target path.') return False return True def detect_fps(target_path: str) -> float: command = ['ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries', 'stream=r_frame_rate', '-of', 'default=noprint_wrappers=1:nokey=1', target_path] output = subprocess.check_output(command).decode().strip().split('/') try: numerator, denominator = map(int, output) return numerator / denominator except Exception: pass return 30 def run_ffmpeg(args: List[str]) -> bool: commands = ['ffmpeg', '-hide_banner', '-loglevel', 'error'] commands.extend(args) try: subprocess.check_output(commands, stderr=subprocess.STDOUT) return True except Exception: pass return False def extract_frames(target_path: str, temp_directory_path: str, fps: float = 30) -> bool: temp_frame_quality = 0 return run_ffmpeg(['-hwaccel', 'auto', '-i', target_path, '-q:v', str(temp_frame_quality), '-pix_fmt', 'rgb24', '-vf', 'fps=' + str(fps), os.path.join(temp_directory_path, '%04d.' + 'png')]) def suggest_execution_threads() -> int: if 'CUDAExecutionProvider' in onnxruntime.get_available_providers(): return 8 return 1 execution_threads = suggest_execution_threads() 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 update_progress(progress, progress_bar, progress_text, total, queue_per_future): process = psutil.Process(os.getpid()) memory_usage = process.memory_info().rss / 1024 / 1024 / 1024 progress += queue_per_future if progress < total: progress_bar.progress(progress / total) progress_text.text(f"Progress: {progress}/{total} - Memory usage: {memory_usage:.2f} GB") else: progress_text.empty() # Remove text value progress_bar.empty() return progress def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], update) -> None: with ThreadPoolExecutor(max_workers=execution_threads) as executor: futures = [] queue = create_queue(temp_frame_paths) queue_per_future = max(len(temp_frame_paths) // execution_threads, 1) progress = 0 total = len(temp_frame_paths) progress_bar = st.progress(0) progress_text = st.empty() while not queue.empty(): future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future)) futures.append(future) for future in as_completed(futures): progress = update(progress, progress_bar, progress_text, total, queue_per_future) def process_video(source_path: str, frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None]) -> None: multi_process_frame(source_path, frame_paths, process_frames, update_progress) def encode_execution_providers(execution_providers: List[str]) -> List[str]: return [execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers] def decode_execution_providers(execution_providers: List[str]) -> List[str]: return [provider for provider, encoded_execution_provider in zip(onnxruntime.get_available_providers(), encode_execution_providers(onnxruntime.get_available_providers())) if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers)] def suggest_execution_providers() -> List[str]: return encode_execution_providers(onnxruntime.get_available_providers()) execution_providers = decode_execution_providers(suggest_execution_providers()) def get_face_swapper() -> Any: global FACE_SWAPPER with THREAD_LOCK: if FACE_SWAPPER is None: model_path = face_swapper_path FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=execution_providers) return FACE_SWAPPER def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame: return get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True) similar_face_distance = 0.85 def find_similar_face(frame: Frame, reference_face: Face) -> Optional[Face]: many_faces = get_many_faces(frame) if many_faces: for face in many_faces: if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'): distance = numpy.sum(numpy.square(face.normed_embedding - reference_face.normed_embedding)) if distance < similar_face_distance: return face return None def process_frame(source_face: Face, reference_face: Face, temp_frame: Frame) -> Frame: target_face = find_similar_face(temp_frame, reference_face) if target_face: temp_frame = swap_face(source_face, target_face, temp_frame) return temp_frame def process_frames(source_path: str, temp_frame_paths: List[str]) -> None: source_face = get_one_face(cv2.imread(source_path)) reference_frame = cv2.imread(temp_frame_paths[0]) reference_face = get_one_face(reference_frame) for temp_frame_path in temp_frame_paths: temp_frame = cv2.imread(temp_frame_path) result = process_frame(source_face, reference_face, temp_frame) cv2.imwrite(temp_frame_path, result) # def get_face_enhancer() -> Any: # global FACE_ENHANCER # with THREAD_LOCK: # if FACE_ENHANCER is None: # model_path = '/content/drive/MyDrive/FaceSwap3/GFPGANv1.4.pth' # # todo: set models path -> https://github.com/TencentARC/GFPGAN/issues/399 # FACE_ENHANCER = GFPGANer(model_path=model_path, upscale=1, device=get_device()) # return FACE_ENHANCER def get_device() -> str: if 'CUDAExecutionProvider' in execution_providers: return 'cuda' if 'CoreMLExecutionProvider' in execution_providers: return 'mps' return 'cpu' 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) temp_face = temp_frame[start_y:end_y, start_x:end_x] if temp_face.size: with THREAD_SEMAPHORE: _, _, temp_face = get_face_enhancer().enhance(temp_face, paste_back=True) temp_frame[start_y:end_y, start_x:end_x] = temp_face return temp_frame def enhance_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 enhance_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 = enhance_frame(None, None, temp_frame) cv2.imwrite(temp_frame_path, result) if update: update() def create_video(target_path: str, temp_directory_path: str, fps: float = 30) -> bool: temp_output_path = os.path.join(temp_directory_path, 'temp.mp4') output_video_quality = (35 + 1) * 51 // 100 commands = ['-hwaccel', 'auto', '-r', str(fps), '-i', os.path.join(temp_directory_path, '%04d.' + 'png'), '-c:v', output_video_encoder] if output_video_encoder in ['libx264', 'libx265', 'libvpx']: commands.extend(['-crf', str(output_video_quality)]) if output_video_encoder in ['h264_nvenc', 'hevc_nvenc']: commands.extend(['-cq', str(output_video_quality)]) commands.extend(['-pix_fmt', 'yuv420p', '-vf', 'colorspace=bt709:iall=bt601-6-625:fast=1', '-y', temp_output_path]) return run_ffmpeg(commands) def move_temp(temp_output_path: str, output_path: str) -> None: if os.path.isfile(temp_output_path): if os.path.isfile(output_path): os.remove(output_path) shutil.move(temp_output_path, output_path) def restore_audio(target_path: str, temp_directory_path: str, output_path: str) -> None: temp_output_path = os.path.join(temp_directory_path, 'temp.mp4') done = run_ffmpeg(['-i', temp_output_path, '-i', target_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path]) if not done: move_temp(temp_output_path, output_path) def start(quality: bool) -> Optional[str]: if not pre_start(): return with st.spinner("Preparing..."): target_directory_path = os.path.dirname(target_path) temp_directory_path = os.path.join(target_directory_path, 'temp') Path(temp_directory_path).mkdir(parents=True, exist_ok=True) fps = detect_fps(target_path) with st.spinner(f'Extracting frames with {fps} FPS...'): extract_frames(target_path, temp_directory_path, fps) temp_frame_paths = glob.glob((os.path.join(glob.escape(temp_directory_path), '*.' + 'png'))) if temp_frame_paths: with st.spinner('Swapping Progressing...'): process_video(source_path, temp_frame_paths, process_frames) if quality: with st.spinner('Enhancing Progressing...'): process_video(None, temp_frame_paths, enhance_frames) else: st.write('Frames not found...') return with st.spinner(f'Creating video with {fps} FPS...'): create_video(target_path, temp_directory_path, fps) with st.spinner('Restoring audio...'): restore_audio(target_path, temp_directory_path, output_path) if is_video(output_path): st.video(output_path) else: st.write('Processing to video failed!') def conditional_download(download_directory_path: str, urls: List[str]) -> Optional[str]: if not os.path.exists(download_directory_path): os.makedirs(download_directory_path) for url in urls: download_file_path = os.path.join(download_directory_path, os.path.basename(url)) if not os.path.exists(download_file_path): request = urllib.request.urlopen(url) # type: ignore[attr-defined] total = int(request.headers.get('Content-Length', 0)) with tqdm(total=total, desc='Downloading', unit='B', unit_scale=True, unit_divisor=1024) as progress: urllib.request.urlretrieve(url, download_file_path, reporthook=lambda count, block_size, total_size: progress.update(block_size)) # type: ignore[attr-defined] def run(quality: bool) -> Optional[str]: if not pre_check(): return with st.spinner("Preprocessing..."): conditional_download(face_swapper_path, ['https://huggingface.co/CountFloyd/deepfake/resolve/main/inswapper_128.onnx']) conditional_download(face_enhancer_path, ['https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth']) start(quality)