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