Initial Commit

handler.py

@@ -0,0 +1,398 @@
+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.
+
+        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):
+        """Download and setup all required models."""
+        print("Setting up models...")
+
+        # Download Wan2.1-I2V-14B-480P model
+        self.wan_model_path = snapshot_download(
+            repo_id="Wan-AI/Wan2.1-I2V-14B-480P",
+            local_dir="./weights/Wan2.1-I2V-14B-480P",
+        )
+
+        # Download Chinese wav2vec2 model
+        self.wav2vec_path = snapshot_download(
+            repo_id="TencentGameMate/chinese-wav2vec2-base",
+            local_dir="./weights/chinese-wav2vec2-base",
+        )
+
+        # Download MeiGen MultiTalk weights
+        self.multitalk_path = snapshot_download(
+            repo_id="MeiGen-AI/MeiGen-MultiTalk",
+            local_dir="./weights/MeiGen-MultiTalk",
+        )
+
+        self.setup_model_files()
+
+    def setup_model_files(self):
+        """Setup model files by copying and renaming as needed."""
+        base_model_dir = "./weights/Wan2.1-I2V-14B-480P"
+        multitalk_dir = "./weights/MeiGen-MultiTalk"
+
+        # 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):
+            os.rename(original_index, backup_index)
+            print("Renamed original index file to .json_old")
+
+        # Copy updated index file from MultiTalk
+        shutil.copy2(
+            os.path.join(
+                multitalk_dir, "diffusion_pytorch_model.safetensors.index.json"
+            ),
+            base_model_dir,
+        )
+
+        # Copy MultiTalk model weights
+        shutil.copy2(
+            os.path.join(multitalk_dir, "multitalk.safetensors"), 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",
+                "weights/Wan2.1-I2V-14B-480P",
+                "--wav2vec_dir",
+                "weights/chinese-wav2vec2-base",
+                "--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)}