import os import shutil import subprocess import tempfile from pathlib import Path import cv2 import gradio as gr import onnxruntime as ort from huggingface_hub import hf_hub_download from insightface.app import FaceAnalysis from insightface.model_zoo import get_model APP_DIR = Path(__file__).resolve().parent MODELS_DIR = APP_DIR / "models" TMP_DIR = APP_DIR / "tmp" MODELS_DIR.mkdir(exist_ok=True) TMP_DIR.mkdir(exist_ok=True) MODEL_REPO_ID = os.environ.get("MODEL_REPO_ID", "ezioruan/inswapper_128.onnx") MODEL_FILENAME = os.environ.get("MODEL_FILENAME", "inswapper_128.onnx") MODEL_REVISION = os.environ.get("MODEL_REVISION") HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN") FACE_MODEL_NAME = os.environ.get("FACE_MODEL_NAME", "buffalo_l") DETECTION_SIZE = int(os.environ.get("DETECTION_SIZE", "640")) MAX_FRAMES = int(os.environ.get("MAX_FRAMES", "0")) FORCE_CPU = os.environ.get("FORCE_CPU") == "1" face_analyzer = None face_swapper = None model_path = None def ensure_ffmpeg(): if shutil.which("ffmpeg") is None: raise RuntimeError("ffmpeg is required but not available in PATH.") def get_execution_providers(): providers = [] if not FORCE_CPU: providers.append("CUDAExecutionProvider") providers.append("CPUExecutionProvider") return providers def provider_status(): available = ort.get_available_providers() return available, ("CUDAExecutionProvider" in available) def ensure_swap_model() -> str: global model_path if model_path and os.path.exists(model_path): return model_path local_target = MODELS_DIR / MODEL_FILENAME if local_target.exists(): model_path = str(local_target) return model_path downloaded = hf_hub_download( repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME, revision=MODEL_REVISION, token=HF_TOKEN, local_dir=str(MODELS_DIR), ) model_path = downloaded return model_path def load_models(): global face_analyzer, face_swapper if face_analyzer is not None and face_swapper is not None: return face_analyzer, face_swapper swap_model_path = ensure_swap_model() face_analyzer = FaceAnalysis(name=FACE_MODEL_NAME, providers=get_execution_providers()) face_analyzer.prepare(ctx_id=0, det_size=(DETECTION_SIZE, DETECTION_SIZE)) face_swapper = get_model(swap_model_path, providers=get_execution_providers()) return face_analyzer, face_swapper def pick_largest_face(faces): if not faces: return None return max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1])) def extract_audio(input_video: str, audio_path: str): subprocess.run(["ffmpeg", "-y", "-i", input_video, "-vn", "-acodec", "copy", audio_path], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) def mux_audio(video_path: str, audio_path: str, output_path: str): if not os.path.exists(audio_path): shutil.move(video_path, output_path) return subprocess.run(["ffmpeg", "-y", "-i", video_path, "-i", audio_path, "-c:v", "copy", "-c:a", "aac", "-shortest", output_path], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) os.remove(video_path) def swap_frame(frame_bgr, source_face, analyzer, swapper): target_faces = analyzer.get(frame_bgr) if not target_faces: return frame_bgr result = frame_bgr.copy() for target_face in target_faces: result = swapper.get(result, target_face, source_face, paste_back=True) return result def process_video(face_image_path, video_path, progress=gr.Progress(track_tqdm=False)): available_providers, cuda_available = provider_status() if FORCE_CPU: raise RuntimeError(f"FORCE_CPU=1 is set. Available providers: {available_providers}") if not cuda_available: raise RuntimeError(f"CUDAExecutionProvider is not available. Available providers: {available_providers}") ensure_ffmpeg() analyzer, swapper = load_models() src = cv2.imread(face_image_path) if src is None: raise ValueError("Could not read the source face image.") source_face = pick_largest_face(analyzer.get(src)) if source_face is None: raise ValueError("No source face detected in the uploaded image.") work_dir = Path(tempfile.mkdtemp(dir=TMP_DIR)) silent_video = str(work_dir / "silent.mp4") audio_file = str(work_dir / "audio.aac") final_video = str(work_dir / "faceswapped.mp4") cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError("Could not open the uploaded video.") fps = cap.get(cv2.CAP_PROP_FPS) or 24.0 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1 limit = min(frame_count, MAX_FRAMES) if MAX_FRAMES > 0 else frame_count writer = cv2.VideoWriter(silent_video, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)) extract_audio(video_path, audio_file) idx = 0 while idx < limit: ok, frame = cap.read() if not ok: break writer.write(swap_frame(frame, source_face, analyzer, swapper)) idx += 1 progress(idx / max(limit, 1), desc=f"Processed {idx}/{limit} frames") cap.release() writer.release() mux_audio(silent_video, audio_file, final_video) return final_video, f"Done. Processed {idx} frame(s). Providers: {available_providers}" def startup_message(): available, cuda_ok = provider_status() return f"ONNX Runtime providers: {available} | CUDA available: {cuda_ok}" with gr.Blocks(title="Video FaceSwap GPU") as demo: gr.Markdown("""# Video FaceSwap GPU Upload a source face image and a target video. The app downloads the swap model from the Hugging Face Hub at runtime, then swaps the main source face onto detected faces in the video.""") status_banner = gr.Textbox(label="Startup status", value=startup_message(), interactive=False) with gr.Row(): face_image = gr.Image(type="filepath", label="Source face image") target_video = gr.Video(label="Target video") run_btn = gr.Button("Run face swap", variant="primary") output_video = gr.Video(label="Output video") status = gr.Textbox(label="Status") run_btn.click(fn=process_video, inputs=[face_image, target_video], outputs=[output_video, status], api_name="faceswap_video") if __name__ == "__main__": demo.queue(max_size=8).launch(server_name="0.0.0.0", server_port=7860, share=False)