app.py

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)