Blame for handler_runningpy

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"
        self.setup_models()
        self.setup_gpu_config()

    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,
            ]

            # Choose command based on number of GPUs
            if self.num_gpus > 1:
                cmd = [
                    "torchrun",
                    f"--nproc_per_node={self.num_gpus}",
                    "--standalone",
                    "generate_multitalk.py",
                    "--dit_fsdp",
                    "--t5_fsdp",
                    "--ulysses_size",
                    str(self.num_gpus),
                ] + common_args
            else:
                cmd = [
                    "python3",
                    "generate_multitalk.py",
                    "--num_persistent_param_in_dit",
                    str(self.used_vram_params),
                ] + common_args

            # Run inference
            with open("inference.log", "w") as log_file:
                process = subprocess.Popen(
                    cmd,
                    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.")

            return f"{result_filename}.mp4"

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