Upgrade face swap quality with HyperSwap and GFPGAN

app.py

@@ -11,8 +11,8 @@ from pathlib import Path
 
 import cv2
 import gradio as gr
-import insightface
 import numpy as np
+import onnxruntime as ort
 import spaces
 from huggingface_hub import hf_hub_download
 from insightface.app import FaceAnalysis
@@ -25,15 +25,43 @@ MODEL_ROOT = Path(tempfile.gettempdir()) / "insightface"
 OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
 MODEL_ROOT.mkdir(parents=True, exist_ok=True)
 
-# The smaller FP16 model cuts startup/download time while preserving the usual
-# InSwapper 128 output quality.
+# HyperSwap works at 256x256 instead of the previous 128x128 InSwapper crop.
+# Both models are downloaded from FaceFusion's official Hugging Face repos.
 SWAPPER_MODEL_PATH = hf_hub_download(
-    repo_id="ninjawick/webui-faceswap-unlocked",
-    filename="inswapper_128_fp16.onnx",
+    repo_id="facefusion/models-3.3.0",
+    filename="hyperswap_1a_256.onnx",
+)
+ENHANCER_MODEL_PATH = hf_hub_download(
+    repo_id="facefusion/models-3.0.0",
+    filename="gfpgan_1.4.onnx",
 )
 
 _MODEL_LOCK = threading.Lock()
-_MODEL_CACHE: tuple[FaceAnalysis, object] | None = None
+_MODEL_CACHE: tuple[FaceAnalysis, ort.InferenceSession, ort.InferenceSession] | None = None
+
+QUALITY_BALANCED = "Balanced — HyperSwap 256"
+QUALITY_BEST = "Best quality — HyperSwap 256 + GFPGAN 512"
+
+ARCFACE_128_TEMPLATE = np.array(
+    [
+        [0.36167656, 0.40387734],
+        [0.63696719, 0.40235469],
+        [0.50019687, 0.56044219],
+        [0.38710391, 0.72160547],
+        [0.61507734, 0.72034453],
+    ],
+    dtype=np.float32,
+)
+FFHQ_512_TEMPLATE = np.array(
+    [
+        [0.37691676, 0.46864664],
+        [0.62285697, 0.46912813],
+        [0.50123859, 0.61331904],
+        [0.39308822, 0.72541100],
+        [0.61150205, 0.72490465],
+    ],
+    dtype=np.float32,
+)
 
 
 def _cleanup_old_outputs(max_age_seconds: int = 3600) -> None:
