Initial Commit

ec67cd4 over 1 year ago

21.4 kB

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