app.py

"""
Face Swap Space
----------------
1. Upload a reference face.
2. Upload a target image and draw a circle over the face you want replaced.
3. InsightFace (inswapper_128) transplants the reference identity into that face.
4. Z-Image-Turbo (ZImageInpaintPipeline) does a low-strength inpaint pass over the
   same masked region to blend lighting/skin tone/edges so the swap looks native
   instead of pasted.

Z-Image-Turbo does NOT do identity transfer by itself -- it only knows how to
follow a text prompt + mask. The actual "put this person's face here" step is
done by the classic face-swap model; Z-Image is only the polish pass.
"""

import cv2
import numpy as np
import torch
import gradio as gr
from PIL import Image

import insightface
from insightface.app import FaceAnalysis
from huggingface_hub import hf_hub_download

from diffusers import ZImageInpaintPipeline

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# inswapper_128.onnx was pulled from InsightFace's own release assets, but it's
# mirrored on the HF Hub, so we just pull it from there on first run instead of
# making the user go find it manually.
SWAPPER_REPO = "Aitrepreneur/insightface"
SWAPPER_FILE = "inswapper_128.onnx"

# ---------------------------------------------------------------------------
# Lazy-loaded models (so the Space boots fast and only pays GPU cost on first use)
# ---------------------------------------------------------------------------
_face_app = None
_swapper = None
_zpipe = None


def get_face_app():
    global _face_app
    if _face_app is None:
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if DEVICE == "cuda" else ["CPUExecutionProvider"]
        _face_app = FaceAnalysis(name="buffalo_l", providers=providers)
        _face_app.prepare(ctx_id=0 if DEVICE == "cuda" else -1, det_size=(640, 640))
    return _face_app


def get_swapper():
    global _swapper
    if _swapper is None:
        local_path = hf_hub_download(repo_id=SWAPPER_REPO, filename=SWAPPER_FILE)
        _swapper = insightface.model_zoo.get_model(local_path)
    return _swapper


def get_zpipe():
    global _zpipe
    if _zpipe is None:
        _zpipe = ZImageInpaintPipeline.from_pretrained(
            "Tongyi-MAI/Z-Image-Turbo",
            torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
        )
        _zpipe.to(DEVICE)
    return _zpipe


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def extract_mask(editor_value: dict, size_wh) -> np.ndarray:
    """Turn the gr.ImageEditor brush layer(s) into a single 0/255 mask, sized to `size_wh` (w, h)."""
    w, h = size_wh
    mask = np.zeros((h, w), dtype=np.uint8)
    for layer in editor_value.get("layers", []):
        layer = np.array(layer)
        if layer.ndim == 3 and layer.shape[-1] == 4:
            alpha = layer[:, :, 3]
            if alpha.shape[:2] != (h, w):
                alpha = cv2.resize(alpha, (w, h), interpolation=cv2.INTER_NEAREST)
            mask = np.maximum(mask, (alpha > 10).astype(np.uint8) * 255)
    return mask


def dilate_mask(mask: np.ndarray, px: int = 15) -> np.ndarray:
    if px <= 0:
        return mask
    kernel = np.ones((px, px), np.uint8)
    return cv2.dilate(mask, kernel, iterations=1)


def pick_target_face(faces, mask: np.ndarray):
    """Pick whichever detected face sits closest to the center of the drawn mask."""
    ys, xs = np.where(mask > 0)
    if len(xs) == 0:
        raise gr.Error("Draw a circle over the face you want to replace before running.")
    cx, cy = xs.mean(), ys.mean()

    def dist(f):
        fx, fy = (f.bbox[0] + f.bbox[2]) / 2, (f.bbox[1] + f.bbox[3]) / 2
        return (fx - cx) ** 2 + (fy - cy) ** 2

    return min(faces, key=dist)


