Update app.py

app.py

@@ -2,9 +2,9 @@ import os
 import cv2
 import math
 import shutil
-import subprocess
 import tempfile
 import traceback
+import subprocess
 import urllib.request
 
 import gradio as gr
@@ -14,13 +14,12 @@ import onnxruntime as ort
 
 
 # ============================================================
-# CONFIG
+# CONFIGURATION
 # ============================================================
 
 MODEL_PATH = "wav2lip_models/wav2lip_gan.onnx"
 
 OUTPUT_DIR = "outputs"
-
 FACE_DIR = "face_detector"
 
 FACE_PROTO = os.path.join(
@@ -33,13 +32,13 @@ FACE_MODEL = os.path.join(
     "res10_300x300_ssd_iter_140000.caffemodel"
 )
 
-PROTO_URL = (
+FACE_PROTO_URL = (
     "https://raw.githubusercontent.com/"
     "opencv/opencv/master/samples/dnn/"
     "face_detector/deploy.prototxt"
 )
 
-MODEL_URL = (
+FACE_MODEL_URL = (
     "https://github.com/opencv/opencv_3rdparty/"
     "raw/dnn_samples_face_detector_20170830/"
     "res10_300x300_ssd_iter_140000.caffemodel"
@@ -47,7 +46,7 @@ MODEL_URL = (
 
 
 # ============================================================
-# WAV2LIP AUDIO PARAMETERS
+# AUDIO
 # ============================================================
 
 SAMPLE_RATE = 16000
@@ -63,39 +62,52 @@ FMAX = 7600
 
 MEL_STEP_SIZE = 16
 
-# Original Wav2Lip uses approximately 80 mel frames/sec
+# Wav2Lip audio timing
 MEL_FPS = 80.0
 
 
 # ============================================================
-# VIDEO / FACE PARAMETERS
+# FACE
 # ============================================================
 
 FACE_SIZE = 96
 
 FACE_CONFIDENCE = 0.50
 
-# Extra face crop around detected face
-FACE_PADDING = 0.20
+# Do NOT make this too large.
+#
+# A large crop followed by 96x96 resizing causes
+# the mouth to become soft/distorted.
+FACE_PADDING = 0.08
 
-# Temporal face smoothing
-FACE_SMOOTHING = 0.15
+# Face tracking smoothing
+FACE_SMOOTHING = 0.25
 
-# Mouth blend strength
-#
-# 0.70 = stronger lip replacement
-# 0.55 = balanced
-# 0.40 = more original face
-MOUTH_BLEND = 0.62
+# How frequently to re-detect face
+REDETECT_SECONDS = 0.40
 
 
 # ============================================================
-# OUTPUT QUALITY
+# BLENDING
 # ============================================================
 
-FFMPEG_CRF = 17
+# 1.0 = strongest generated mouth
+# 0.7 = balanced
+# 0.55 = preserve original texture
+BLEND_STRENGTH = 0.72
+
+
+# Lower-face mask starts around this percentage
+# of the crop.
+LOWER_FACE_START = 0.47
+
 
-FFMPEG_PRESET = "medium"
+# ============================================================
+# VIDEO ENCODING
+# ============================================================
+
+CRF = 17
+PRESET = "medium"
 
 
 os.makedirs(
@@ -120,52 +132,56 @@ print("=" * 70)
 if not os.path.exists(MODEL_PATH):
 
     raise FileNotFoundError(
-        "\nWav2Lip ONNX model not found:\n"
+        "\n\nWav2Lip ONNX model not found:\n\n"
         f"{MODEL_PATH}\n\n"
-        "Expected:\n"
-        "wav2lip_models/wav2lip_gan.onnx"
+        "Your Space should contain:\n\n"
+        "wav2lip_models/\n"
+        "    wav2lip_gan.onnx\n"
     )
 
 
 # ============================================================
-# ONNX PROVIDERS
+# ONNX
 # ============================================================
 
 available_providers = (
     ort.get_available_providers()
 )
 
+print(
+    "Available ONNX providers:",
+    available_providers
+)
+
+
 if "CUDAExecutionProvider" in available_providers:
 
-    providers = [
+    PROVIDERS = [
         "CUDAExecutionProvider",
         "CPUExecutionProvider"
     ]
 
 else:
 
-    providers = [
+    PROVIDERS = [
         "CPUExecutionProvider"
     ]
 
 
 print(
-    "ONNX providers:",
-    providers
+    "Using providers:",
+    PROVIDERS
 )
 
 
-# ============================================================
-# LOAD WAV2LIP
-# ============================================================
-
 print(
     "Loading Wav2Lip ONNX model..."
 )
 
+
 SESSION = ort.InferenceSession(
     MODEL_PATH,
-    providers=providers
+    providers=PROVIDERS
 )
 
 
@@ -203,10 +219,10 @@ for out in SESSION.get_outputs():
 
 
 # ============================================================
-# FACE DETECTOR
+# FACE DETECTOR DOWNLOAD
 # ============================================================
 
-def download_file(
+def download_if_needed(
     url,
     path
 ):
@@ -228,49 +244,53 @@ def download_file(
     )
 
 
-    urllib.request.urlretrieve(
-        url,
-        path
-    )
-
+    try:
 
-def load_face_detector():
+        urllib.request.urlretrieve(
+            url,
+            path
+        )
 
-    download_file(
-        PROTO_URL,
-        FACE_PROTO
-    )
+    except Exception as e:
 
-    download_file(
-        MODEL_URL,
-        FACE_MODEL
-    )
+        raise RuntimeError(
+            f"Could not download {url}\n{e}"
+        )
 
 
-    print(
-        "Loading OpenCV DNN face detector..."
-    )
+download_if_needed(
+    FACE_PROTO_URL,
+    FACE_PROTO
+)
 
+download_if_needed(
+    FACE_MODEL_URL,
+    FACE_MODEL
+)
 
-    detector = cv2.dnn.readNetFromCaffe(
-        FACE_PROTO,
-        FACE_MODEL
-    )
 
+# ============================================================
+# FACE DETECTOR
+# ============================================================
 
-    print(
-        "Face detector loaded."
-    )
+print(
+    "Loading OpenCV DNN face detector..."
+)
 
 
-    return detector
+FACE_DETECTOR = cv2.dnn.readNetFromCaffe(
+    FACE_PROTO,
+    FACE_MODEL
+)
 
 
-FACE_DETECTOR = load_face_detector()
+print(
+    "Face detector loaded."
+)
 
 
 # ============================================================
-# AUDIO LOADING
+# AUDIO
 # ============================================================
 
 def load_audio(
@@ -280,31 +300,31 @@ def load_audio(
     if not path:
 
         raise ValueError(
-            "Audio file is required."
+            "Target audio is required."
         )
 
 
-    wav, sr = librosa.load(
+    audio, sr = librosa.load(
         path,
         sr=SAMPLE_RATE,
         mono=True
     )
 
 
-    if wav is None or len(wav) == 0:
+    if audio is None or len(audio) == 0:
 
         raise ValueError(
-            "Could not read audio."
+            "Could not read target audio."
         )
 
 
-    wav = wav.astype(
+    audio = audio.astype(
         np.float32
     )
 
 
     duration = (
-        len(wav)
+        len(audio)
         / SAMPLE_RATE
     )
 
@@ -316,62 +336,45 @@ def load_audio(
     )
 
 
-    return wav, duration
+    return audio, duration
 
 
 # ============================================================
-# WAV2LIP-COMPATIBLE MEL SPECTROGRAM
+# WAV2LIP MEL
 # ============================================================
 
-def wav2lip_mel(
-    wav
+def make_mel(
+    audio
 ):
 
     """
-    Wav2Lip-style mel preprocessing.
+    Wav2Lip-style mel spectrogram.
+
+    Important:
+    The ONNX model expects exactly:
 
-    This is intentionally closer to the original Wav2Lip
-    audio preprocessing than the previous implementation.
+        [1, 1, 80, 16]
     """
 
 
-    # Pre-emphasis / STFT
-    stft = librosa.stft(
-        wav,
+    mel = librosa.feature.melspectrogram(
+        y=audio,
+        sr=SAMPLE_RATE,
         n_fft=N_FFT,
         hop_length=HOP_LENGTH,
         win_length=WIN_LENGTH,
         window="hann",
-        center=True
-    )
-
-
-    # Magnitude
-    magnitudes = np.abs(
-        stft
-    )
-
-
-    # Mel filter bank
-    mel_filters = librosa.filters.mel(
-        sr=SAMPLE_RATE,
-        n_fft=N_FFT,
         n_mels=N_MELS,
         fmin=FMIN,
-        fmax=FMAX
-    )
-
-
-    mel = np.dot(
-        mel_filters,
-        magnitudes
+        fmax=FMAX,
+        power=1.0,
+        center=True
     )
 
 
-    # Wav2Lip-style log compression
     mel = np.maximum(
-        1e-5,
-        mel
+        mel,
+        1e-5
     )
 
 
@@ -389,13 +392,11 @@ def wav2lip_mel(
 # FACE DETECTION
 # ============================================================
 
-def detect_faces(
+def detect_face(
     frame
 ):
 
-    height, width = (
-        frame.shape[:2]
-    )
+    h, w = frame.shape[:2]
 
 
     blob = cv2.dnn.blobFromImage(
@@ -406,9 +407,9 @@ def detect_faces(
         1.0,
         (300, 300),
         (
-            104.0,
-            177.0,
-            123.0
+            104,
+            177,
+            123
         )
     )
 
@@ -423,7 +424,8 @@ def detect_faces(
     )
 
 
-    faces = []
+    best = None
+    best_score = 0.0
 
 
     for i in range(
@@ -454,10 +456,10 @@ def detect_faces(
             ]
             * np.array(
                 [
-                    width,
-                    height,
-                    width,
-                    height
+                    w,
+                    h,
+                    w,
+                    h
                 ]
             )
         )
@@ -472,7 +474,7 @@ def detect_faces(
             0,
             min(
                 x1,
-                width - 1
+                w - 1
             )
         )
 
@@ -481,7 +483,7 @@ def detect_faces(
             0,
             min(
                 y1,
-                height - 1
+                h - 1
             )
         )
 
@@ -490,7 +492,7 @@ def detect_faces(
             0,
             min(
                 x2,
-                width
+                w
             )
         )
 
@@ -499,164 +501,48 @@ def detect_faces(
             0,
             min(
                 y2,
-                height
+                h
             )
         )
 
 
-        w = x2 - x1
-        h = y2 - y1
+        bw = x2 - x1
+        bh = y2 - y1
 
 
-        if w < 50 or h < 50:
+        if bw < 40 or bh < 40:
 
             continue
 
 
-        faces.append(
-            {
-                "x": x1,
-                "y": y1,
-                "w": w,
-                "h": h,
-                "confidence": confidence
-            }
+        area = (
+            bw * bh
         )
 
 
-    return faces
-
-
-# ============================================================
-# BEST FACE
-# ============================================================
-
-def best_face(
-    frame
-):
-
-    faces = detect_faces(
-        frame
-    )
-
-
-    if not faces:
-
-        return None
-
-
-    faces.sort(
-        key=lambda f: (
-            f["w"] * f["h"],
-            f["confidence"]
-        ),
-        reverse=True
-    )
-
-
-    f = faces[0]
-
-
-    return (
-        f["x"],
-        f["y"],
-        f["w"],
-        f["h"]
-    )
-
-
-# ============================================================
-# FACE BOX WITH PADDING
-# ============================================================
-
-def padded_square_box(
-    frame,
-    box
-):
-
-    height, width = (
-        frame.shape[:2]
-    )
-
-
-    x, y, w, h = box
-
-
-    # Wav2Lip needs a little surrounding context
-    size = max(
-        w,
-        h
-    )
-
-
-    size = int(
-        size
-        * (
-            1.0
-            + FACE_PADDING
+        score = (
+            confidence
+            * math.sqrt(area)
         )
-    )
 
 
-    size = min(
-        size,
-        width,
-        height
-    )
-
-
-    cx = (
-        x
-        + w / 2.0
-    )
+        if score > best_score:
 
+            best_score = score
 
-    cy = (
-        y
-        + h / 2.0
-    )
-
-
-    x1 = int(
-        cx
-        - size / 2
-    )
-
-
-    y1 = int(
-        cy
-        - size / 2
-    )
-
-
-    x1 = max(
-        0,
-        min(
-            x1,
-            width - size
-        )
-    )
-
-
-    y1 = max(
-        0,
-        min(
-            y1,
-            height - size
-        )
-    )
+            best = (
+                x1,
+                y1,
+                bw,
+                bh
+            )
 
 
-    return (
-        int(x1),
-        int(y1),
-        int(size),
-        int(size)
-    )
+    return best
 
 
 # ============================================================
-# INITIAL FACE BOX
+# FIND FACE FROM SEVERAL FRAMES
 # ============================================================
 
 def find_initial_face(
@@ -671,7 +557,7 @@ def find_initial_face(
     if not cap.isOpened():
 
         raise RuntimeError(
-            "Could not open video."
+            "Could not open source video."
         )
 
 
@@ -687,12 +573,12 @@ def find_initial_face(
         cap.release()
 
         raise RuntimeError(
-            "Video has no frames."
+            "Source video contains no frames."
         )
 
 
-    sample_count = min(
-        15,
+    samples = min(
+        12,
         total
     )
 
@@ -701,19 +587,19 @@ def find_initial_face(
 
 
     for i in range(
-        sample_count
+        samples
     ):
 
-        index = int(
+        frame_number = int(
             i
             * total
-            / sample_count
+            / samples
         )
 
 
         cap.set(
             cv2.CAP_PROP_POS_FRAMES,
-            index
+            frame_number
         )
 
 
@@ -727,7 +613,7 @@ def find_initial_face(
             continue
 
 
-        box = best_face(
+        box = detect_face(
             frame
         )
 
@@ -746,20 +632,20 @@ def find_initial_face(
 
         raise RuntimeError(
             "No face detected.\n\n"
-            "Please use a video with a clearly "
-            "visible face."
+            "Please use a video where the face is "
+            "clearly visible."
         )
 
 
-    boxes = np.asarray(
+    arr = np.asarray(
         boxes,
         dtype=np.float32
     )
 
 
-    # Median is more stable than taking the first frame
+    # Median is much more stable than first frame.
     median = np.median(
-        boxes,
+        arr,
         axis=0
     )
 
@@ -770,27 +656,27 @@ def find_initial_face(
 
 
 # ============================================================
-# FACE BOX SMOOTHING
+# SMOOTH FACE BOX
 # ============================================================
 
-def smooth_box(
-    old_box,
-    new_box
+def smooth_face(
+    old,
+    new
 ):
 
-    if old_box is None:
+    if old is None:
 
-        return new_box
+        return new
 
 
     old = np.asarray(
-        old_box,
+        old,
         dtype=np.float32
     )
 
 
     new = np.asarray(
-        new_box,
+        new,
         dtype=np.float32
     )
 
@@ -813,253 +699,160 @@ def smooth_box(
 
 
 # ============================================================
-# MOUTH MASK
+# MAKE FACE CROP
 # ============================================================
 
-def create_mouth_mask(
-    size
+def make_crop(
+    frame,
+    face_box
 ):
 
-    mask = np.zeros(
-        (
-            size,
-            size
-        ),
-        dtype=np.float32
+    frame_h, frame_w = (
+        frame.shape[:2]
     )
 
 
-    # Important:
-    # Don't use a rectangular mouth region.
-    #
-    # An ellipse gives much smoother transitions and
-    # avoids visible square/rectangular artifacts.
-
-
-    center_x = int(
-        size * 0.50
-    )
+    x, y, w, h = face_box
 
 
-    center_y = int(
-        size * 0.665
+    # Keep crop close to actual face.
+    size = int(
+        max(w, h)
+        * (
+            1.0
+            + FACE_PADDING
+        )
     )
 
 
-    axis_x = int(
-        size * 0.285
+    # Prevent crop from becoming enormous.
+    size = min(
+        size,
+        frame_w,
+        frame_h
     )
 
 
-    axis_y = int(
-        size * 0.155
+    # Face center
+    cx = (
+        x
+        + w / 2.0
     )
 
 
-    cv2.ellipse(
-        mask,
-        (
-            center_x,
-            center_y
-        ),
-        (
-            axis_x,
-            axis_y
-        ),
-        0,
-        0,
-        360,
-        1.0,
-        -1
+    cy = (
+        y
+        + h / 2.0
     )
 
 
-    # Extra blur creates a natural transition
-    sigma = max(
-        2,
-        int(
-            size * 0.045
-        )
+    x1 = int(
+        cx
+        - size / 2
     )
 
 
-    mask = cv2.GaussianBlur(
-        mask,
-        (0, 0),
-        sigma
+    y1 = int(
+        cy
+        - size / 2
     )
 
 
-    return np.clip(
-        mask,
+    # Keep inside frame
+    x1 = max(
         0,
-        1
-    )
-
-
-MOUTH_MASK = create_mouth_mask(
-    FACE_SIZE
-)
-
-
-# ============================================================
-# MOUTH-ONLY COLOR CORRECTION
-# ============================================================
-
-def color_correct(
-    generated,
-    original,
-    mask
-):
-
-    result = generated.astype(
-        np.float32
-    )
-
-
-    original_float = (
-        original.astype(
-            np.float32
+        min(
+            x1,
+            frame_w - size
         )
     )
 
 
-    strong = (
-        mask > 0.35
-    )
-
-
-    if np.count_nonzero(
-        strong
-    ) < 50:
-
-        return generated
-
-
-    for c in range(3):
-
-        src = (
-            original_float[
-                :,
-                :,
-                c
-            ][strong]
-        )
-
-
-        dst = (
-            result[
-                :,
-                :,
-                c
-            ][strong]
-        )
-
-
-        if (
-            len(src) == 0
-            or len(dst) == 0
-        ):
-
-            continue
-
-
-        src_mean = np.mean(
-            src
-        )
-
-
-        dst_mean = np.mean(
-            dst
+    y1 = max(
+        0,
+        min(
+            y1,
+            frame_h - size
         )
+    )
 
 
-        # Only mild correction.
-        # Strong correction can make lips look unnatural.
-        correction = (
-            src_mean
-            - dst_mean
-        ) * 0.35
-
-
-        result[
-            :,
-            :,
-            c
-        ] += correction
-
-
-    return np.clip(
-        result,
-        0,
-        255
-    ).astype(
-        np.uint8
+    return (
+        int(x1),
+        int(y1),
+        int(size),
+        int(size)
     )
 
 
 # ============================================================
-# WAV2LIP INFERENCE
+# WAV2LIP MODEL
 # ============================================================
 
-def run_wav2lip(
+def inference(
     face,
-    mel
+    mel_chunk
 ):
 
     # --------------------------------------------------------
-    # Resize to model input
+    # 96x96
     # --------------------------------------------------------
 
-    face_96 = cv2.resize(
+    face96 = cv2.resize(
         face,
         (
-            FACE_SIZE,
-            FACE_SIZE
+            96,
+            96
         ),
         interpolation=cv2.INTER_AREA
     )
 
 
     # --------------------------------------------------------
-    # Mask lower half
+    # ORIGINAL WAV2LIP MASKING
+    #
+    # Upper half remains original.
+    # Lower half is masked.
     # --------------------------------------------------------
 
-    masked = face_96.copy()
+    masked = face96.copy()
 
 
     masked[
-        FACE_SIZE // 2:
+        48:,
+        :
     ] = 0
 
 
     # --------------------------------------------------------
     # 6 channels
     #
-    # [masked face RGB]
-    # [original face RGB]
+    # [masked RGB]
+    # [original RGB]
     # --------------------------------------------------------
 
-    combined = np.concatenate(
+    input_image = np.concatenate(
         [
             masked,
-            face_96
+            face96
         ],
         axis=2
     )
 
 
-    combined = (
-        combined.astype(
-            np.float32
-        )
+    # --------------------------------------------------------
+    # HWC -> CHW
+    # --------------------------------------------------------
+
+    input_image = (
+        input_image
+        .astype(np.float32)
         / 255.0
     )
 
 
-    # HWC -> CHW
-    combined = (
-        combined.transpose(
+    input_image = (
+        input_image.transpose(
             2,
             0,
             1
@@ -1067,8 +860,8 @@ def run_wav2lip(
     )
 
 
-    video_input = (
-        combined[
+    input_image = (
+        input_image[
             np.newaxis,
             ...
         ]
@@ -1079,11 +872,11 @@ def run_wav2lip(
 
 
     # --------------------------------------------------------
-    # Mel
+    # AUDIO
     # --------------------------------------------------------
 
     mel_input = (
-        mel[
+        mel_chunk[
             np.newaxis,
             np.newaxis,
             :,
@@ -1099,23 +892,19 @@ def run_wav2lip(
     # ONNX
     # --------------------------------------------------------
 
-    outputs = SESSION.run(
+    result = SESSION.run(
         None,
         {
             "mel_spectrogram":
                 mel_input,
 
             "video_frames":
-                video_input
+                input_image
         }
     )
 
 
-    prediction = outputs[0]
-
-
-    # [1,3,96,96] -> [3,96,96]
-    prediction = prediction[0]
+    prediction = result[0][0]
 
 
     # --------------------------------------------------------
@@ -1132,43 +921,154 @@ def run_wav2lip(
 
 
     prediction = np.clip(
-        prediction * 255.0,
+        prediction,
         0,
-        255
+        1
+    )
+
+
+    prediction = (
+        prediction
+        * 255.0
     ).astype(
         np.uint8
     )
 
 
-    # Model output is RGB
+    # Model returns RGB
     prediction = cv2.cvtColor(
         prediction,
         cv2.COLOR_RGB2BGR
     )
 
 
-    # --------------------------------------------------------
-    # Back to original crop resolution
-    # --------------------------------------------------------
+    return prediction
 
-    prediction = cv2.resize(
-        prediction,
+
+# ============================================================
+# SOFT LOWER FACE MASK
+# ============================================================
+
+def make_blend_mask(
+    h,
+    w
+):
+
+    mask = np.zeros(
         (
-            face.shape[1],
-            face.shape[0]
+            h,
+            w
         ),
-        interpolation=cv2.INTER_CUBIC
+        dtype=np.float32
     )
 
 
-    return prediction
+    start = int(
+        h
+        * LOWER_FACE_START
+    )
+
+
+    # Start from zero at upper edge.
+    # Gradually increase toward mouth.
+    #
+    # This avoids a visible rectangular patch.
+
+
+    for y in range(
+        start,
+        h
+    ):
+
+        position = (
+            y - start
+        ) / max(
+            1,
+            h - start
+        )
+
+
+        # Smoothstep
+        value = (
+            position
+            * position
+            * (
+                3
+                - 2 * position
+            )
+        )
+
+
+        mask[y, :] = value
+
+
+    # Stronger around mouth center
+    yy, xx = np.mgrid[
+        0:h,
+        0:w
+    ]
+
+
+    cx = w * 0.50
+    cy = h * 0.68
+
+    rx = w * 0.42
+    ry = h * 0.30
+
+
+    ellipse = (
+        (
+            (xx - cx) / rx
+        ) ** 2
+        +
+        (
+            (yy - cy) / ry
+        ) ** 2
+    )
+
+
+    ellipse_mask = np.exp(
+        -ellipse * 1.8
+    )
+
+
+    # Combine lower-face and mouth masks.
+    mask = (
+        mask
+        * (
+            0.45
+            + 0.55
+            * ellipse_mask
+        )
+    )
+
+
+    # Smooth edges.
+    mask = cv2.GaussianBlur(
+        mask,
+        (
+            0,
+            0
+        ),
+        sigmaX=max(
+            1.0,
+            w * 0.035
+        )
+    )
+
+
+    return np.clip(
+        mask,
+        0,
+        1
+    )
 
 
 # ============================================================
-# BLEND RESULT
+# BLEND
 # ============================================================
 
-def blend_mouth(
+def blend_prediction(
     original,
     generated
 ):
@@ -1188,43 +1088,114 @@ def blend_mouth(
     )
 
 
-    mask = cv2.resize(
-        MOUTH_MASK,
-        (
-            w,
-            h
-        ),
-        interpolation=cv2.INTER_CUBIC
+    mask = make_blend_mask(
+        h,
+        w
     )
 
 
-    # Mild color correction
-    generated = color_correct(
-        generated,
-        original,
-        mask
+    # --------------------------------------------------------
+    # Slight color matching.
+    #
+    # Do NOT aggressively change colors.
+    # --------------------------------------------------------
+
+    original_float = (
+        original.astype(
+            np.float32
+        )
+    )
+
+
+    generated_float = (
+        generated.astype(
+            np.float32
+        )
+    )
+
+
+    mouth_area = (
+        mask > 0.55
+    )
+
+
+    if np.count_nonzero(
+        mouth_area
+    ) > 100:
+
+        for c in range(3):
+
+            src = (
+                original_float[
+                    :,
+                    :,
+                    c
+                ][mouth_area]
+            )
+
+
+            gen = (
+                generated_float[
+                    :,
+                    :,
+                    c
+                ][mouth_area]
+            )
+
+
+            if (
+                len(src) > 0
+                and len(gen) > 0
+            ):
+
+                src_mean = np.mean(
+                    src
+                )
+
+                gen_mean = np.mean(
+                    gen
+                )
+
+
+                correction = (
+                    src_mean
+                    - gen_mean
+                ) * 0.20
+
+
+                generated_float[
+                    :,
+                    :,
+                    c
+                ] += correction
+
+
+    generated_float = np.clip(
+        generated_float,
+        0,
+        255
     )
 
 
+    # --------------------------------------------------------
+    # Blend
+    # --------------------------------------------------------
+
     alpha = (
         mask[
             :,
             :,
             np.newaxis
         ]
-        * MOUTH_BLEND
+        * BLEND_STRENGTH
     )
 
 
-    result = (
-        generated.astype(
-            np.float32
-        )
+    output = (
+        generated_float
         * alpha
         +
-        original.astype(
-            np.float32
-        )
+        original_float
         * (
             1.0
             - alpha
@@ -1233,7 +1204,7 @@ def blend_mouth(
 
 
     return np.clip(
-        result,
+        output,
         0,
         255
     ).astype(
@@ -1245,34 +1216,34 @@ def blend_mouth(
 # FFMPEG
 # ============================================================
 
-def get_ffmpeg():
+def ffmpeg_path():
 
-    ffmpeg = shutil.which(
+    path = shutil.which(
         "ffmpeg"
     )
 
 
-    if ffmpeg is None:
+    if path is None:
 
         raise RuntimeError(
             "FFmpeg is not installed."
         )
 
 
-    return ffmpeg
+    return path
 
 
 # ============================================================
 # MUX AUDIO
 # ============================================================
 
-def mux_audio(
-    video,
+def add_audio(
+    silent_video,
     audio,
     output
 ):
 
-    ffmpeg = get_ffmpeg()
+    ffmpeg = ffmpeg_path()
 
 
     command = [
@@ -1280,7 +1251,7 @@ def mux_audio(
         "-y",
 
         "-i",
-        video,
+        silent_video,
 
         "-i",
         audio,
@@ -1295,10 +1266,10 @@ def mux_audio(
         "libx264",
 
         "-preset",
-        FFMPEG_PRESET,
+        PRESET,
 
         "-crf",
-        str(FFMPEG_CRF),
+        str(CRF),
 
         "-pix_fmt",
         "yuv420p",
@@ -1334,32 +1305,79 @@ def mux_audio(
 
 
         raise RuntimeError(
-            "FFmpeg audio muxing failed."
+            "Could not add audio to output video."
         )
 
 
 # ============================================================
-# MAIN PROCESS
+# PROCESS
 # ============================================================
 
 def process_video(
-    video_path,
-    audio_path,
-    text,
+    video_input,
+    audio_input,
+    text_input,
     progress=gr.Progress()
 ):
 
+    # ========================================================
+    # GET VIDEO PATH
+    # ========================================================
+
+    video_path = video_input
+
+
+    # Gradio can return a dictionary in some configurations.
+    if isinstance(
+        video_input,
+        dict
+    ):
+
+        video_path = (
+            video_input.get(
+                "path"
+            )
+            or
+            video_input.get(
+                "name"
+            )
+        )
+
+
+    # ========================================================
+    # AUDIO PATH
+    # ========================================================
+
+    audio_path = audio_input
+
+
+    if isinstance(
+        audio_input,
+        dict
+    ):
+
+        audio_path = (
+            audio_input.get(
+                "path"
+            )
+            or
+            audio_input.get(
+                "name"
+            )
+        )
+
+
     if not video_path:
 
         raise gr.Error(
-            "Please upload a video."
+            "Please upload a source video."
         )
 
 
     if not audio_path:
 
         raise gr.Error(
-            "Please upload audio."
+            "Please upload target audio."
         )
 
 
@@ -1374,9 +1392,9 @@ def process_video(
     )
 
 
-    output_video = os.path.join(
+    output_path = os.path.join(
         OUTPUT_DIR,
-        "wav2lip_hq_result.mp4"
+        "wav2lip_result.mp4"
     )
 
 
@@ -1388,11 +1406,11 @@ def process_video(
 
         progress(
             0.02,
-            desc="Loading audio..."
+            desc="Loading target audio..."
         )
 
 
-        wav, audio_duration = (
+        audio, audio_duration = (
             load_audio(
                 audio_path
             )
@@ -1400,13 +1418,13 @@ def process_video(
 
 
         progress(
-            0.06,
-            desc="Creating Wav2Lip mel spectrogram..."
+            0.07,
+            desc="Creating mel spectrogram..."
         )
 
 
-        mel = wav2lip_mel(
-            wav
+        mel = make_mel(
+            audio
         )
 
 
@@ -1416,13 +1434,21 @@ def process_video(
         )
 
 
+        if mel.shape[0] != 80:
+
+            raise RuntimeError(
+                "Mel spectrogram has wrong number "
+                f"of bands: {mel.shape[0]}"
+            )
+
+
         # ====================================================
         # VIDEO
         # ====================================================
 
         progress(
             0.10,
-            desc="Opening video..."
+            desc="Opening source video..."
         )
 
 
@@ -1434,7 +1460,7 @@ def process_video(
         if not cap.isOpened():
 
             raise RuntimeError(
-                "Could not open video."
+                "Could not open source video."
             )
 
 
@@ -1444,7 +1470,7 @@ def process_video(
 
 
         if (
-            fps is None
+            not fps
             or fps <= 0
         ):
 
@@ -1472,15 +1498,6 @@ def process_video(
         )
 
 
-        if total_frames <= 0:
-
-            cap.release()
-
-            raise RuntimeError(
-                "Video contains no frames."
-            )
-
-
         print(
             "FPS:",
             fps
@@ -1488,7 +1505,7 @@ def process_video(
 
 
         print(
-            "Resolution:",
+            "Video:",
             width,
             "x",
             height
@@ -1496,22 +1513,22 @@ def process_video(
 
 
         print(
-            "Total frames:",
+            "Frames:",
             total_frames
         )
 
 
         # ====================================================
-        # FACE
+        # FIND FACE
         # ====================================================
 
         progress(
             0.14,
-            desc="Finding face..."
+            desc="Locating face..."
         )
 
 
-        initial_face = (
+        current_face = (
             find_initial_face(
                 video_path
             )
@@ -1520,12 +1537,12 @@ def process_video(
 
         print(
             "Initial face:",
-            initial_face
+            current_face
         )
 
 
         # ====================================================
-        # CREATE VIDEO WRITER
+        # VIDEO WRITER
         # ====================================================
 
         fourcc = (
@@ -1556,10 +1573,10 @@ def process_video(
 
 
         # ====================================================
-        # PROCESS FRAME COUNT
+        # HOW MANY VIDEO FRAMES
         # ====================================================
 
-        audio_frame_count = int(
+        audio_frames = int(
             math.ceil(
                 audio_duration
                 * fps
@@ -1567,43 +1584,53 @@ def process_video(
         )
 
 
-        process_frames = min(
+        frames_to_process = min(
             total_frames,
-            audio_frame_count
+            audio_frames
         )
 
 
         print(
             "Frames to process:",
-            process_frames
+            frames_to_process
         )
 
 
         # ====================================================
-        # FACE TRACKING
+        # FACE REDETECTION
         # ====================================================
 
-        current_face = initial_face
-
-
-        # How often to re-detect
-        #
-        # This prevents the face from drifting when the
-        # person moves.
-        REDETECT_EVERY = max(
+        redetect_every = max(
             1,
             int(
-                fps * 0.5
+                fps
+                * REDETECT_SECONDS
             )
         )
 
 
         # ====================================================
-        # PROCESS
+        # AUDIO MAPPING
+        # ====================================================
+
+        mel_idx_multiplier = (
+            MEL_FPS
+            / fps
+        )
+
+
+        print(
+            "Mel index multiplier:",
+            mel_idx_multiplier
+        )
+
+
+        # ====================================================
+        # FRAME LOOP
         # ====================================================
 
         for frame_index in range(
-            process_frames
+            frames_to_process
         ):
 
             ok, frame = (
@@ -1617,15 +1644,16 @@ def process_video(
 
 
             # ------------------------------------------------
-            # Re-detect periodically
+            # FACE TRACKING
             # ------------------------------------------------
 
             if (
-                frame_index % REDETECT_EVERY
+                frame_index
+                % redetect_every
                 == 0
             ):
 
-                detected = best_face(
+                detected = detect_face(
                     frame
                 )
 
@@ -1633,7 +1661,7 @@ def process_video(
                 if detected is not None:
 
                     current_face = (
-                        smooth_box(
+                        smooth_face(
                             current_face,
                             detected
                         )
@@ -1641,27 +1669,29 @@ def process_video(
 
 
             # ------------------------------------------------
-            # Face crop
+            # FACE CROP
             # ------------------------------------------------
 
-            crop_box = (
-                padded_square_box(
-                    frame,
-                    current_face
-                )
+            crop = make_crop(
+                frame,
+                current_face
             )
 
 
-            x, y, w, h = crop_box
+            x, y, cw, ch = crop
 
 
             face = frame[
-                y:y+h,
-                x:x+w
+                y:y+ch,
+                x:x+cw
             ]
 
 
-            if face.size == 0:
+            if (
+                face.size == 0
+                or face.shape[0] < 20
+                or face.shape[1] < 20
+            ):
 
                 writer.write(
                     frame
@@ -1671,23 +1701,8 @@ def process_video(
 
 
             # ------------------------------------------------
-            # EXACT AUDIO/VIDEO FRAME MAPPING
+            # CORRECT MEL WINDOW
             # ------------------------------------------------
-            #
-            # This is more important for sync than a manually
-            # chosen arbitrary offset.
-            #
-            # Original Wav2Lip style:
-            #
-            # mel_idx_multiplier = 80 / fps
-            #
-            # ------------------------------------------------
-
-            mel_idx_multiplier = (
-                MEL_FPS
-                / fps
-            )
-
 
             mel_start = int(
                 frame_index
@@ -1702,7 +1717,7 @@ def process_video(
 
 
             # ------------------------------------------------
-            # Handle final audio frames
+            # End of audio
             # ------------------------------------------------
 
             if mel_start >= (
@@ -1729,7 +1744,7 @@ def process_video(
 
 
             # ------------------------------------------------
-            # Pad final chunk
+            # EXACTLY 16 MEL FRAMES
             # ------------------------------------------------
 
             if (
@@ -1750,8 +1765,8 @@ def process_video(
 
                     mel_chunk = np.zeros(
                         (
-                            N_MELS,
-                            MEL_STEP_SIZE
+                            80,
+                            16
                         ),
                         dtype=np.float32
                     )
@@ -1774,33 +1789,48 @@ def process_video(
                     )
 
 
+            elif (
+                mel_chunk.shape[1]
+                > MEL_STEP_SIZE
+            ):
+
+                mel_chunk = (
+                    mel_chunk[
+                        :,
+                        :MEL_STEP_SIZE
+                    ]
+                )
+
+
             # ------------------------------------------------
-            # Generate
+            # MODEL
             # ------------------------------------------------
 
-            generated = run_wav2lip(
+            generated = inference(
                 face,
                 mel_chunk
             )
 
 
             # ------------------------------------------------
-            # Mouth-only blend
+            # BLEND
             # ------------------------------------------------
 
-            result_face = blend_mouth(
-                face,
-                generated
+            result_face = (
+                blend_prediction(
+                    face,
+                    generated
+                )
             )
 
 
             # ------------------------------------------------
-            # Paste back
+            # PUT FACE BACK
             # ------------------------------------------------
 
             frame[
-                y:y+h,
-                x:x+w
+                y:y+ch,
+                x:x+cw
             ] = result_face
 
 
@@ -1810,7 +1840,7 @@ def process_video(
 
 
             # ------------------------------------------------
-            # Progress
+            # PROGRESS
             # ------------------------------------------------
 
             if (
@@ -1818,22 +1848,26 @@ def process_video(
                 == 0
             ):
 
-                ratio = (
+                fraction = (
                     frame_index
-                    / max(
+                    /
+                    max(
                         1,
-                        process_frames
+                        frames_to_process
                     )
                 )
 
 
                 progress(
                     0.18
-                    + ratio * 0.70,
+                    +
+                    fraction
+                    * 0.70,
+
                     desc=(
                         "Lip syncing "
                         f"{frame_index + 1}/"
-                        f"{process_frames}"
+                        f"{frames_to_process}"
                     )
                 )
 
@@ -1844,7 +1878,7 @@ def process_video(
 
 
         # ====================================================
-        # AUDIO
+        # ADD AUDIO
         # ====================================================
 
         progress(
@@ -1853,10 +1887,10 @@ def process_video(
         )
 
 
-        mux_audio(
+        add_audio(
             silent_video,
             audio_path,
-            output_video
+            output_path
         )
 
 
@@ -1866,46 +1900,47 @@ def process_video(
 
         progress(
             1.0,
-            desc="Finished!"
+            desc="Completed!"
         )
 
 
-        file_size = (
+        size_mb = (
             os.path.getsize(
-                output_video
-            )
-            / (
-                1024
-                * 1024
+                output_path
             )
+            /
+            1024
+            /
+            1024
         )
 
 
         status = (
             "## ✅ Lip-sync completed\n\n"
-            f"**Audio:** `{audio_duration:.2f}s`\n\n"
-            f"**FPS:** `{fps:.2f}`\n\n"
-            f"**Frames processed:** `{process_frames}`\n\n"
-            f"**Face tracking:** `Enabled`\n\n"
-            f"**Mouth-only blending:** `Enabled`\n\n"
-            f"**Output quality:** `CRF {FFMPEG_CRF}`\n\n"
-            f"**File size:** `{file_size:.2f} MB`"
+            f"**Video FPS:** `{fps:.2f}`\n\n"
+            f"**Audio duration:** `{audio_duration:.2f}s`\n\n"
+            f"**Processed frames:** `{frames_to_process}`\n\n"
+            "**Audio mapping:** `80 FPS mel mapping`\n\n"
+            "**Face tracking:** `Enabled`\n\n"
+            "**Soft mouth blending:** `Enabled`\n\n"
+            f"**Video quality:** `CRF {CRF}`\n\n"
+            f"**Output size:** `{size_mb:.2f} MB`"
         )
 
 
         if (
-            text
-            and text.strip()
+            text_input
+            and text_input.strip()
         ):
 
             status += (
                 "\n\n### 📝 Text reference\n\n"
-                + text.strip()
+                + text_input.strip()
             )
 
 
         return (
-            output_video,
+            output_path,
             status
         )
 
@@ -1936,7 +1971,7 @@ def process_video(
 
 
 # ============================================================
-# GRADIO
+# CSS
 # ============================================================
 
 CSS = """
@@ -1945,14 +1980,18 @@ CSS = """
     margin: auto !important;
 }
 
-#generate {
-    min-height: 58px;
+#generate-button {
+    min-height: 60px;
     font-size: 18px;
     font-weight: 700;
 }
 """
 
 
+# ============================================================
+# UI
+# ============================================================
+
 with gr.Blocks(
     title="Wav2Lip AI - High Quality Lip Sync"
 ) as demo:
@@ -1961,16 +2000,20 @@ with gr.Blocks(
         """
 # 🎬 Wav2Lip AI — High Quality Lip Sync
 
-### Upload a video + target audio → generate synchronized speech
+Upload a source video and target speech audio.
+The face will be tracked and the mouth synchronized
+to the target speech.
 """
     )
 
 
     gr.Markdown(
         """
-**Improved version:** better Wav2Lip audio preprocessing,
-more accurate frame-to-audio mapping, face tracking,
-and smooth mouth-only blending to reduce mouth distortion.
+### Improved pipeline
+
+Correct Wav2Lip-style audio preprocessing •
+accurate audio/video mapping • face tracking •
+soft mouth blending • high-quality H.264 output
 """
     )
 
@@ -1997,7 +2040,7 @@ and smooth mouth-only blending to reduce mouth distortion.
 
 
             text_input = gr.Textbox(
-                label="📝 Audio Text (Optional)",
+                label="📝 Audio Text — Optional",
                 placeholder=(
                     "Example: Hello, welcome to my video."
                 ),
@@ -2006,9 +2049,9 @@ and smooth mouth-only blending to reduce mouth distortion.
 
 
             generate_button = gr.Button(
-                "🎬 GENERATE HIGH-QUALITY LIP SYNC",
+                "🎬 GENERATE LIP-SYNC VIDEO",
                 variant="primary",
-                elem_id="generate"
+                elem_id="generate-button"
             )
 
 
@@ -2019,44 +2062,37 @@ and smooth mouth-only blending to reduce mouth distortion.
         with gr.Column():
 
             output_video = gr.Video(
-                label="🎬 Lip-Synced Result"
+                label="🎬 Lip-Synced Video"
             )
 
 
-            status_output = gr.Markdown(
-                label="Status"
-            )
+            status_output = gr.Markdown()
 
 
     # ========================================================
-    # TIPS
+    # INFO
     # ========================================================
 
     gr.Markdown(
         """
-### 💡 For the best mouth quality
-
-- Use **one person** in the video.
-- Face should be clearly visible.
-- Front-facing or slightly angled faces work best.
-- Avoid hands covering the mouth.
-- Avoid very small faces.
-- Use clean speech audio.
-- Keep the video reasonably stable.
-- The target audio should contain the speech you want the
-  person to perform.
-
-### ⚙️ Processing
-
-The system uses the uploaded **audio waveform**, not the
-text field, to generate lip movement. The text box is only
-provided as an optional reference.
+### 📌 Best input
+
+- One clearly visible face
+- Face looking toward the camera
+- Good lighting
+- Mouth not covered
+- 720p or higher source video recommended
+- Clean speech audio
+- Avoid very fast head movement
+
+**Note:** The text box is only a reference. The actual
+lip movement is generated from the uploaded audio waveform.
 """
     )
 
 
     # ========================================================
-    # BUTTON
+    # EVENT
     # ========================================================
 
     generate_button.click(
@@ -2082,7 +2118,7 @@ provided as an optional reference.
 if __name__ == "__main__":
 
     print("=" * 70)
-    print("Starting Wav2Lip High Quality Gradio Space")
+    print("Starting Wav2Lip Gradio Space")
     print("=" * 70)