# # Hugging Face Inference Endpoint Handler for MultiTalk # import os # import subprocess # import time # import json # import tempfile # import logging # import sys # import warnings # import shutil # import base64 # from typing import Optional, Dict, Any, Union # from datetime import datetime # from types import SimpleNamespace # from io import BytesIO # # Suppress warnings for cleaner output # warnings.filterwarnings("ignore") # 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]), # ] # ) # new_human_speech2 = np.concatenate( # [ # np.zeros(human_speech_array1.shape[0]), # human_speech_array2[: human_speech_array2.shape[0]], # ] # ) # 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 len(embeddings) == 0: # print("⚠️ Failed to extract audio embedding") # return None # 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] # # command = ["wget", "-O", dest, url] # # 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") # def download_weights(url: str, dest: str) -> None: # """ # Download weights using huggingface snapshot_download when appropriate, # otherwise use direct HTTP download (e.g., for .tar or .pth files). # """ # start = time.time() # print("[!] Initiating download from URL: ", url) # print("[~] Destination path: ", dest) # # If it's a .tar or direct file, treat as file URL # if url.endswith(".tar") or url.startswith("http"): # try: # # Create parent directory if missing # os.makedirs(os.path.dirname(dest), exist_ok=True) # print("[~] Downloading using direct URL...") # import requests # with requests.get(url, stream=True) as r: # r.raise_for_status() # with open(dest, "wb") as f: # shutil.copyfileobj(r.raw, f) # print("[+] File downloaded to:", dest) # except Exception as e: # print(f"[ERROR] Failed to download from {url}") # raise e # else: # # Assume it's a huggingface repo ID # print("[~] Downloading using HuggingFace snapshot...") # try: # local_dir = snapshot_download( # repo_id=url, # use full Hugging Face repo ID # local_dir=dest, # local_dir_use_symlinks=False # ) # print("[+] Snapshot downloaded to:", local_dir) # except Exception as e: # print(f"[ERROR] Failed to snapshot_download: {url}") # raise e # print("[+] Download completed in:", time.time() - start, "seconds") # 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] # # Decode and save # with open(output_path, "wb") as f: # f.write(base64.b64decode(base64_string)) # return output_path # except Exception as e: # raise ValueError(f"Failed to decode base64 file: {str(e)}") # 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)}") # 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}/" # ) # # 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", # ] # 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)], # ) # # 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 # ) # # 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") # if vram_gb > 40: # High VRAM setup # print( # "🚀 High-VRAM detected: Enabling maximum performance optimizations" # ) # # 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 # ) # 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]: # """ # 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: # # Parse input data # inputs = data.get("inputs", {}) # parameters = data.get("parameters", {}) # # 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)" # ) # # 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 # ) # 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", # ) # # 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, # ) # # 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: # print(f"❌ Error during inference: {str(e)}") # return {"error": str(e), "type": type(e).__name__} # Hugging Face Inference Endpoint Handler for MultiTalk import os import subprocess import time import json import tempfile import logging import sys import warnings import shutil import base64 import requests from typing import Optional, Dict, Any, Union from datetime import datetime from types import SimpleNamespace from io import BytesIO # Suppress warnings for cleaner output warnings.filterwarnings("ignore") 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]), ] ) new_human_speech2 = np.concatenate( [ np.zeros(human_speech_array1.shape[0]), human_speech_array2[: human_speech_array2.shape[0]], ] ) 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 len(embeddings) == 0: print("⚠️ Failed to extract audio embedding") return None 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_with_resume( url: str, dest: str, max_retries: int = 3, chunk_size: int = 8192 * 1024 ) -> None: """ Download file with resume capability and robust error handling. Args: url: URL to download from dest: Destination file path max_retries: Maximum number of retry attempts chunk_size: Size of chunks to download (8MB default) """ print(f"[!] Downloading from: {url}") print(f"[~] Destination: {dest}") # Create destination directory if it doesn't exist os.makedirs(os.path.dirname(dest), exist_ok=True) # Get initial file size if resuming initial_size = 0 if os.path.exists(dest): initial_size = os.path.getsize(dest) print(f"[~] Resuming download from {initial_size} bytes") for attempt in range(max_retries): try: # Set up headers for resume headers = {} if initial_size > 0: headers["Range"] = f"bytes={initial_size}-" # Start download with streaming response = requests.get(url, headers=headers, stream=True, timeout=60) response.raise_for_status() # Get total file size total_size = initial_size if "content-length" in response.headers: total_size = initial_size + int(response.headers["content-length"]) elif "content-range" in response.headers: # Parse content-range header: bytes start-end/total content_range = response.headers["content-range"] if "/" in content_range: total_size = int(content_range.split("/")[-1]) print(f"[~] Total size: {total_size / (1024**3):.2f} GB") # Download with progress tracking start_time = time.time() last_progress_time = start_time # Open file in append mode if resuming, write mode otherwise mode = "ab" if initial_size > 0 else "wb" with open(dest, mode) as f: downloaded = initial_size for chunk in response.iter_content(chunk_size=chunk_size): if chunk: f.write(chunk) downloaded += len(chunk) # Progress update every 10 seconds current_time = time.time() if current_time - last_progress_time >= 10: if total_size > 0: progress = (downloaded / total_size) * 100 speed = ( downloaded / (current_time - start_time) / (1024**2) ) # MB/s print( f"[~] Progress: {progress:.1f}% ({downloaded/(1024**3):.2f}/{total_size/(1024**3):.2f} GB) - Speed: {speed:.1f} MB/s" ) last_progress_time = current_time # Verify download completed successfully final_size = os.path.getsize(dest) if total_size > 0 and final_size < total_size: print(f"[!] Download incomplete: {final_size}/{total_size} bytes") initial_size = final_size continue elapsed_time = time.time() - start_time print(f"[+] Download completed in {elapsed_time:.2f} seconds") print(f"[+] Final size: {final_size / (1024**3):.2f} GB") return except requests.exceptions.RequestException as e: print(f"[!] Download attempt {attempt + 1} failed: {str(e)}") if attempt < max_retries - 1: print(f"[~] Retrying in 5 seconds...") time.sleep(5) initial_size = os.path.getsize(dest) if os.path.exists(dest) else 0 else: print(f"[ERROR] Failed to download after {max_retries} attempts") raise except Exception as e: print(f"[ERROR] Unexpected error during download: {str(e)}") raise def extract_tar_file(tar_path: str, extract_dir: str) -> None: """Extract tar file with error handling""" print(f"[~] Extracting {tar_path} to {extract_dir}") try: import tarfile with tarfile.open(tar_path, "r") as tar: tar.extractall(path=extract_dir) print(f"[+] Extraction completed") # Remove tar file to save space os.remove(tar_path) print(f"[+] Cleaned up {tar_path}") except Exception as e: print(f"[ERROR] Failed to extract {tar_path}: {str(e)}") raise def download_weights(url: str, dest: str) -> None: """ Download and extract model weights with robust error handling. Args: url: URL to download from dest: Destination path (should end with .tar for tar files) """ start_time = time.time() try: # Determine if this is a tar file that needs extraction is_tar = dest.endswith(".tar") extract_dir = dest.replace(".tar", "") if is_tar else dest # Skip if already extracted if is_tar and os.path.exists(extract_dir) and os.path.isdir(extract_dir): print(f"[~] Model already exists at {extract_dir}, skipping download") return # Download the file download_with_resume(url, dest) # Extract if it's a tar file if is_tar: extract_tar_file(dest, os.path.dirname(dest)) total_time = time.time() - start_time print(f"[+] Model setup completed in {total_time:.2f} seconds") except Exception as e: print(f"[ERROR] Failed to download weights from {url}: {str(e)}") # Clean up partial files if os.path.exists(dest): os.remove(dest) raise 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] # Decode and save with open(output_path, "wb") as f: f.write(base64.b64decode(base64_string)) return output_path except Exception as e: raise ValueError(f"Failed to decode base64 file: {str(e)}") 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)}") 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}/" ) # 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 = [ "chinese-wav2vec2-base.tar", # Download smaller model first "MeiGen-MultiTalk.tar", "Wan2.1-I2V-14B-480P.tar", # Download largest model last ] print("[~] Starting model downloads...") for model_file in model_files: url = self.BASE_URL + model_file dest_path = os.path.join(self.MODEL_CACHE, model_file) 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)], ) # 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 ) # 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") if vram_gb > 40: # High VRAM setup print( "🚀 High-VRAM detected: Enabling maximum performance optimizations" ) # 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 ) 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]: """ 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: # Parse input data inputs = data.get("inputs", {}) parameters = data.get("parameters", {}) # 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)" ) # 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 ) 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", ) # 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, ) # 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: print(f"❌ Error during inference: {str(e)}") return {"error": str(e), "type": type(e).__name__}