# # 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__}
# # 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]],
# ]
# )