Blame for handler.py

# # 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__}