app.py

import os

# Speed up Hub downloads a lot (parallel chunked transfer instead of a
# single stream). Must be set before huggingface_hub is imported/used.
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

import cv2
import gradio as gr
import numpy as np
import insightface
from insightface.app import FaceAnalysis
from huggingface_hub import hf_hub_download
import subprocess
import tempfile

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
# inswapper_128.onnx isn't bundled in this repo (it's a large binary model
# file). We try to fetch it automatically from a couple of known public
# mirrors on the Hugging Face Hub. If that ever stops working (mirrors can
# disappear), just upload `inswapper_128.onnx` yourself to the Space's root
# folder and the app will pick it up automatically. See README.md.

INSWAPPER_MIRRORS = [
    ("ezioruan/inswapper_128.onnx", "inswapper_128.onnx"),
    ("Devrim/inswapper_128", "inswapper_128.onnx"),
    ("countfloyd/deepfake", "inswapper_128.onnx"),
]

LOCAL_MODEL_NAME = "inswapper_128.onnx"


def get_swapper_model_path() -> str:
    if os.path.exists(LOCAL_MODEL_NAME):
        return LOCAL_MODEL_NAME

    last_err = None
    for repo_id, filename in INSWAPPER_MIRRORS:
        try:
            print(f"[model-download] trying {repo_id} ...")
            path = hf_hub_download(repo_id=repo_id, filename=filename)
            print(f"[model-download] success from {repo_id} -> {path}")
            return path
        except Exception as e:  # noqa: BLE001
            print(f"[model-download] failed from {repo_id}: {e}")
            last_err = e

    raise RuntimeError(
        "Could not auto-download inswapper_128.onnx from any mirror. "
        "Please upload the file manually to the Space's root directory "
        f"(same folder as app.py). Last error: {last_err}"
    )


_face_analyser = None
_swapper = None


def get_models():
    """Lazily load and cache the face-analysis and face-swap models."""
    global _face_analyser, _swapper

    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]

    if _face_analyser is None:
        _face_analyser = FaceAnalysis(name="buffalo_l", providers=providers)
        _face_analyser.prepare(ctx_id=0, det_size=(640, 640))

    if _swapper is None:
        model_path = get_swapper_model_path()
        _swapper = insightface.model_zoo.get_model(model_path, providers=providers)

    return _face_analyser, _swapper


# ---------------------------------------------------------------------------
# Core face-swap logic
# ---------------------------------------------------------------------------

def get_source_face(face_analyser, image_bgr):
    faces = face_analyser.get(image_bgr)
    if not faces:
        return None
    # If several faces are in the source photo, use the largest one.
    faces = sorted(
        faces,
        key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]),
        reverse=True,
    )
    return faces[0]


def swap_face_in_frame(face_analyser, swapper, source_face, frame):
    faces = face_analyser.get(frame)
    if not faces:
        return frame
    result = frame
    for face in faces:
        result = swapper.get(result, face, source_face, paste_back=True)
    return result


def merge_audio(original_video_path, silent_video_path, output_path) -> bool:
    """Copy the audio track from the original video onto the swapped video."""
    try:
        cmd = [
            "ffmpeg", "-y",
            "-i", silent_video_path,
            "-i", original_video_path,
            "-c:v", "copy",
            "-map", "0:v:0",
            "-map", "1:a:0?",
            "-c:a", "aac",
            "-shortest",
            output_path,
        ]
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0 or not os.path.exists(output_path):
            print("[ffmpeg] audio merge failed:", result.stderr)
            return False
        return True
    except Exception as e:  # noqa: BLE001
        print("[ffmpeg] audio merge error:", e)
        return False


def process_video(source_image, target_video_path, progress=gr.Progress()):
    if source_image is None or target_video_path is None:
        raise gr.Error("Please upload both a source face image and a target video.")

    progress(0, desc="Loading models (first run downloads them, ~1-2 min)...")
    face_analyser, swapper = get_models()

    source_bgr = cv2.cvtColor(source_image, cv2.COLOR_RGB2BGR)
    source_face = get_source_face(face_analyser, source_bgr)
    if source_face is None:
        raise gr.Error(
            "No face was detected in the uploaded image. "
            "Try a clearer, front-facing, well-lit photo."
        )

    cap = cv2.VideoCapture(target_video_path)
    if not cap.isOpened():
        raise gr.Error("Could not open the uploaded video.")

    fps = cap.get(cv2.CAP_PROP_FPS) or 25
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0

    tmp_dir = tempfile.mkdtemp()
    silent_path = os.path.join(tmp_dir, "swapped_silent.mp4")
    final_path = os.path.join(tmp_dir, "swapped_final.mp4")

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(silent_path, fourcc, fps, (width, height))

    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        swapped = swap_face_in_frame(face_analyser, swapper, source_face, frame)
        writer.write(swapped)
        frame_idx += 1
        if total_frames:
            progress(
                min(frame_idx / total_frames, 0.95),
                desc=f"Swapping frame {frame_idx}/{total_frames}",
            )
        else:
            progress(0.5, desc=f"Swapping frame {frame_idx}")

    cap.release()
    writer.release()

    progress(0.97, desc="Merging original audio track...")
    ok = merge_audio(target_video_path, silent_path, final_path)
    result_path = final_path if ok else silent_path

    progress(1.0, desc="Done!")
    return result_path


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="Face Swap: Image → Video") as demo:
    gr.Markdown(
        """
        # 🎭 Face Swap — Image onto Video
        Upload a clear photo of a face and a target video. The face from your
        photo will be swapped onto every face detected in the video.

        **Please only use this on content you have the rights and consent to edit.**
        This tool is for creative and educational use — don't use it to
        impersonate real people without their permission.
        """
    )
    with gr.Row():
        with gr.Column():
            source_image = gr.Image(label="Source face image", type="numpy")
            target_video = gr.Video(label="Target video", sources=["upload"])
            run_btn = gr.Button("Swap Face", variant="primary")
        with gr.Column():
            output_video = gr.Video(label="Result")

    run_btn.click(
        fn=process_video,
        inputs=[source_image, target_video],
        outputs=[output_video],
    )

    gr.Markdown(
        """
        ---
        ### Notes
        - First run downloads the face-analysis + face-swap models (~300 MB
          total), which can take a minute or two.
        - This app needs a **GPU Space** (T4 or better) for reasonable speed.
          On CPU it will still work but be much slower.
        - Every frame of the video is processed, so long videos take a while.
        - Output keeps the original video's audio track when possible.
        """
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("GRADIO_SERVER_PORT", 7860)),
    )