@@ -74,6 +102,7 @@ def _estimate_gpu_seconds(
     target_face_index: int,
     clip_length: int,
     output_resolution: str,
+    quality_mode: str,
     keep_audio: bool,
     add_watermark: bool,
     consent: bool,
@@ -86,10 +115,11 @@ def _estimate_gpu_seconds(
             seconds = min(_video_metadata(target_video)[3] or seconds, seconds)
         except Exception:
             pass
-    return int(max(90, min(300, 45 + seconds * 5.5)))
+    seconds_per_second = 8.0 if quality_mode == QUALITY_BEST else 5.5
+    return int(max(120, min(300, 60 + seconds * seconds_per_second)))
 
 
-def _load_models() -> tuple[FaceAnalysis, object]:
+def _load_models() -> tuple[FaceAnalysis, ort.InferenceSession, ort.InferenceSession]:
     """Create CUDA ONNX sessions only after ZeroGPU has allocated a GPU."""
     global _MODEL_CACHE
     with _MODEL_LOCK:
@@ -104,14 +134,131 @@ def _load_models() -> tuple[FaceAnalysis, object]:
             providers=providers,
         )
         analyser.prepare(ctx_id=0, det_size=(640, 640))
-        swapper = insightface.model_zoo.get_model(
-            SWAPPER_MODEL_PATH,
-            providers=providers,
+        session_options = ort.SessionOptions()
+        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        swapper = ort.InferenceSession(
+            SWAPPER_MODEL_PATH, sess_options=session_options, providers=providers
         )
-        _MODEL_CACHE = (analyser, swapper)
+        enhancer = ort.InferenceSession(
+            ENHANCER_MODEL_PATH, sess_options=session_options, providers=providers
+        )
+        _MODEL_CACHE = (analyser, swapper, enhancer)
         return _MODEL_CACHE
 
 
+def _input_dtype(session: ort.InferenceSession, input_name: str) -> np.dtype:
+    input_type = next(item.type for item in session.get_inputs() if item.name == input_name)
+    if "float16" in input_type:
+        return np.dtype(np.float16)
+    if "double" in input_type:
+        return np.dtype(np.float64)
+    return np.dtype(np.float32)
+
+
+def _warp_face(
+    frame: np.ndarray,
+    landmarks: np.ndarray,
+    template: np.ndarray,
+    size: tuple[int, int],
+) -> tuple[np.ndarray, np.ndarray]:
+    destination = template * np.array(size, dtype=np.float32)
+    matrix = cv2.estimateAffinePartial2D(
+        np.asarray(landmarks, dtype=np.float32),
+        destination,
+        method=cv2.RANSAC,
+        ransacReprojThreshold=100,
+    )[0]
+    if matrix is None:
+        raise RuntimeError("Could not align the target face.")
+    crop = cv2.warpAffine(
+        frame,
+        matrix,
+        size,
+        flags=cv2.INTER_LINEAR,
+        borderMode=cv2.BORDER_REPLICATE,
+    )
+    return crop, matrix
+
+
+def _soft_face_mask(size: tuple[int, int], blur: float = 0.12) -> np.ndarray:
+    width, height = size
+    mask = np.zeros((height, width), dtype=np.float32)
+    center = (width // 2, int(height * 0.53))
+    axes = (int(width * 0.43), int(height * 0.48))
+    cv2.ellipse(mask, center, axes, 0, 0, 360, 1.0, -1, cv2.LINE_AA)
+    sigma = max(1.0, width * blur * 0.35)
+    return cv2.GaussianBlur(mask, (0, 0), sigma).clip(0, 1)
+
+
+def _paste_face(
+    frame: np.ndarray,
+    crop: np.ndarray,
+    mask: np.ndarray,
+    affine_matrix: np.ndarray,
+) -> np.ndarray:
+    height, width = frame.shape[:2]
+    inverse_matrix = cv2.invertAffineTransform(affine_matrix)
+    pasted = cv2.warpAffine(
+        crop,
+        inverse_matrix,
+        (width, height),
+        flags=cv2.INTER_LINEAR,
+        borderMode=cv2.BORDER_REPLICATE,
+    )
+    pasted_mask = cv2.warpAffine(
+        mask,
+        inverse_matrix,
+        (width, height),
+        flags=cv2.INTER_LINEAR,
+    ).clip(0, 1)[..., None]
+    return (frame * (1 - pasted_mask) + pasted * pasted_mask).astype(np.uint8)
+
+
+def _swap_face_hq(
+    frame: np.ndarray,
+    landmarks: np.ndarray,
+    source_embedding: np.ndarray,
+    swapper: ort.InferenceSession,
+) -> np.ndarray:
+    size = (256, 256)
+    crop, matrix = _warp_face(frame, landmarks, ARCFACE_128_TEMPLATE, size)
+    target = crop[:, :, ::-1].astype(np.float32) / 255.0
+    target = ((target - 0.5) / 0.5).transpose(2, 0, 1)[None]
+
+    inputs: dict[str, np.ndarray] = {}
+    for item in swapper.get_inputs():
+        if item.name == "source":
+            inputs[item.name] = source_embedding.astype(_input_dtype(swapper, item.name))
+        elif item.name == "target":
+            inputs[item.name] = target.astype(_input_dtype(swapper, item.name))
+    output = swapper.run(None, inputs)[0][0].transpose(1, 2, 0)
+    output = ((output * 0.5 + 0.5).clip(0, 1)[:, :, ::-1] * 255).astype(np.uint8)
+    return _paste_face(frame, output, _soft_face_mask(size), matrix)
+
+
+def _enhance_face(
+    frame: np.ndarray,
+    landmarks: np.ndarray,
+    enhancer: ort.InferenceSession,
+    blend: float = 0.75,
+) -> np.ndarray:
+    size = (512, 512)
+    crop, matrix = _warp_face(frame, landmarks, FFHQ_512_TEMPLATE, size)
+    prepared = crop[:, :, ::-1].astype(np.float32) / 255.0
+    prepared = ((prepared - 0.5) / 0.5).transpose(2, 0, 1)[None]
+
+    inputs: dict[str, np.ndarray] = {}
+    for item in enhancer.get_inputs():
+        if item.name == "input":
+            inputs[item.name] = prepared.astype(_input_dtype(enhancer, item.name))
+        elif item.name == "weight":
+            inputs[item.name] = np.array([0.5], dtype=_input_dtype(enhancer, item.name))
+    output = enhancer.run(None, inputs)[0][0].clip(-1, 1)
+    output = (((output + 1) * 0.5).transpose(1, 2, 0)[:, :, ::-1] * 255).astype(np.uint8)
+    enhanced = _paste_face(frame, output, _soft_face_mask(size, blur=0.10), matrix)
+    return cv2.addWeighted(frame, 1.0 - blend, enhanced, blend, 0)
+
+
 def _resize_to_limit(frame: np.ndarray, resolution: str) -> np.ndarray:
     limits = {"720p (fast)": 1280, "1080p": 1920, "Original": None}
     limit = limits.get(resolution, 1280)
@@ -217,6 +364,7 @@ def swap_video_face(
     target_face_index: int,
     clip_length: int,
     output_resolution: str,
+    quality_mode: str,
     keep_audio: bool,
     add_watermark: bool,
     consent: bool,
@@ -235,7 +383,7 @@ def swap_video_face(
         requested_seconds = min(requested_seconds, duration)
     frame_limit = max(1, int(round(requested_seconds * fps)))
 
-    analyser, swapper = _load_models()
+    analyser, swapper, enhancer = _load_models()
     source_frame = cv2.imread(source_image)
     if source_frame is None:
         raise gr.Error("The source image format is unsupported.")
@@ -244,6 +392,7 @@ def swap_video_face(
     if not source_faces:
         raise gr.Error("No face was detected in the source image. Use a clear, front-facing portrait.")
     source_face = _largest_face(source_faces)
+    source_embedding = _embedding(source_face).reshape(1, -1)
 
     job_dir = OUTPUT_ROOT / uuid.uuid4().hex
     job_dir.mkdir(parents=True, exist_ok=False)
@@ -253,6 +402,8 @@ def swap_video_face(
     cap = cv2.VideoCapture(target_video)
     writer: cv2.VideoWriter | None = None
     tracked_embedding: np.ndarray | None = None
+    smoothed_landmarks: np.ndarray | None = None
+    missed_face_frames = 0
     swapped_frames = 0
     processed_frames = 0
     face_index = max(1, int(target_face_index))
@@ -288,8 +439,23 @@ def swap_video_face(
                     target_face = _select_tracked_face(faces, tracked_embedding)
 
                 if target_face is not None:
-                    frame = swapper.get(frame, target_face, source_face, paste_back=True)
+                    current_landmarks = np.asarray(target_face.kps, dtype=np.float32)
+                    if smoothed_landmarks is None or missed_face_frames > 2:
+                        smoothed_landmarks = current_landmarks
+                    else:
+                        smoothed_landmarks = 0.72 * current_landmarks + 0.28 * smoothed_landmarks
+                    missed_face_frames = 0
+                    frame = _swap_face_hq(
+                        frame,
+                        smoothed_landmarks,
+                        source_embedding,
+                        swapper,
+                    )
+                    if quality_mode == QUALITY_BEST:
+                        frame = _enhance_face(frame, smoothed_landmarks, enhancer)
                     swapped_frames += 1
+            else:
+                missed_face_frames += 1
 
             if add_watermark:
                 frame = _add_ai_watermark(frame)
@@ -329,7 +495,7 @@ def swap_video_face(
 
     status = (
         f"Completed {processed_seconds:.1f}s at {fps:.2f} FPS — "
-        f"face replaced in {swapped_frames}/{processed_frames} frames."
+        f"face replaced in {swapped_frames}/{processed_frames} frames using {quality_mode}."
     )
     return str(final_path), status
 
@@ -346,8 +512,8 @@ with gr.Blocks(title=APP_TITLE, theme=gr.themes.Soft(), css=CSS) as demo:
         """
         <div class="hero">
           <h1>🎭 Dream Video Face Swap</h1>
-          <p>ZeroGPU face replacement with identity tracking and original-audio preservation.</p>
-          <p dir="rtl">تبديل الوجه في الفيديو باستخدام ZeroGPU مع تتبع الشخص والاحتفاظ بالصوت الأصلي.</p>
+          <p>High-quality ZeroGPU face replacement with HyperSwap, enhancement, identity tracking, and audio preservation.</p>
+          <p dir="rtl">تبديل وجه عالي الجودة عبر ZeroGPU باستخدام HyperSwap وتحسين تفاصيل الوجه مع تتبع الشخص والاحتفاظ بالصوت.</p>
         </div>
         """
     )
@@ -383,9 +549,15 @@ with gr.Blocks(title=APP_TITLE, theme=gr.themes.Soft(), css=CSS) as demo:
             )
             output_resolution = gr.Radio(
                 choices=["720p (fast)", "1080p", "Original"],
-                value="720p (fast)",
+                value="1080p",
                 label="Output resolution",
             )
+            quality_mode = gr.Radio(
+                choices=[QUALITY_BEST, QUALITY_BALANCED],
+                value=QUALITY_BEST,
+                label="Quality / الجودة",
+                info="Best quality restores facial detail at 512×512; Balanced is faster.",
+            )
             keep_audio = gr.Checkbox(value=True, label="Keep original audio")
             add_watermark = gr.Checkbox(value=True, label="Add ‘AI FACE SWAP’ watermark")
             consent = gr.Checkbox(
@@ -401,7 +573,7 @@ with gr.Blocks(title=APP_TITLE, theme=gr.themes.Soft(), css=CSS) as demo:
         """
 <div class="notice">
 
-**Tips / نصائح:** Use a sharp, front-facing source portrait. For multiple people, select the target by its left-to-right position in the first clear frame. Start with 720p and a short clip for the fastest ZeroGPU queue.
+**Tips / نصائح:** Use a sharp, front-facing source portrait. For multiple people, select the target by its left-to-right position in the first clear frame. “Best quality” uses HyperSwap 256 plus GFPGAN 512; start with a 5–10 second clip.
 
 استخدم صورة أمامية واضحة للوجه. إذا ظهر أكثر من شخص، اختر رقم الشخص حسب ترتيبه من اليسار إلى اليمين في أول لقطة واضحة.
 
@@ -419,6 +591,7 @@ Uploads are processed temporarily and old results are automatically removed. Do
             target_face_index,
             clip_length,
             output_resolution,
+            quality_mode,
             keep_audio,
             add_watermark,
             consent,