Initial Commit

handler.py

@@ -1,470 +1,559 @@
-import torch
+# Hugging Face Inference Endpoint Handler for MultiTalk
 import os
-import shutil
 import subprocess
+import time
 import json
 import tempfile
+import logging
+import sys
+import warnings
+import shutil
 import base64
-from typing import Dict, Any, Optional, List
-import soundfile as sf
+from typing import Optional, Dict, Any, Union
 from datetime import datetime
-from huggingface_hub import snapshot_download
-
-
-class EndpointHandler:
-    def __init__(self, path=""):
-        """
-        Initialize the handler for Meigen MultiTalk inference.
-        Optimized for serverless deployment with pre-downloaded models.
-
-        Args:
-            path: Path to the model directory (not used in this implementation)
-        """
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+from types import SimpleNamespace
+from io import BytesIO
 
-        # Set up the script path - try multiple possible locations
-        self.script_path = self.find_script_path()
-        if not self.script_path:
-            raise RuntimeError("Could not locate generate_multitalk.py script")
+# Suppress warnings for cleaner output
+warnings.filterwarnings("ignore")
 
-        print(f"Found script at: {self.script_path}")
-
-        self.setup_models()
-        self.setup_gpu_config()
-
-    def find_script_path(self) -> Optional[str]:
-        """
-        Find the location of the generate_multitalk.py script.
-
-        Returns:
-            Path to the script if found, None otherwise
-        """
-        possible_paths = [
-            "generate_multitalk.py",  # Current directory
-            "/app/generate_multitalk.py",  # App directory
-            "./src/generate_multitalk.py",  # Src directory
-            "../generate_multitalk.py",  # Parent directory
-            "./weights/MeiGen-MultiTalk/generate_multitalk.py",  # In model directory
-            "./MeiGen-MultiTalk/generate_multitalk.py",  # Alternative model directory
-        ]
-
-        for path in possible_paths:
-            if os.path.exists(path):
-                return os.path.abspath(path)
-
-        # Try to find it recursively in common directories
-        search_dirs = ["/app", ".", "./weights", "./src"]
-        for search_dir in search_dirs:
-            if os.path.exists(search_dir):
-                for root, dirs, files in os.walk(search_dir):
-                    if "generate_multitalk.py" in files:
-                        return os.path.join(root, "generate_multitalk.py")
-
-        return None
-
-    def setup_models(self):
-        """Setup models - download only if not already present."""
-        print("Setting up models...")
-
-        # Define model paths
-        self.wan_model_path = "./weights/Wan2.1-I2V-14B-480P"
-        self.wav2vec_path = "./weights/chinese-wav2vec2-base"
-        self.multitalk_path = "./weights/MeiGen-MultiTalk"
-
-        # Only download if models don't exist (for local development)
-        # In serverless deployment, these should be pre-baked into the container
-        if not os.path.exists(self.wan_model_path):
-            print("Downloading Wan2.1-I2V-14B-480P model...")
-            self.wan_model_path = snapshot_download(
-                repo_id="Wan-AI/Wan2.1-I2V-14B-480P",
-                local_dir=self.wan_model_path,
-            )
-
-        if not os.path.exists(self.wav2vec_path):
-            print("Downloading Chinese wav2vec2 model...")
-            self.wav2vec_path = snapshot_download(
-                repo_id="TencentGameMate/chinese-wav2vec2-base",
-                local_dir=self.wav2vec_path,
-            )
-
-        if not os.path.exists(self.multitalk_path):
-            print("Downloading MeiGen MultiTalk weights...")
-            self.multitalk_path = snapshot_download(
-                repo_id="MeiGen-AI/MeiGen-MultiTalk",
-                local_dir=self.multitalk_path,
-            )
-
-        # Setup model files if not already done
-        self.setup_model_files()
-
-    def setup_model_files(self):
-        """Setup model files by copying and renaming as needed."""
-        base_model_dir = self.wan_model_path
-        multitalk_dir = self.multitalk_path
-
-        # Check if setup is already done
-        multitalk_weights_path = os.path.join(base_model_dir, "multitalk.safetensors")
-        if os.path.exists(multitalk_weights_path):
-            print("Model files already set up, skipping setup.")
-            return
-
-        # File to rename
-        original_index = os.path.join(
-            base_model_dir, "diffusion_pytorch_model.safetensors.index.json"
+import torch
+import numpy as np
+import random
+import soundfile as sf
+from PIL import Image
+
+# Import MultiTalk components
+import wan
+from wan.configs import WAN_CONFIGS
+from transformers import Wav2Vec2FeatureExtractor
+from src.audio_analysis.wav2vec2 import Wav2Vec2Model
+import librosa
+import pyloudnorm as pyln
+from einops import rearrange
+from wan.utils.multitalk_utils import save_video_ffmpeg
+
+
+def loudness_norm(audio_array, sr=16000, lufs=-23):
+    meter = pyln.Meter(sr)
+    loudness = meter.integrated_loudness(audio_array)
+    if abs(loudness) > 100:
+        return audio_array
+    normalized_audio = pyln.normalize.loudness(audio_array, loudness, lufs)
+    return normalized_audio
+
+
+def extract_audio_from_video(filename, sample_rate=16000):
+    """Extract audio from video file with robust error handling"""
+    raw_audio_path = f"{os.path.splitext(os.path.basename(filename))[0]}.wav"
+    ffmpeg_command = [
+        "ffmpeg",
+        "-y",
+        "-i",
+        str(filename),
+        "-vn",
+        "-acodec",
+        "pcm_s16le",
+        "-ar",
+        str(sample_rate),
+        "-ac",
+        "2",
+        raw_audio_path,
+    ]
+    subprocess.run(ffmpeg_command, check=True, capture_output=True)
+    human_speech_array, sr = librosa.load(raw_audio_path, sr=sample_rate)
+    human_speech_array = loudness_norm(human_speech_array, sr)
+    os.remove(raw_audio_path)
+    return human_speech_array
+
+
+def audio_prepare_single(audio_path, sample_rate=16000):
+    ext = os.path.splitext(audio_path)[1].lower()
+    if ext in [".mp4", ".mov", ".avi", ".mkv"]:
+        human_speech_array = extract_audio_from_video(audio_path, sample_rate)
+        return human_speech_array
+    else:
+        human_speech_array, sr = librosa.load(audio_path, sr=sample_rate)
+        human_speech_array = loudness_norm(human_speech_array, sr)
+        return human_speech_array
+
+
+def audio_prepare_multi(left_path, right_path, audio_type, sample_rate=16000):
+    human_speech_array1 = audio_prepare_single(left_path)
+    human_speech_array2 = audio_prepare_single(right_path)
+
+    if audio_type == "para":
+        new_human_speech1 = human_speech_array1
+        new_human_speech2 = human_speech_array2
+    elif audio_type == "add":
+        new_human_speech1 = np.concatenate(
+            [
+                human_speech_array1[: human_speech_array1.shape[0]],
+                np.zeros(human_speech_array2.shape[0]),
+            ]
         )
-        backup_index = os.path.join(
-            base_model_dir, "diffusion_pytorch_model.safetensors.index.json_old"
+        new_human_speech2 = np.concatenate(
+            [
+                np.zeros(human_speech_array1.shape[0]),
+                human_speech_array2[: human_speech_array2.shape[0]],
+            ]
         )
-
-        # Rename the original index file
-        if os.path.exists(original_index) and not os.path.exists(backup_index):
-            os.rename(original_index, backup_index)
-            print("Renamed original index file to .json_old")
-
-        # Copy updated index file from MultiTalk
-        multitalk_index = os.path.join(
-            multitalk_dir, "diffusion_pytorch_model.safetensors.index.json"
+    sum_human_speechs = new_human_speech1 + new_human_speech2
+    return new_human_speech1, new_human_speech2, sum_human_speechs
+
+
+def get_embedding(
+    speech_array, wav2vec_feature_extractor, audio_encoder, sr=16000, device="cpu"
+):
+    """Extract audio embeddings optimized for GPU processing"""
+    audio_duration = len(speech_array) / sr
+    video_length = audio_duration * 25  # Assume the video fps is 25
+
+    # Extract audio features
+    audio_feature = np.squeeze(
+        wav2vec_feature_extractor(speech_array, sampling_rate=sr).input_values
+    )
+    audio_feature = torch.from_numpy(audio_feature).float().to(device=device)
+    audio_feature = audio_feature.unsqueeze(0)
+
+    # Generate embeddings on appropriate device
+    with torch.no_grad():
+        embeddings = audio_encoder(
+            audio_feature, seq_len=int(video_length), output_hidden_states=True
         )
-        if os.path.exists(multitalk_index):
-            shutil.copy2(multitalk_index, base_model_dir)
-
-        # Copy MultiTalk model weights
-        multitalk_weights_src = os.path.join(multitalk_dir, "multitalk.safetensors")
-        if os.path.exists(multitalk_weights_src):
-            shutil.copy2(multitalk_weights_src, base_model_dir)
-
-        print("Copied MultiTalk files into base model directory.")
-
-    def setup_gpu_config(self):
-        """Setup GPU configuration and VRAM parameters."""
-        self.num_gpus = torch.cuda.device_count()
-        print(f"GPU AVAILABLE: {self.num_gpus}")
-
-        if torch.cuda.is_available():
-            self.gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
-            print(f"Current GPU: {self.gpu_name}")
-
-            # GPU to VRAM mapping
-            GPU_TO_VRAM_PARAMS = {
-                "NVIDIA A100": 11000000000,
-                "NVIDIA A100-SXM4-40GB": 11000000000,
-                "NVIDIA A100-SXM4-80GB": 22000000000,
-                "NVIDIA L4": 5000000000,
-                "NVIDIA L40S": 11000000000,
-            }
-
-            self.used_vram_params = GPU_TO_VRAM_PARAMS.get(self.gpu_name, 5000000000)
-            print(f"Using {self.used_vram_params} for num_persistent_param_in_dit")
-        else:
-            raise RuntimeError("No CUDA-compatible GPU found. A GPU is required.")
 
-    def trim_audio_to_5s_temp(self, audio_path: str, sample_rate: int = 16000) -> str:
-        """
-        Trim audio to 5 seconds and save to temporary file.
-
-        Args:
-            audio_path: Path to input audio file
-            sample_rate: Target sample rate
+    if len(embeddings) == 0:
+        print("⚠️ Failed to extract audio embedding")
+        return None
 
-        Returns:
-            Path to trimmed audio file
-        """
-        max_duration_sec = 5
-        audio, sr = sf.read(audio_path)
+    audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
+    audio_emb = rearrange(audio_emb, "b s d -> s b d")
+
+    # Keep on CPU for compatibility with downstream processing
+    audio_emb = audio_emb.cpu().detach()
+    return audio_emb
+
+
+def download_weights(url: str, dest: str) -> None:
+    start = time.time()
+    print("[!] Initiating download from URL: ", url)
+    print("[~] Destination path: ", dest)
+    if ".tar" in dest:
+        dest = os.path.dirname(dest)
+    command = ["pget", "-vf" + ("x" if ".tar" in url else ""), url, dest]
+    try:
+        print(f"[~] Running command: {' '.join(command)}")
+        subprocess.check_call(command, close_fds=False)
+    except subprocess.CalledProcessError as e:
+        print(
+            f"[ERROR] Failed to download weights. Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}."
+        )
+        raise
+    print("[+] Download completed in: ", time.time() - start, "seconds")
 
-        if sr != sample_rate:
-            sample_rate = sr
 
-        max_samples = max_duration_sec * sample_rate
-        if len(audio) > max_samples:
-            audio = audio[:max_samples]
+def decode_base64_file(base64_string: str, output_path: str) -> str:
+    """Decode base64 string and save to file"""
+    try:
+        # Remove data URL prefix if present
+        if "," in base64_string:
+            base64_string = base64_string.split(",")[1]
 
-        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
-        base_name = os.path.splitext(os.path.basename(audio_path))[0]
-        temp_filename = f"{base_name}_trimmed_{timestamp}.wav"
-        temp_path = os.path.join(tempfile.gettempdir(), temp_filename)
+        # Decode and save
+        with open(output_path, "wb") as f:
+            f.write(base64.b64decode(base64_string))
 
-        sf.write(temp_path, audio, samplerate=sample_rate)
-        return temp_path
+        return output_path
+    except Exception as e:
+        raise ValueError(f"Failed to decode base64 file: {str(e)}")
 
-    def save_base64_file(self, base64_data: str, file_extension: str) -> str:
-        """
-        Save base64 encoded file to temporary location.
 
-        Args:
-            base64_data: Base64 encoded file data
-            file_extension: File extension (e.g., '.png', '.wav')
+def encode_file_to_base64(file_path: str) -> str:
+    """Encode file to base64 string"""
+    try:
+        with open(file_path, "rb") as f:
+            return base64.b64encode(f.read()).decode("utf-8")
+    except Exception as e:
+        raise ValueError(f"Failed to encode file to base64: {str(e)}")
 
-        Returns:
-            Path to saved file
-        """
-        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
-        temp_filename = f"temp_file_{timestamp}{file_extension}"
-        temp_path = os.path.join(tempfile.gettempdir(), temp_filename)
-
-        with open(temp_path, "wb") as f:
-            f.write(base64.b64decode(base64_data))
-
-        return temp_path
-
-    def create_temp_input_json(
-        self,
-        prompt: str,
-        cond_image_path: str,
-        cond_audio_path_spk1: str,
-        cond_audio_path_spk2: Optional[str] = None,
-    ) -> str:
-        """
-        Create a temporary JSON file with the user-provided prompt, image, and audio paths.
 
-        Args:
-            prompt: Text prompt for generation
-            cond_image_path: Path to conditioning image
-            cond_audio_path_spk1: Path to first speaker audio
-            cond_audio_path_spk2: Path to second speaker audio (optional)
+class EndpointHandler:
+    def __init__(self, path=""):
+        """Initialize the handler"""
+        # Model configuration
+        self.MODEL_CACHE = "weights"
+        self.BASE_URL = (
+            f"https://weights.replicate.delivery/default/multitalk/{self.MODEL_CACHE}/"
+        )
 
-        Returns:
-            Path to temporary JSON file
-        """
-        if cond_audio_path_spk2 is None:
-            data = {
-                "prompt": prompt,
-                "cond_image": cond_image_path,
-                "cond_audio": {"person1": cond_audio_path_spk1},
-            }
-        else:
-            data = {
-                "prompt": prompt,
-                "cond_image": cond_image_path,
-                "audio_type": "para",
-                "cond_audio": {
-                    "person1": cond_audio_path_spk1,
-                    "person2": cond_audio_path_spk2,
-                },
-            }
+        # Set environment variables
+        os.environ["HF_HOME"] = self.MODEL_CACHE
+        os.environ["TORCH_HOME"] = self.MODEL_CACHE
+        os.environ["HF_DATASETS_CACHE"] = self.MODEL_CACHE
+        os.environ["TRANSFORMERS_CACHE"] = self.MODEL_CACHE
+        os.environ["HUGGINGFACE_HUB_CACHE"] = self.MODEL_CACHE
+
+        # Create model cache directory if it doesn't exist
+        os.makedirs(self.MODEL_CACHE, exist_ok=True)
+
+        # Download model weights if needed
+        model_files = [
+            "MeiGen-MultiTalk.tar",
+            "Wan2.1-I2V-14B-480P.tar",
+            "chinese-wav2vec2-base.tar",
+        ]
 
-        temp_json = tempfile.NamedTemporaryFile(
-            delete=False, suffix=".json", mode="w", encoding="utf-8"
+        for model_file in model_files:
+            url = self.BASE_URL + model_file
+            filename = url.split("/")[-1]
+            dest_path = os.path.join(self.MODEL_CACHE, filename)
+            if not os.path.exists(dest_path.replace(".tar", "")):
+                download_weights(url, dest_path)
+
+        # Set up logging
+        logging.basicConfig(
+            level=logging.INFO,
+            format="[%(asctime)s] %(levelname)s: %(message)s",
+            handlers=[logging.StreamHandler(stream=sys.stdout)],
         )
-        json.dump(data, temp_json, indent=4)
-        temp_json_path = temp_json.name
-        temp_json.close()
-
-        print(f"Temporary input JSON saved to: {temp_json_path}")
-        return temp_json_path
-
-    def run_inference(
-        self,
-        prompt: str,
-        cond_image_path: str,
-        cond_audio_path_spk1: str,
-        cond_audio_path_spk2: Optional[str] = None,
-        sample_steps: int = 12,
-    ) -> str:
-        """
-        Run the MultiTalk inference.
 
-        Args:
-            prompt: Text prompt
-            cond_image_path: Path to conditioning image
-            cond_audio_path_spk1: Path to first speaker audio
-            cond_audio_path_spk2: Path to second speaker audio (optional)
-            sample_steps: Number of sampling steps
+        # Model paths
+        self.ckpt_dir = "weights/Wan2.1-I2V-14B-480P"
+        self.wav2vec_dir = "weights/chinese-wav2vec2-base"
+        self.multitalk_dir = "weights/MeiGen-MultiTalk"
+
+        # Initialize device for single GPU
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {self.device}")
+
+        # Load wav2vec models optimized for high VRAM
+        print("Loading wav2vec models...")
+        audio_device = (
+            self.device
+            if torch.cuda.is_available()
+            and torch.cuda.get_device_properties(0).total_memory > 40 * 1024**3
+            else "cpu"
+        )
+        print(f"Loading audio encoder on: {audio_device}")
+
+        self.audio_encoder = Wav2Vec2Model.from_pretrained(
+            self.wav2vec_dir, local_files_only=True, attn_implementation="eager"
+        ).to(audio_device)
+        self.audio_encoder.feature_extractor._freeze_parameters()
+        self.wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
+            self.wav2vec_dir, local_files_only=True
+        )
+        self.audio_device = audio_device
+
+        # Load MultiTalk pipeline
+        print("Loading MultiTalk pipeline...")
+        self.cfg = WAN_CONFIGS["multitalk-14B"]
+        self.wan_i2v = wan.MultiTalkPipeline(
+            config=self.cfg,
+            checkpoint_dir=self.ckpt_dir,
+            device_id=0,
+            rank=0,
+            t5_fsdp=False,
+            dit_fsdp=False,
+            use_usp=False,
+            t5_cpu=False,  # Keep T5 on GPU for speed
+        )
 
-        Returns:
-            Path to generated video
-        """
-        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
-        result_filename = f"meigen_multitalk_result_{sample_steps}_steps_{timestamp}"
-        temp_files_to_cleanup = []
+        # GPU optimizations for high-VRAM setup (A100/H100/H200)
+        if torch.cuda.is_available():
+            vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+            print(f"🔍 Detected {vram_gb:.1f}GB VRAM")
 
-        try:
-            # Trim audio files to 5 seconds
-            trimmed_audio_path_spk1 = self.trim_audio_to_5s_temp(cond_audio_path_spk1)
-            if trimmed_audio_path_spk1 != cond_audio_path_spk1:
-                cond_audio_path_spk1 = trimmed_audio_path_spk1
-                temp_files_to_cleanup.append(trimmed_audio_path_spk1)
-
-            if cond_audio_path_spk2 is not None:
-                trimmed_audio_path_spk2 = self.trim_audio_to_5s_temp(
-                    cond_audio_path_spk2
+            if vram_gb > 40:  # High VRAM setup
+                print(
+                    "🚀 High-VRAM detected: Enabling maximum performance optimizations"
                 )
-                if trimmed_audio_path_spk2 != cond_audio_path_spk2:
-                    cond_audio_path_spk2 = trimmed_audio_path_spk2
-                    temp_files_to_cleanup.append(trimmed_audio_path_spk2)
-
-            # Prepare input JSON
-            input_json_path = self.create_temp_input_json(
-                prompt, cond_image_path, cond_audio_path_spk1, cond_audio_path_spk2
-            )
-            temp_files_to_cleanup.append(input_json_path)
-
-            # Base command arguments
-            common_args = [
-                "--ckpt_dir",
-                self.wan_model_path,
-                "--wav2vec_dir",
-                self.wav2vec_path,
-                "--input_json",
-                input_json_path,
-                "--sample_steps",
-                str(sample_steps),
-                "--mode",
-                "streaming",
-                "--use_teacache",
-                "--save_file",
-                result_filename,
-            ]
-
-            # Get the directory containing the script to set as working directory
-            script_dir = os.path.dirname(self.script_path)
-            script_name = os.path.basename(self.script_path)
-
-            # Choose command based on number of GPUs
-            if self.num_gpus > 1:
-                cmd = [
-                    "torchrun",
-                    f"--nproc_per_node={self.num_gpus}",
-                    "--standalone",
-                    script_name,  # Use just the script name since we'll set cwd
-                    "--dit_fsdp",
-                    "--t5_fsdp",
-                    "--ulysses_size",
-                    str(self.num_gpus),
-                ] + common_args
-            else:
-                cmd = [
-                    "python3",
-                    script_name,  # Use just the script name since we'll set cwd
-                    "--num_persistent_param_in_dit",
-                    str(self.used_vram_params),
-                ] + common_args
-
-            print(f"Running command: {' '.join(cmd)}")
-            print(f"Working directory: {script_dir}")
-
-            # Run inference with correct working directory
-            with open("inference.log", "w") as log_file:
-                process = subprocess.Popen(
-                    cmd,
-                    cwd=script_dir,  # Set working directory to script location
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.STDOUT,
-                    text=True,
-                    bufsize=1,
+                # Enable advanced GPU features for maximum speed
+                torch.backends.cuda.enable_flash_sdp(True)
+                torch.backends.cudnn.benchmark = (
+                    True  # Optimize for consistent input sizes
+                )
+                torch.backends.cuda.matmul.allow_tf32 = (
+                    True  # Enable TF32 for faster matmul
                 )
-                for line in process.stdout:
-                    print(line, end="")
-                    log_file.write(line)
-                process.wait()
-
-            if process.returncode != 0:
-                raise RuntimeError("Inference failed. Check inference.log for details.")
-
-            # The output video might be in the script directory
-            output_video_path = os.path.join(script_dir, f"{result_filename}.mp4")
-            if not os.path.exists(output_video_path):
-                # Try current directory as fallback
-                output_video_path = f"{result_filename}.mp4"
-
-            if not os.path.exists(output_video_path):
-                raise RuntimeError(f"Output video not found: {output_video_path}")
-
-            return output_video_path
-
-        finally:
-            # Clean up temporary files
-            for f in temp_files_to_cleanup:
-                try:
-                    if os.path.exists(f):
-                        os.remove(f)
-                        print(f"[INFO] Removed temporary file: {f}")
-                except Exception as e:
-                    print(f"[WARNING] Could not remove {f}: {e}")
+                torch.backends.cudnn.allow_tf32 = True  # Enable TF32 for convolutions
+                torch.cuda.empty_cache()  # Clear any existing memory
+                print(
+                    "⚡ Enabled Flash-SDP, TF32, and cuDNN benchmarking for maximum throughput"
+                )
+            else:
+                print("🔧 Standard GPU optimizations enabled")
+                torch.backends.cuda.enable_flash_sdp(True)
+                torch.cuda.empty_cache()
+
+        print("✅ Model setup completed successfully!")
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
-        Main inference method called by Hugging Face Inference API.
-
-        Args:
-            data: Input data dictionary containing:
-                - inputs: Main input data
-                - parameters: Optional parameters
-
-        Returns:
-            Dictionary containing the generated video
+        Handle inference requests
+
+        Expected input format:
+        {
+            "inputs": {
+                "image": "base64_encoded_image_string",
+                "first_audio": "base64_encoded_audio_string",
+                "prompt": "Text prompt describing the desired interaction",
+                "second_audio": "base64_encoded_audio_string (optional)",
+                "num_frames": 81,
+                "sampling_steps": 40,
+                "seed": 42,
+                "turbo": true
+            },
+            "parameters": {
+                // Optional parameters can also be passed here
+            }
+        }
         """
         try:
-            # Extract inputs
+            # Parse input data
             inputs = data.get("inputs", {})
             parameters = data.get("parameters", {})
 
-            # Required inputs
-            prompt = inputs.get("prompt", "")
-            if not prompt:
-                raise ValueError("Prompt is required")
-
-            # Image input (base64 encoded)
-            image_data = inputs.get("image")
-            if not image_data:
-                raise ValueError("Image is required")
-
-            # Audio inputs (base64 encoded)
-            audio_spk1_data = inputs.get("audio_spk1")
-            if not audio_spk1_data:
-                raise ValueError("Audio for speaker 1 is required")
-
-            audio_spk2_data = inputs.get("audio_spk2")  # Optional
-
-            # Parameters
-            sample_steps = parameters.get("sample_steps", 12)
-            sample_steps = max(2, min(25, sample_steps))  # Clamp between 2 and 25
-
-            # Save base64 files to temporary locations
-            temp_files_to_cleanup = []
-
-            # Save image
-            image_path = self.save_base64_file(image_data, ".png")
-            temp_files_to_cleanup.append(image_path)
-
-            # Save audio files
-            audio_spk1_path = self.save_base64_file(audio_spk1_data, ".wav")
-            temp_files_to_cleanup.append(audio_spk1_path)
-
-            audio_spk2_path = None
-            if audio_spk2_data:
-                audio_spk2_path = self.save_base64_file(audio_spk2_data, ".wav")
-                temp_files_to_cleanup.append(audio_spk2_path)
-
-            # Run inference
-            output_video_path = self.run_inference(
-                prompt=prompt,
-                cond_image_path=image_path,
-                cond_audio_path_spk1=audio_spk1_path,
-                cond_audio_path_spk2=audio_spk2_path,
-                sample_steps=sample_steps,
+            # Merge parameters with inputs (parameters take precedence)
+            config = {**inputs, **parameters}
+
+            # Validate required inputs
+            if "image" not in config:
+                raise ValueError("Missing required input: image")
+            if "first_audio" not in config:
+                raise ValueError("Missing required input: first_audio")
+
+            # Extract parameters with defaults
+            image_b64 = config["image"]
+            first_audio_b64 = config["first_audio"]
+            prompt = config.get(
+                "prompt",
+                "A smiling man and woman wearing headphones sit in front of microphones, appearing to host a podcast.",
             )
+            second_audio_b64 = config.get("second_audio", None)
+            num_frames = config.get("num_frames", 81)
+            sampling_steps = config.get("sampling_steps", 40)
+            seed = config.get("seed", None)
+            turbo = config.get("turbo", True)
+
+            # Validate parameters
+            if not isinstance(num_frames, int) or num_frames < 25 or num_frames > 201:
+                raise ValueError("num_frames must be an integer between 25 and 201")
+            if (
+                not isinstance(sampling_steps, int)
+                or sampling_steps < 2
+                or sampling_steps > 100
+            ):
+                raise ValueError("sampling_steps must be an integer between 2 and 100")
+
+            # Auto-correct frame count to nearest valid value (4n+1 format)
+            original_frames = num_frames
+            if (num_frames - 1) % 4 != 0:
+                # Find the nearest valid values
+                n_lower = (num_frames - 1) // 4
+                n_upper = n_lower + 1
+
+                frames_lower = 4 * n_lower + 1
+                frames_upper = 4 * n_upper + 1
+
+                # Choose the closer one
+                if abs(num_frames - frames_lower) <= abs(num_frames - frames_upper):
+                    num_frames = frames_lower
+                else:
+                    num_frames = frames_upper
+
+                # Ensure it's within bounds [25, 201]
+                num_frames = max(25, min(num_frames, 201))
+
+                # Final safety check and adjustment if needed
+                while (num_frames - 1) % 4 != 0 and num_frames <= 201:
+                    num_frames += 1
+
+                print(
+                    f"📐 Auto-corrected num_frames from {original_frames} to {num_frames} (required format: 4n+1)"
+                )
 
-            # Read and encode output video
-            with open(output_video_path, "rb") as f:
-                video_data = base64.b64encode(f.read()).decode()
+            # Set random seed
+            if seed is None:
+                seed = random.randint(0, 99999999)
+
+            print(f"🎬 Generating video with seed: {seed}")
+
+            with tempfile.TemporaryDirectory() as temp_dir:
+                # Save input files from base64
+                image_path = os.path.join(temp_dir, "input_image.jpg")
+                first_audio_path = os.path.join(temp_dir, "first_audio.wav")
+
+                decode_base64_file(image_b64, image_path)
+                decode_base64_file(first_audio_b64, first_audio_path)
+
+                second_audio_path = None
+                if second_audio_b64:
+                    second_audio_path = os.path.join(temp_dir, "second_audio.wav")
+                    decode_base64_file(second_audio_b64, second_audio_path)
+
+                is_multi_person = second_audio_path is not None
+                audio_save_dir = os.path.join(temp_dir, "audio_embeddings")
+                os.makedirs(audio_save_dir, exist_ok=True)
+
+                # Process audio and generate embeddings
+                if is_multi_person:
+                    print("🎤 Processing multi-person audio...")
+                    audio_type = "add"  # Sequential by default
+                    speech1, speech2, combined_speech = audio_prepare_multi(
+                        first_audio_path, second_audio_path, audio_type
+                    )
+
+                    # Generate embeddings on optimal device
+                    embedding1 = get_embedding(
+                        speech1,
+                        self.wav2vec_feature_extractor,
+                        self.audio_encoder,
+                        device=self.audio_device,
+                    )
+                    embedding2 = get_embedding(
+                        speech2,
+                        self.wav2vec_feature_extractor,
+                        self.audio_encoder,
+                        device=self.audio_device,
+                    )
+
+                    # Save embeddings and audio
+                    emb1_path = os.path.join(audio_save_dir, "1.pt")
+                    emb2_path = os.path.join(audio_save_dir, "2.pt")
+                    sum_audio_path = os.path.join(audio_save_dir, "sum.wav")
+
+                    torch.save(embedding1, emb1_path)
+                    torch.save(embedding2, emb2_path)
+                    sf.write(sum_audio_path, combined_speech, 16000)
+
+                    # Create input data
+                    input_data = {
+                        "prompt": prompt,
+                        "cond_image": image_path,
+                        "audio_type": audio_type,
+                        "cond_audio": {"person1": emb1_path, "person2": emb2_path},
+                        "video_audio": sum_audio_path,
+                    }
+                else:
+                    print("🎤 Processing single-person audio...")
+                    speech = audio_prepare_single(first_audio_path)
+                    embedding = get_embedding(
+                        speech,
+                        self.wav2vec_feature_extractor,
+                        self.audio_encoder,
+                        device=self.audio_device,
+                    )
+
+                    # Save embedding and audio
+                    emb_path = os.path.join(audio_save_dir, "1.pt")
+                    sum_audio_path = os.path.join(audio_save_dir, "sum.wav")
+
+                    torch.save(embedding, emb_path)
+                    sf.write(sum_audio_path, speech, 16000)
+
+                    # Create input data
+                    input_data = {
+                        "prompt": prompt,
+                        "cond_image": image_path,
+                        "cond_audio": {"person1": emb_path},
+                        "video_audio": sum_audio_path,
+                    }
+
+                print("🎬 Generating video...")
+
+                # Configure generation parameters based on turbo mode and VRAM availability
+                high_vram = (
+                    torch.cuda.is_available()
+                    and torch.cuda.get_device_properties(0).total_memory > 40 * 1024**3
+                )
 
-            # Clean up input files
-            for f in temp_files_to_cleanup:
-                try:
-                    if os.path.exists(f):
-                        os.remove(f)
-                except Exception as e:
-                    print(f"[WARNING] Could not remove {f}: {e}")
+                if turbo:
+                    teacache_thresh = 0.8
+                    text_guide_scale = 3.0
+                    audio_guide_scale = 3.0
+                    shift = 5.0
+                    offload_model = False  # Never offload in turbo mode
+                    print(
+                        f"🚀 TURBO MODE: {sampling_steps} steps, thresh={teacache_thresh}"
+                    )
+                else:
+                    teacache_thresh = 0.3
+                    text_guide_scale = 5.0
+                    audio_guide_scale = 4.0
+                    shift = 7.0
+                    offload_model = (
+                        not high_vram
+                    )  # Don't offload with high VRAM for maximum speed
+                    print(
+                        f"🎬 QUALITY MODE: {sampling_steps} steps{', keeping models in GPU' if high_vram else ''}"
+                    )
+
+                # Configure optimizations using SimpleNamespace
+                extra_args = SimpleNamespace(
+                    use_teacache=True,
+                    teacache_thresh=teacache_thresh,
+                    use_apg=False,
+                    size="multitalk-480",
+                )
 
-            # Clean up output video
-            try:
-                if os.path.exists(output_video_path):
-                    os.remove(output_video_path)
-            except Exception as e:
-                print(f"[WARNING] Could not remove output video: {e}")
+                # Generate video using loaded pipeline
+                video = self.wan_i2v.generate(
+                    input_data,
+                    size_buckget="multitalk-480",
+                    motion_frame=25,
+                    frame_num=num_frames,
+                    shift=shift,
+                    sampling_steps=sampling_steps,
+                    text_guide_scale=text_guide_scale,
+                    audio_guide_scale=audio_guide_scale,
+                    seed=seed,
+                    offload_model=offload_model,
+                    max_frames_num=num_frames,
+                    extra_args=extra_args,
+                )
 
-            return {"generated_video": video_data, "content_type": "video/mp4"}
+                # Save video
+                output_name = f"multitalk_{abs(hash(prompt + str(seed))) % 10000}"
+                print(f"💾 Saving video...")
+                save_video_ffmpeg(video, output_name, [input_data["video_audio"]])
+
+                # Find generated video
+                output_file = f"{output_name}.mp4"
+                if not os.path.exists(output_file):
+                    # Look for any mp4 files with our output name
+                    for file in os.listdir("."):
+                        if output_name in file and file.endswith(".mp4"):
+                            output_file = file
+                            break
+
+                    if not os.path.exists(output_file):
+                        raise RuntimeError(
+                            f"Video generation failed - output file not found"
+                        )
+
+                # Encode video to base64
+                video_b64 = encode_file_to_base64(output_file)
+
+                # Cleanup GPU memory
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+
+                print(f"✅ Video generation completed successfully!")
+
+                return {
+                    "video": video_b64,
+                    "metadata": {
+                        "seed": seed,
+                        "num_frames": num_frames,
+                        "sampling_steps": sampling_steps,
+                        "turbo": turbo,
+                        "is_multi_person": is_multi_person,
+                        "prompt": prompt,
+                        "original_frames_requested": original_frames,
+                    },
+                }
 
         except Exception as e:
-            return {"error": str(e)}
+            print(f"❌ Error during inference: {str(e)}")
+            return {"error": str(e), "type": type(e).__name__}

handler_norun.py

@@ -0,0 +1,470 @@
+import torch
+import os
+import shutil
+import subprocess
+import json
+import tempfile
+import base64
+from typing import Dict, Any, Optional, List
+import soundfile as sf
+from datetime import datetime
+from huggingface_hub import snapshot_download
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        """
+        Initialize the handler for Meigen MultiTalk inference.
+        Optimized for serverless deployment with pre-downloaded models.
+
+        Args:
+            path: Path to the model directory (not used in this implementation)
+        """
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Set up the script path - try multiple possible locations
+        self.script_path = self.find_script_path()
+        if not self.script_path:
+            raise RuntimeError("Could not locate generate_multitalk.py script")
+
+        print(f"Found script at: {self.script_path}")
+
+        self.setup_models()
+        self.setup_gpu_config()
+
+    def find_script_path(self) -> Optional[str]:
+        """
+        Find the location of the generate_multitalk.py script.
+
+        Returns:
+            Path to the script if found, None otherwise
+        """
+        possible_paths = [
+            "generate_multitalk.py",  # Current directory
+            "/app/generate_multitalk.py",  # App directory
+            "./src/generate_multitalk.py",  # Src directory
+            "../generate_multitalk.py",  # Parent directory
+            "./weights/MeiGen-MultiTalk/generate_multitalk.py",  # In model directory
+            "./MeiGen-MultiTalk/generate_multitalk.py",  # Alternative model directory
+        ]
+
+        for path in possible_paths:
+            if os.path.exists(path):
+                return os.path.abspath(path)
+
+        # Try to find it recursively in common directories
+        search_dirs = ["/app", ".", "./weights", "./src"]
+        for search_dir in search_dirs:
+            if os.path.exists(search_dir):
+                for root, dirs, files in os.walk(search_dir):
+                    if "generate_multitalk.py" in files:
+                        return os.path.join(root, "generate_multitalk.py")
+
+        return None
+
+    def setup_models(self):
+        """Setup models - download only if not already present."""
+        print("Setting up models...")
+
+        # Define model paths
+        self.wan_model_path = "./weights/Wan2.1-I2V-14B-480P"
+        self.wav2vec_path = "./weights/chinese-wav2vec2-base"
+        self.multitalk_path = "./weights/MeiGen-MultiTalk"
+
+        # Only download if models don't exist (for local development)
+        # In serverless deployment, these should be pre-baked into the container
+        if not os.path.exists(self.wan_model_path):
+            print("Downloading Wan2.1-I2V-14B-480P model...")
+            self.wan_model_path = snapshot_download(
+                repo_id="Wan-AI/Wan2.1-I2V-14B-480P",
+                local_dir=self.wan_model_path,
+            )
+
+        if not os.path.exists(self.wav2vec_path):
+            print("Downloading Chinese wav2vec2 model...")
+            self.wav2vec_path = snapshot_download(
+                repo_id="TencentGameMate/chinese-wav2vec2-base",
+                local_dir=self.wav2vec_path,
+            )
+
+        if not os.path.exists(self.multitalk_path):
+            print("Downloading MeiGen MultiTalk weights...")
+            self.multitalk_path = snapshot_download(
+                repo_id="MeiGen-AI/MeiGen-MultiTalk",
+                local_dir=self.multitalk_path,
+            )
+
+        # Setup model files if not already done
+        self.setup_model_files()
+
+    def setup_model_files(self):
+        """Setup model files by copying and renaming as needed."""
+        base_model_dir = self.wan_model_path
+        multitalk_dir = self.multitalk_path
+
+        # Check if setup is already done
+        multitalk_weights_path = os.path.join(base_model_dir, "multitalk.safetensors")
+        if os.path.exists(multitalk_weights_path):
+            print("Model files already set up, skipping setup.")
+            return
+
+        # File to rename
+        original_index = os.path.join(
+            base_model_dir, "diffusion_pytorch_model.safetensors.index.json"
+        )
+        backup_index = os.path.join(
+            base_model_dir, "diffusion_pytorch_model.safetensors.index.json_old"
+        )
+
+        # Rename the original index file
+        if os.path.exists(original_index) and not os.path.exists(backup_index):
+            os.rename(original_index, backup_index)
+            print("Renamed original index file to .json_old")
+
+        # Copy updated index file from MultiTalk
+        multitalk_index = os.path.join(
+            multitalk_dir, "diffusion_pytorch_model.safetensors.index.json"
+        )
+        if os.path.exists(multitalk_index):
+            shutil.copy2(multitalk_index, base_model_dir)
+
+        # Copy MultiTalk model weights
+        multitalk_weights_src = os.path.join(multitalk_dir, "multitalk.safetensors")
+        if os.path.exists(multitalk_weights_src):
+            shutil.copy2(multitalk_weights_src, base_model_dir)
+
+        print("Copied MultiTalk files into base model directory.")
+
+    def setup_gpu_config(self):
+        """Setup GPU configuration and VRAM parameters."""
+        self.num_gpus = torch.cuda.device_count()
+        print(f"GPU AVAILABLE: {self.num_gpus}")
+
+        if torch.cuda.is_available():
+            self.gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
+            print(f"Current GPU: {self.gpu_name}")
+
+            # GPU to VRAM mapping
+            GPU_TO_VRAM_PARAMS = {
+                "NVIDIA A100": 11000000000,
+                "NVIDIA A100-SXM4-40GB": 11000000000,
+                "NVIDIA A100-SXM4-80GB": 22000000000,
+                "NVIDIA L4": 5000000000,
+                "NVIDIA L40S": 11000000000,
+            }
+
+            self.used_vram_params = GPU_TO_VRAM_PARAMS.get(self.gpu_name, 5000000000)
+            print(f"Using {self.used_vram_params} for num_persistent_param_in_dit")
+        else:
+            raise RuntimeError("No CUDA-compatible GPU found. A GPU is required.")
+
+    def trim_audio_to_5s_temp(self, audio_path: str, sample_rate: int = 16000) -> str:
+        """
+        Trim audio to 5 seconds and save to temporary file.
+
+        Args:
+            audio_path: Path to input audio file
+            sample_rate: Target sample rate
+
+        Returns:
+            Path to trimmed audio file
+        """
+        max_duration_sec = 5
+        audio, sr = sf.read(audio_path)
+
+        if sr != sample_rate:
+            sample_rate = sr
+
+        max_samples = max_duration_sec * sample_rate
+        if len(audio) > max_samples:
+            audio = audio[:max_samples]
+
+        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
+        base_name = os.path.splitext(os.path.basename(audio_path))[0]
+        temp_filename = f"{base_name}_trimmed_{timestamp}.wav"
+        temp_path = os.path.join(tempfile.gettempdir(), temp_filename)
+
+        sf.write(temp_path, audio, samplerate=sample_rate)
+        return temp_path
+
+    def save_base64_file(self, base64_data: str, file_extension: str) -> str:
+        """
+        Save base64 encoded file to temporary location.
+
+        Args:
+            base64_data: Base64 encoded file data
+            file_extension: File extension (e.g., '.png', '.wav')
+
+        Returns:
+            Path to saved file
+        """
+        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
+        temp_filename = f"temp_file_{timestamp}{file_extension}"
+        temp_path = os.path.join(tempfile.gettempdir(), temp_filename)
+
+        with open(temp_path, "wb") as f:
+            f.write(base64.b64decode(base64_data))
+
+        return temp_path
+
+    def create_temp_input_json(
+        self,
+        prompt: str,
+        cond_image_path: str,
+        cond_audio_path_spk1: str,
+        cond_audio_path_spk2: Optional[str] = None,
+    ) -> str:
+        """
+        Create a temporary JSON file with the user-provided prompt, image, and audio paths.
+
+        Args:
+            prompt: Text prompt for generation
+            cond_image_path: Path to conditioning image
+            cond_audio_path_spk1: Path to first speaker audio
+            cond_audio_path_spk2: Path to second speaker audio (optional)
+
+        Returns:
+            Path to temporary JSON file
+        """
+        if cond_audio_path_spk2 is None:
+            data = {
+                "prompt": prompt,
+                "cond_image": cond_image_path,
+                "cond_audio": {"person1": cond_audio_path_spk1},
+            }
+        else:
+            data = {
+                "prompt": prompt,
+                "cond_image": cond_image_path,
+                "audio_type": "para",
+                "cond_audio": {
+                    "person1": cond_audio_path_spk1,
+                    "person2": cond_audio_path_spk2,
+                },
+            }
+
+        temp_json = tempfile.NamedTemporaryFile(
+            delete=False, suffix=".json", mode="w", encoding="utf-8"
+        )
+        json.dump(data, temp_json, indent=4)
+        temp_json_path = temp_json.name
+        temp_json.close()
+
+        print(f"Temporary input JSON saved to: {temp_json_path}")
+        return temp_json_path
+
+    def run_inference(
+        self,
+        prompt: str,
+        cond_image_path: str,
+        cond_audio_path_spk1: str,
+        cond_audio_path_spk2: Optional[str] = None,
+        sample_steps: int = 12,
+    ) -> str:
+        """
+        Run the MultiTalk inference.
+
+        Args:
+            prompt: Text prompt
+            cond_image_path: Path to conditioning image
+            cond_audio_path_spk1: Path to first speaker audio
+            cond_audio_path_spk2: Path to second speaker audio (optional)
+            sample_steps: Number of sampling steps
+
+        Returns:
+            Path to generated video
+        """
+        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
+        result_filename = f"meigen_multitalk_result_{sample_steps}_steps_{timestamp}"
+        temp_files_to_cleanup = []
+
+        try:
+            # Trim audio files to 5 seconds
+            trimmed_audio_path_spk1 = self.trim_audio_to_5s_temp(cond_audio_path_spk1)
+            if trimmed_audio_path_spk1 != cond_audio_path_spk1:
+                cond_audio_path_spk1 = trimmed_audio_path_spk1
+                temp_files_to_cleanup.append(trimmed_audio_path_spk1)
+
+            if cond_audio_path_spk2 is not None:
+                trimmed_audio_path_spk2 = self.trim_audio_to_5s_temp(
+                    cond_audio_path_spk2
+                )
+                if trimmed_audio_path_spk2 != cond_audio_path_spk2:
+                    cond_audio_path_spk2 = trimmed_audio_path_spk2
+                    temp_files_to_cleanup.append(trimmed_audio_path_spk2)
+
+            # Prepare input JSON
+            input_json_path = self.create_temp_input_json(
+                prompt, cond_image_path, cond_audio_path_spk1, cond_audio_path_spk2
+            )
+            temp_files_to_cleanup.append(input_json_path)
+
+            # Base command arguments
+            common_args = [
+                "--ckpt_dir",
+                self.wan_model_path,
+                "--wav2vec_dir",
+                self.wav2vec_path,
+                "--input_json",
+                input_json_path,
+                "--sample_steps",
+                str(sample_steps),
+                "--mode",
+                "streaming",
+                "--use_teacache",
+                "--save_file",
+                result_filename,
+            ]
+
+            # Get the directory containing the script to set as working directory
+            script_dir = os.path.dirname(self.script_path)
+            script_name = os.path.basename(self.script_path)
+
+            # Choose command based on number of GPUs
+            if self.num_gpus > 1:
+                cmd = [
+                    "torchrun",
+                    f"--nproc_per_node={self.num_gpus}",
+                    "--standalone",
+                    script_name,  # Use just the script name since we'll set cwd
+                    "--dit_fsdp",
+                    "--t5_fsdp",
+                    "--ulysses_size",
+                    str(self.num_gpus),
+                ] + common_args
+            else:
+                cmd = [
+                    "python3",
+                    script_name,  # Use just the script name since we'll set cwd
+                    "--num_persistent_param_in_dit",
+                    str(self.used_vram_params),
+                ] + common_args
+
+            print(f"Running command: {' '.join(cmd)}")
+            print(f"Working directory: {script_dir}")
+
+            # Run inference with correct working directory
+            with open("inference.log", "w") as log_file:
+                process = subprocess.Popen(
+                    cmd,
+                    cwd=script_dir,  # Set working directory to script location
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT,
+                    text=True,
+                    bufsize=1,
+                )
+                for line in process.stdout:
+                    print(line, end="")
+                    log_file.write(line)
+                process.wait()
+
+            if process.returncode != 0:
+                raise RuntimeError("Inference failed. Check inference.log for details.")
+
+            # The output video might be in the script directory
+            output_video_path = os.path.join(script_dir, f"{result_filename}.mp4")
+            if not os.path.exists(output_video_path):
+                # Try current directory as fallback
+                output_video_path = f"{result_filename}.mp4"
+
+            if not os.path.exists(output_video_path):
+                raise RuntimeError(f"Output video not found: {output_video_path}")
+
+            return output_video_path
+
+        finally:
+            # Clean up temporary files
+            for f in temp_files_to_cleanup:
+                try:
+                    if os.path.exists(f):
+                        os.remove(f)
+                        print(f"[INFO] Removed temporary file: {f}")
+                except Exception as e:
+                    print(f"[WARNING] Could not remove {f}: {e}")
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Main inference method called by Hugging Face Inference API.
+
+        Args:
+            data: Input data dictionary containing:
+                - inputs: Main input data
+                - parameters: Optional parameters
+
+        Returns:
+            Dictionary containing the generated video
+        """
+        try:
+            # Extract inputs
+            inputs = data.get("inputs", {})
+            parameters = data.get("parameters", {})
+
+            # Required inputs
+            prompt = inputs.get("prompt", "")
+            if not prompt:
+                raise ValueError("Prompt is required")
+
+            # Image input (base64 encoded)
+            image_data = inputs.get("image")
+            if not image_data:
+                raise ValueError("Image is required")
+
+            # Audio inputs (base64 encoded)
+            audio_spk1_data = inputs.get("audio_spk1")
+            if not audio_spk1_data:
+                raise ValueError("Audio for speaker 1 is required")
+
+            audio_spk2_data = inputs.get("audio_spk2")  # Optional
+
+            # Parameters
+            sample_steps = parameters.get("sample_steps", 12)
+            sample_steps = max(2, min(25, sample_steps))  # Clamp between 2 and 25
+
+            # Save base64 files to temporary locations
+            temp_files_to_cleanup = []
+
+            # Save image
+            image_path = self.save_base64_file(image_data, ".png")
+            temp_files_to_cleanup.append(image_path)
+
+            # Save audio files
+            audio_spk1_path = self.save_base64_file(audio_spk1_data, ".wav")
+            temp_files_to_cleanup.append(audio_spk1_path)
+
+            audio_spk2_path = None
+            if audio_spk2_data:
+                audio_spk2_path = self.save_base64_file(audio_spk2_data, ".wav")
+                temp_files_to_cleanup.append(audio_spk2_path)
+
+            # Run inference
+            output_video_path = self.run_inference(
+                prompt=prompt,
+                cond_image_path=image_path,
+                cond_audio_path_spk1=audio_spk1_path,
+                cond_audio_path_spk2=audio_spk2_path,
+                sample_steps=sample_steps,
+            )
+
+            # Read and encode output video
+            with open(output_video_path, "rb") as f:
+                video_data = base64.b64encode(f.read()).decode()
+
+            # Clean up input files
+            for f in temp_files_to_cleanup:
+                try:
+                    if os.path.exists(f):
+                        os.remove(f)
+                except Exception as e:
+                    print(f"[WARNING] Could not remove {f}: {e}")
+
+            # Clean up output video
+            try:
+                if os.path.exists(output_video_path):
+                    os.remove(output_video_path)
+            except Exception as e:
+                print(f"[WARNING] Could not remove output video: {e}")
+
+            return {"generated_video": video_data, "content_type": "video/mp4"}
+
+        except Exception as e:
+            return {"error": str(e)}

requirements.txt

@@ -1,11 +1,6 @@
-torch==2.6.0
-torchvision==0.21.0
-torchaudio==2.6.0
 opencv-python>=4.9.0.80
 diffusers>=0.31.0
 transformers>=4.49.0
-peft
-xformers
 tokenizers>=0.20.3
 accelerate>=1.1.1
 tqdm
@@ -18,11 +13,15 @@ gradio>=5.0.0
 numpy>=1.23.5,<2
 xfuser>=0.4.1
 pyloudnorm
-soundfile
-librosa
-einops
-numpy<2
+torch>=2.4.0
+torchvision>=0.19.0
+torchaudio>=2.4.0
+soundfile>=0.12.1
+librosa>=0.10.0
+einops>=0.8.0
+huggingface-hub>=0.19.0
+xformers>=0.0.28
 ninja
-psutil
 packaging
+safetensors>=0.3.1
 https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

requirements_bck.txt

@@ -0,0 +1,28 @@
+torch==2.6.0
+torchvision==0.21.0
+torchaudio==2.6.0
+opencv-python>=4.9.0.80
+diffusers>=0.31.0
+transformers>=4.49.0
+peft
+xformers
+tokenizers>=0.20.3
+accelerate>=1.1.1
+tqdm
+imageio
+easydict
+ftfy
+dashscope
+imageio-ffmpeg
+gradio>=5.0.0
+numpy>=1.23.5,<2
+xfuser>=0.4.1
+pyloudnorm
+soundfile
+librosa
+einops
+numpy<2
+ninja
+psutil
+packaging
+https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl