Blame for handler_norun.py

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)}