Update inference2.py

inference2.py

@@ -1,346 +1,346 @@
-# inference.py (Updated)
-
-from os import listdir, path
-import numpy as np
-import scipy, cv2, os, sys, argparse, audio
-import json, subprocess, random, string
-from tqdm import tqdm
-from glob import glob
-import torch # Ensure torch is imported
-try:
-    import face_detection # Assuming this is installed or in a path accessible by your Flask app
-except ImportError:
-    print("face_detection not found. Please ensure it's installed or available in your PYTHONPATH.")
-    # You might want to raise an error or handle this gracefully if face_detection is truly optional.
-
-# Make sure you have a models/Wav2Lip.py or similar structure
-try:
-    from models import Wav2Lip
-except ImportError:
-    print("Wav2Lip model not found. Please ensure models/Wav2Lip.py exists and is correctly configured.")
-    # You might want to raise an error or handle this gracefully.
-
-import platform
-import shutil # For clearing temp directory
-
-
-# These globals are still useful for shared configuration
-mel_step_size = 16
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-print('Inference script using {} for inference.'.format(device))
-
-
-def get_smoothened_boxes(boxes, T):
-    for i in range(len(boxes)):
-        if i + T > len(boxes):
-            window = boxes[len(boxes) - T:]
-        else:
-            window = boxes[i : i + T]
-        boxes[i] = np.mean(window, axis=0)
-    return boxes
-
-def face_detect(images, pads, face_det_batch_size, nosmooth, img_size):
-    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
-                                            flip_input=False, device=device)
-
-    batch_size = face_det_batch_size
-
-    while 1:
-        predictions = []
-        try:
-            for i in tqdm(range(0, len(images), batch_size), desc="Face Detection"):
-                predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
-        except RuntimeError as e:
-            if batch_size == 1:
-                raise RuntimeError(f'Image too big to run face detection on GPU. Error: {e}')
-            batch_size //= 2
-            print('Recovering from OOM error; New face detection batch size: {}'.format(batch_size))
-            continue
-        break
-
-    results = []
-    pady1, pady2, padx1, padx2 = pads
-    for rect, image in zip(predictions, images):
-        if rect is None:
-            # Save the faulty frame for debugging
-            output_dir = 'temp' # Ensure this exists or create it
-            os.makedirs(output_dir, exist_ok=True)
-            cv2.imwrite(os.path.join(output_dir, 'faulty_frame.jpg'), image)
-            raise ValueError('Face not detected! Ensure the video/image contains a face in all the frames or try adjusting pads/box.')
-
-        y1 = max(0, rect[1] - pady1)
-        y2 = min(image.shape[0], rect[3] + pady2)
-        x1 = max(0, rect[0] - padx1)
-        x2 = min(image.shape[1], rect[2] + padx2)
-
-        results.append([x1, y1, x2, y2])
-
-    boxes = np.array(results)
-    if not nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
-    results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
-
-    del detector # Clean up detector
-    return results
-
-def datagen(frames, mels, box, static, wav2lip_batch_size, img_size, pads, face_det_batch_size, nosmooth):
-    img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
-
-    if box[0] == -1:
-        if not static:
-            face_det_results = face_detect(frames, pads, face_det_batch_size, nosmooth, img_size) # BGR2RGB for CNN face detection
-        else:
-            face_det_results = face_detect([frames[0]], pads, face_det_batch_size, nosmooth, img_size)
-    else:
-        print('Using the specified bounding box instead of face detection...')
-        y1, y2, x1, x2 = box
-        face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
-
-    for i, m in enumerate(mels):
-        idx = 0 if static else i % len(frames)
-        frame_to_save = frames[idx].copy()
-        face, coords = face_det_results[idx].copy()
-
-        face = cv2.resize(face, (img_size, img_size))
-
-        img_batch.append(face)
-        mel_batch.append(m)
-        frame_batch.append(frame_to_save)
-        coords_batch.append(coords)
-
-        if len(img_batch) >= wav2lip_batch_size:
-            img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
-
-            img_masked = img_batch.copy()
-            img_masked[:, img_size//2:] = 0
-
-            img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
-            mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
-
-            yield img_batch, mel_batch, frame_batch, coords_batch
-            img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
-
-    if len(img_batch) > 0:
-        img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
-
-        img_masked = img_batch.copy()
-        img_masked[:, img_size//2:] = 0
-
-        img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
-        mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
-
-        yield img_batch, mel_batch, frame_batch, coords_batch
-
-def _load(checkpoint_path):
-    # Use torch.jit.load for TorchScript archives
-    if device == 'cuda':
-        model = torch.jit.load(checkpoint_path)
-    else:
-        # Accepts string or torch.device, not a lambda
-        model = torch.jit.load(checkpoint_path, map_location='cpu')
-    return model
-
-def load_model(path):
-    print("Loading scripted model from:", path)
-    model = _load(path) # returns the TorchScript Module
-    model = model.to(device) # move to CPU or GPU
-    return model.eval() # set to eval() mode
-
-
-# New function to be called from Flask app
-def run_inference(
-    checkpoint_path: str,
-    face_path: str,
-    audio_path: str,
-    output_filename: str,
-    static: bool = False,
-    fps: float = 25.,
-    pads: list = [0, 10, 0, 0],
-    face_det_batch_size: int = 16,
-    wav2lip_batch_size: int = 128,
-    resize_factor: int = 1,
-    crop: list = [0, -1, 0, -1],
-    box: list = [-1, -1, -1, -1],
-    rotate: bool = False,
-    nosmooth: bool = False,
-    img_size: int = 96 # Fixed for Wav2Lip
-) -> str:
-    """
-    Runs the Wav2Lip inference process.
-
-    Args:
-        checkpoint_path (str): Path to the Wav2Lip model checkpoint.
-        face_path (str): Path to the input video/image file with a face.
-        audio_path (str): Path to the input audio file.
-        output_filename (str): Name of the output video file (e.g., 'result.mp4').
-        static (bool): If True, use only the first video frame for inference.
-        fps (float): Frames per second for static image input.
-        pads (list): Padding for face detection (top, bottom, left, right).
-        face_det_batch_size (int): Batch size for face detection.
-        wav2lip_batch_size (int): Batch size for Wav2Lip model(s).
-        resize_factor (int): Reduce the resolution by this factor.
-        crop (list): Crop video to a smaller region (top, bottom, left, right).
-        box (list): Constant bounding box for the face.
-        rotate (bool): Rotate video right by 90deg.
-        nosmooth (bool): Prevent smoothing face detections.
-        img_size (int): Image size for the model.
-
-    Returns:
-        str: The path to the generated output video file.
-    """
-    print(f"Starting inference with: face='{face_path}', audio='{audio_path}', checkpoint='{checkpoint_path}', outfile='{output_filename}'")
-
-    # Create necessary directories
-    output_dir = 'results'
-    temp_dir = 'temp'
-    os.makedirs(output_dir, exist_ok=True)
-    os.makedirs(temp_dir, exist_ok=True)
-
-    # Clear temp directory for fresh run
-    for item in os.listdir(temp_dir):
-        item_path = os.path.join(temp_dir, item)
-        if os.path.isfile(item_path):
-            os.remove(item_path)
-        elif os.path.isdir(item_path):
-            shutil.rmtree(item_path)
-
-    # Determine if input is static based on file extension
-    is_static_input = static or (os.path.isfile(face_path) and face_path.split('.')[-1].lower() in ['jpg', 'png', 'jpeg'])
-
-    full_frames = []
-    if is_static_input:
-        full_frames = [cv2.imread(face_path)]
-        if full_frames[0] is None:
-            raise ValueError(f"Could not read face image at: {face_path}")
-    else:
-        video_stream = cv2.VideoCapture(face_path)
-        if not video_stream.isOpened():
-            raise ValueError(f"Could not open video file at: {face_path}")
-        fps = video_stream.get(cv2.CAP_PROP_FPS)
-
-        print('Reading video frames...')
-        while 1:
-            still_reading, frame = video_stream.read()
-            if not still_reading:
-                video_stream.release()
-                break
-            if resize_factor > 1:
-                frame = cv2.resize(frame, (frame.shape[1]//resize_factor, frame.shape[0]//resize_factor))
-
-            if rotate:
-                frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
-
-            y1, y2, x1, x2 = crop
-            if x2 == -1: x2 = frame.shape[1]
-            if y2 == -1: y2 = frame.shape[0]
-
-            frame = frame[y1:y2, x1:x2]
-            full_frames.append(frame)
-
-    print ("Number of frames available for inference: "+str(len(full_frames)))
-    if not full_frames:
-        raise ValueError("No frames could be read from the input face file.")
-
-    temp_audio_path = os.path.join(temp_dir, 'temp_audio.wav')
-    if not audio_path.endswith('.wav'):
-        print('Extracting raw audio...')
-        command = f'ffmpeg -y -i "{audio_path}" -strict -2 "{temp_audio_path}"'
-        try:
-            subprocess.run(command, shell=True, check=True, capture_output=True)
-            audio_path = temp_audio_path
-        except subprocess.CalledProcessError as e:
-            print(f"FFmpeg error: {e.stderr.decode()}")
-            raise RuntimeError(f"Failed to extract audio from {audio_path}. Error: {e.stderr.decode()}")
-    else:
-        # Copy the wav file to temp if it's already wav to maintain consistency in naming
-        shutil.copy(audio_path, temp_audio_path)
-        audio_path = temp_audio_path
-
-
-    wav = audio.load_wav(audio_path, 16000)
-    mel = audio.melspectrogram(wav)
-    print("Mel spectrogram shape:", mel.shape)
-
-    if np.isnan(mel.reshape(-1)).sum() > 0:
-        raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
-
-    mel_chunks = []
-    mel_idx_multiplier = 80./fps
-    i = 0
-    while 1:
-        start_idx = int(i * mel_idx_multiplier)
-        if start_idx + mel_step_size > len(mel[0]):
-            mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
-            break
-        mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
-        i += 1
-
-    print("Length of mel chunks: {}".format(len(mel_chunks)))
-
-    # Ensure full_frames matches mel_chunks length, or loop if static
-    if not is_static_input:
-        full_frames = full_frames[:len(mel_chunks)]
-    else:
-        # If static, replicate the first frame for the duration of the audio
-        full_frames = [full_frames[0]] * len(mel_chunks)
-
-
-    gen = datagen(full_frames.copy(), mel_chunks, box, is_static_input, wav2lip_batch_size, img_size, pads, face_det_batch_size, nosmooth)
-
-    output_avi_path = os.path.join(temp_dir, 'result.avi')
-
-    model_loaded = False
-    model = None
-    frame_h, frame_w = 0, 0
-    out = None
-
-    for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, desc="Wav2Lip Inference",
-                                            total=int(np.ceil(float(len(mel_chunks))/wav2lip_batch_size)))):
-        if not model_loaded:
-            model = load_model(checkpoint_path)
-            model_loaded = True
-            print ("Model loaded successfully")
-
-            frame_h, frame_w = full_frames[0].shape[:-1]
-            out = cv2.VideoWriter(output_avi_path,
-                                    cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
-        if out is None: # In case no frames were generated for some reason
-            raise RuntimeError("Video writer could not be initialized.")
-
-
-        img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
-        mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
-
-        with torch.no_grad():
-            pred = model(mel_batch, img_batch)
-
-        pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
-
-        for p, f, c in zip(pred, frames, coords):
-            y1, y2, x1, x2 = c
-            p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
-
-            f[y1:y2, x1:x2] = p
-            out.write(f)
-
-    if out:
-        out.release()
-    else:
-        print("Warning: Video writer was not initialized or no frames were processed.")
-
-
-    final_output_path = os.path.join(output_dir, output_filename)
-    command = f'ffmpeg -y -i "{audio_path}" -i "{output_avi_path}" -strict -2 -q:v 1 "{final_output_path}"'
-
-    try:
-        subprocess.run(command, shell=True, check=True, capture_output=True)
-        print(f"Output saved to: {final_output_path}")
-    except subprocess.CalledProcessError as e:
-        print(f"FFmpeg final merge error: {e.stderr.decode()}")
-        raise RuntimeError(f"Failed to merge audio and video. Error: {e.stderr.decode()}")
-
-    # Clean up temporary files (optional, but good practice)
-    # shutil.rmtree(temp_dir) # Be careful with this if you want to inspect temp files
-
-    return final_output_path
-
+# inference.py (Updated)
+import audio
+from os import listdir, path
+import numpy as np
+import scipy, cv2, os, sys, argparse, audio
+import json, subprocess, random, string
+from tqdm import tqdm
+from glob import glob
+import torch # Ensure torch is imported
+try:
+    import face_detection # Assuming this is installed or in a path accessible by your Flask app
+except ImportError:
+    print("face_detection not found. Please ensure it's installed or available in your PYTHONPATH.")
+    # You might want to raise an error or handle this gracefully if face_detection is truly optional.
+
+# Make sure you have a models/Wav2Lip.py or similar structure
+try:
+    from models import Wav2Lip
+except ImportError:
+    print("Wav2Lip model not found. Please ensure models/Wav2Lip.py exists and is correctly configured.")
+    # You might want to raise an error or handle this gracefully.
+
+import platform
+import shutil # For clearing temp directory
+
+
+# These globals are still useful for shared configuration
+mel_step_size = 16
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print('Inference script using {} for inference.'.format(device))
+
+
+def get_smoothened_boxes(boxes, T):
+    for i in range(len(boxes)):
+        if i + T > len(boxes):
+            window = boxes[len(boxes) - T:]
+        else:
+            window = boxes[i : i + T]
+        boxes[i] = np.mean(window, axis=0)
+    return boxes
+
+def face_detect(images, pads, face_det_batch_size, nosmooth, img_size):
+    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
+                                            flip_input=False, device=device)
+
+    batch_size = face_det_batch_size
+
+    while 1:
+        predictions = []
+        try:
+            for i in tqdm(range(0, len(images), batch_size), desc="Face Detection"):
+                predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
+        except RuntimeError as e:
+            if batch_size == 1:
+                raise RuntimeError(f'Image too big to run face detection on GPU. Error: {e}')
+            batch_size //= 2
+            print('Recovering from OOM error; New face detection batch size: {}'.format(batch_size))
+            continue
+        break
+
+    results = []
+    pady1, pady2, padx1, padx2 = pads
+    for rect, image in zip(predictions, images):
+        if rect is None:
+            # Save the faulty frame for debugging
+            output_dir = 'temp' # Ensure this exists or create it
+            os.makedirs(output_dir, exist_ok=True)
+            cv2.imwrite(os.path.join(output_dir, 'faulty_frame.jpg'), image)
+            raise ValueError('Face not detected! Ensure the video/image contains a face in all the frames or try adjusting pads/box.')
+
+        y1 = max(0, rect[1] - pady1)
+        y2 = min(image.shape[0], rect[3] + pady2)
+        x1 = max(0, rect[0] - padx1)
+        x2 = min(image.shape[1], rect[2] + padx2)
+
+        results.append([x1, y1, x2, y2])
+
+    boxes = np.array(results)
+    if not nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
+    results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+
+    del detector # Clean up detector
+    return results
+
+def datagen(frames, mels, box, static, wav2lip_batch_size, img_size, pads, face_det_batch_size, nosmooth):
+    img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
+
+    if box[0] == -1:
+        if not static:
+            face_det_results = face_detect(frames, pads, face_det_batch_size, nosmooth, img_size) # BGR2RGB for CNN face detection
+        else:
+            face_det_results = face_detect([frames[0]], pads, face_det_batch_size, nosmooth, img_size)
+    else:
+        print('Using the specified bounding box instead of face detection...')
+        y1, y2, x1, x2 = box
+        face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
+
+    for i, m in enumerate(mels):
+        idx = 0 if static else i % len(frames)
+        frame_to_save = frames[idx].copy()
+        face, coords = face_det_results[idx].copy()
+
+        face = cv2.resize(face, (img_size, img_size))
+
+        img_batch.append(face)
+        mel_batch.append(m)
+        frame_batch.append(frame_to_save)
+        coords_batch.append(coords)
+
+        if len(img_batch) >= wav2lip_batch_size:
+            img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
+
+            img_masked = img_batch.copy()
+            img_masked[:, img_size//2:] = 0
+
+            img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
+            mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
+
+            yield img_batch, mel_batch, frame_batch, coords_batch
+            img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
+
+    if len(img_batch) > 0:
+        img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
+
+        img_masked = img_batch.copy()
+        img_masked[:, img_size//2:] = 0
+
+        img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
+        mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
+
+        yield img_batch, mel_batch, frame_batch, coords_batch
+
+def _load(checkpoint_path):
+    # Use torch.jit.load for TorchScript archives
+    if device == 'cuda':
+        model = torch.jit.load(checkpoint_path)
+    else:
+        # Accepts string or torch.device, not a lambda
+        model = torch.jit.load(checkpoint_path, map_location='cpu')
+    return model
+
+def load_model(path):
+    print("Loading scripted model from:", path)
+    model = _load(path) # returns the TorchScript Module
+    model = model.to(device) # move to CPU or GPU
+    return model.eval() # set to eval() mode
+
+
+# New function to be called from Flask app
+def run_inference(
+    checkpoint_path: str,
+    face_path: str,
+    audio_path: str,
+    output_filename: str,
+    static: bool = False,
+    fps: float = 25.,
+    pads: list = [0, 10, 0, 0],
+    face_det_batch_size: int = 16,
+    wav2lip_batch_size: int = 128,
+    resize_factor: int = 1,
+    crop: list = [0, -1, 0, -1],
+    box: list = [-1, -1, -1, -1],
+    rotate: bool = False,
+    nosmooth: bool = False,
+    img_size: int = 96 # Fixed for Wav2Lip
+) -> str:
+    """
+    Runs the Wav2Lip inference process.
+
+    Args:
+        checkpoint_path (str): Path to the Wav2Lip model checkpoint.
+        face_path (str): Path to the input video/image file with a face.
+        audio_path (str): Path to the input audio file.
+        output_filename (str): Name of the output video file (e.g., 'result.mp4').
+        static (bool): If True, use only the first video frame for inference.
+        fps (float): Frames per second for static image input.
+        pads (list): Padding for face detection (top, bottom, left, right).
+        face_det_batch_size (int): Batch size for face detection.
+        wav2lip_batch_size (int): Batch size for Wav2Lip model(s).
+        resize_factor (int): Reduce the resolution by this factor.
+        crop (list): Crop video to a smaller region (top, bottom, left, right).
+        box (list): Constant bounding box for the face.
+        rotate (bool): Rotate video right by 90deg.
+        nosmooth (bool): Prevent smoothing face detections.
+        img_size (int): Image size for the model.
+
+    Returns:
+        str: The path to the generated output video file.
+    """
+    print(f"Starting inference with: face='{face_path}', audio='{audio_path}', checkpoint='{checkpoint_path}', outfile='{output_filename}'")
+
+    # Create necessary directories
+    output_dir = 'results'
+    temp_dir = 'temp'
+    os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(temp_dir, exist_ok=True)
+
+    # Clear temp directory for fresh run
+    for item in os.listdir(temp_dir):
+        item_path = os.path.join(temp_dir, item)
+        if os.path.isfile(item_path):
+            os.remove(item_path)
+        elif os.path.isdir(item_path):
+            shutil.rmtree(item_path)
+
+    # Determine if input is static based on file extension
+    is_static_input = static or (os.path.isfile(face_path) and face_path.split('.')[-1].lower() in ['jpg', 'png', 'jpeg'])
+
+    full_frames = []
+    if is_static_input:
+        full_frames = [cv2.imread(face_path)]
+        if full_frames[0] is None:
+            raise ValueError(f"Could not read face image at: {face_path}")
+    else:
+        video_stream = cv2.VideoCapture(face_path)
+        if not video_stream.isOpened():
+            raise ValueError(f"Could not open video file at: {face_path}")
+        fps = video_stream.get(cv2.CAP_PROP_FPS)
+
+        print('Reading video frames...')
+        while 1:
+            still_reading, frame = video_stream.read()
+            if not still_reading:
+                video_stream.release()
+                break
+            if resize_factor > 1:
+                frame = cv2.resize(frame, (frame.shape[1]//resize_factor, frame.shape[0]//resize_factor))
+
+            if rotate:
+                frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
+
+            y1, y2, x1, x2 = crop
+            if x2 == -1: x2 = frame.shape[1]
+            if y2 == -1: y2 = frame.shape[0]
+
+            frame = frame[y1:y2, x1:x2]
+            full_frames.append(frame)
+
+    print ("Number of frames available for inference: "+str(len(full_frames)))
+    if not full_frames:
+        raise ValueError("No frames could be read from the input face file.")
+
+    temp_audio_path = os.path.join(temp_dir, 'temp_audio.wav')
+    if not audio_path.endswith('.wav'):
+        print('Extracting raw audio...')
+        command = f'ffmpeg -y -i "{audio_path}" -strict -2 "{temp_audio_path}"'
+        try:
+            subprocess.run(command, shell=True, check=True, capture_output=True)
+            audio_path = temp_audio_path
+        except subprocess.CalledProcessError as e:
+            print(f"FFmpeg error: {e.stderr.decode()}")
+            raise RuntimeError(f"Failed to extract audio from {audio_path}. Error: {e.stderr.decode()}")
+    else:
+        # Copy the wav file to temp if it's already wav to maintain consistency in naming
+        shutil.copy(audio_path, temp_audio_path)
+        audio_path = temp_audio_path
+
+
+    wav = audio.load_wav(audio_path, 16000)
+    mel = audio.melspectrogram(wav)
+    print("Mel spectrogram shape:", mel.shape)
+
+    if np.isnan(mel.reshape(-1)).sum() > 0:
+        raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
+
+    mel_chunks = []
+    mel_idx_multiplier = 80./fps
+    i = 0
+    while 1:
+        start_idx = int(i * mel_idx_multiplier)
+        if start_idx + mel_step_size > len(mel[0]):
+            mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
+            break
+        mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
+        i += 1
+
+    print("Length of mel chunks: {}".format(len(mel_chunks)))
+
+    # Ensure full_frames matches mel_chunks length, or loop if static
+    if not is_static_input:
+        full_frames = full_frames[:len(mel_chunks)]
+    else:
+        # If static, replicate the first frame for the duration of the audio
+        full_frames = [full_frames[0]] * len(mel_chunks)
+
+
+    gen = datagen(full_frames.copy(), mel_chunks, box, is_static_input, wav2lip_batch_size, img_size, pads, face_det_batch_size, nosmooth)
+
+    output_avi_path = os.path.join(temp_dir, 'result.avi')
+
+    model_loaded = False
+    model = None
+    frame_h, frame_w = 0, 0
+    out = None
+
+    for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, desc="Wav2Lip Inference",
+                                            total=int(np.ceil(float(len(mel_chunks))/wav2lip_batch_size)))):
+        if not model_loaded:
+            model = load_model(checkpoint_path)
+            model_loaded = True
+            print ("Model loaded successfully")
+
+            frame_h, frame_w = full_frames[0].shape[:-1]
+            out = cv2.VideoWriter(output_avi_path,
+                                    cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
+        if out is None: # In case no frames were generated for some reason
+            raise RuntimeError("Video writer could not be initialized.")
+
+
+        img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
+        mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
+
+        with torch.no_grad():
+            pred = model(mel_batch, img_batch)
+
+        pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
+
+        for p, f, c in zip(pred, frames, coords):
+            y1, y2, x1, x2 = c
+            p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
+
+            f[y1:y2, x1:x2] = p
+            out.write(f)
+
+    if out:
+        out.release()
+    else:
+        print("Warning: Video writer was not initialized or no frames were processed.")
+
+
+    final_output_path = os.path.join(output_dir, output_filename)
+    command = f'ffmpeg -y -i "{audio_path}" -i "{output_avi_path}" -strict -2 -q:v 1 "{final_output_path}"'
+
+    try:
+        subprocess.run(command, shell=True, check=True, capture_output=True)
+        print(f"Output saved to: {final_output_path}")
+    except subprocess.CalledProcessError as e:
+        print(f"FFmpeg final merge error: {e.stderr.decode()}")
+        raise RuntimeError(f"Failed to merge audio and video. Error: {e.stderr.decode()}")
+
+    # Clean up temporary files (optional, but good practice)
+    # shutil.rmtree(temp_dir) # Be careful with this if you want to inspect temp files
+
+    return final_output_path
+
 # No `if __name__ == '__main__':` block here, as it's meant to be imported