def face_swap(reference_img: Image.Image, target_img: Image.Image, mask: np.ndarray) -> Image.Image:
    fa = get_face_app()
    sw = get_swapper()

    ref_bgr = np.array(reference_img.convert("RGB"))[:, :, ::-1]
    tgt_bgr = np.array(target_img.convert("RGB"))[:, :, ::-1].copy()

    ref_faces = fa.get(ref_bgr)
    if not ref_faces:
        raise gr.Error("No face detected in the reference image.")
    ref_face = max(ref_faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))

    tgt_faces = fa.get(tgt_bgr)
    if not tgt_faces:
        raise gr.Error("No face detected in the target image.")
    tgt_face = pick_target_face(tgt_faces, mask)

    swapped_bgr = sw.get(tgt_bgr, tgt_face, ref_face, paste_back=True)
    return Image.fromarray(swapped_bgr[:, :, ::-1])


def blend_with_zimage(swapped_img: Image.Image, mask: np.ndarray, prompt: str, steps: int, strength: float) -> Image.Image:
    pipe = get_zpipe()
    w, h = swapped_img.size

    # Z-Image works best at 1024-ish, multiple-of-64 sizes; work on a resized copy
    # and paste the result back at native resolution.
    long_side = 1024
    scale = long_side / max(w, h)
    rw, rh = max(64, int(round(w * scale / 64)) * 64), max(64, int(round(h * scale / 64)) * 64)

    init = swapped_img.resize((rw, rh))
    mask_small = cv2.resize(mask, (rw, rh), interpolation=cv2.INTER_NEAREST)
    mask_img = Image.fromarray(mask_small).convert("L")

    out = pipe(
        prompt,
        image=init,
        mask_image=mask_img,
        strength=strength,
        num_inference_steps=steps,
        guidance_scale=0.0,
        generator=torch.Generator(DEVICE).manual_seed(0),
    ).images[0]

    return out.resize((w, h))


# ---------------------------------------------------------------------------
# Main callback
# ---------------------------------------------------------------------------
def run(reference_img, editor_value, prompt, do_blend, strength, steps, mask_grow):
    if reference_img is None:
        raise gr.Error("Upload a reference face image.")
    if editor_value is None or editor_value.get("background") is None:
        raise gr.Error("Upload a target image.")

    bg = editor_value["background"]
    target_img = Image.fromarray(bg) if isinstance(bg, np.ndarray) else bg
    target_img = target_img.convert("RGB")

    mask = extract_mask(editor_value, target_img.size)
    mask = dilate_mask(mask, mask_grow)

    swapped = face_swap(reference_img, target_img, mask)

    if do_blend:
        result = blend_with_zimage(swapped, mask, prompt, int(steps), float(strength))
    else:
        result = swapped

    return result


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Face Swap + Z-Image Turbo blend") as demo:
    gr.Markdown(
        """
        # Face Swap
        1. Upload the **reference face** (the identity you want to put in).
        2. Upload the **target image** and paint a circle over the face to replace.
        3. Hit **Swap face**.

        The actual identity transplant is done by InsightFace's `inswapper_128`.
        Z-Image-Turbo is optionally used afterward for a light blend pass so the
        edges/lighting match the rest of the photo.
        """
    )
    with gr.Row():
        ref_img = gr.Image(label="Reference face", type="pil")
        editor = gr.ImageEditor(
            label="Target image — draw a circle over the face to replace",
            type="numpy",
            brush=gr.Brush(colors=["#ffffff"], color_mode="fixed", default_size=40),
        )

    with gr.Accordion("Blend settings", open=False):
        do_blend = gr.Checkbox(label="Refine/blend with Z-Image Turbo", value=True)
        prompt = gr.Textbox(
            label="Blend prompt",
            value="photorealistic face, natural skin texture, consistent lighting, seamless, high detail",
        )
        strength = gr.Slider(0.05, 0.6, value=0.25, step=0.05, label="Blend strength (higher = more change, more identity drift)")
        steps = gr.Slider(4, 12, value=8, step=1, label="Steps")
        mask_grow = gr.Slider(0, 40, value=15, step=1, label="Mask grow (px)")

    btn = gr.Button("Swap face", variant="primary")
    output = gr.Image(label="Result")

    btn.click(
        run,
        inputs=[ref_img, editor, prompt, do_blend, strength, steps, mask_grow],
        outputs=output,
    )

if __name__ == "__main__":
    demo.launch()