""" 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()