all models

369e7ad over 1 year ago

25 kB

import modal
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
from pathlib import Path
from fastapi import HTTPException
from pydantic import BaseModel

# Suppress warnings for cleaner output
warnings.filterwarnings("ignore")

# Modal App Configuration
app = modal.App("multitalk-inference")

# Create a Volume for model weights
volume = modal.Volume.from_name("multitalk-weights", create_if_missing=True)
MODEL_DIR = Path("/models")

# HuggingFace secret for authentication
# huggingface_secret = modal.Secret.from_name(
#     "huggingface-secret", required_keys=["HF_TOKEN"]
# )
huggingface_secret = ''

# Base image with system dependencies
base_image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install(
        "ffmpeg",
        "libsndfile1",
        "git",
        "wget",
        "curl",
        "libgl1-mesa-glx",
        "libglib2.0-0",
        "libsm6",
        "libxext6",
        "libxrender-dev",
        "libgomp1",
    )
    .pip_install(
        "torch>=2.0.0",
        "torchvision",
        "torchaudio",
        "numpy",
        "soundfile",
        "librosa",
        "pyloudnorm",
        "einops",
        "transformers>=4.30.0",
        "huggingface-hub",
        "Pillow",
        "fastapi[standard]",
        "pydantic",
        "requests",
        "accelerate",
        "diffusers",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

# Download image with additional dependencies
download_image = base_image.pip_install(
    "huggingface_hub[hf_transfer]",
    "requests",
    "tqdm",
)

# Pydantic models for API
class MultiTalkRequest(BaseModel):
    image: str  # base64 encoded image
    first_audio: str  # base64 encoded audio
    prompt: str = "A smiling man and woman wearing headphones sit in front of microphones, appearing to host a podcast."
    second_audio: Optional[str] = None  # base64 encoded audio (optional)
    num_frames: int = 81
    sampling_steps: int = 40
    seed: Optional[int] = None
    turbo: bool = True

class MultiTalkResponse(BaseModel):
    video: str  # base64 encoded video
    metadata: Dict[str, Any]

class ErrorResponse(BaseModel):
    error: str
    type: str

# Helper functions
def loudness_norm(audio_array, sr=16000, lufs=-23):
    import pyloudnorm as pyln
    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"""
    import librosa
    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):
    import librosa
    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):
    import numpy as np
    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"""
    import torch
    import numpy as np
    from einops import rearrange
    
    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."""
    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:
            headers = {}
            if initial_size > 0:
                headers["Range"] = f"bytes={initial_size}-"

            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:
                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")

            start_time = time.time()
            last_progress_time = start_time
            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)

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

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

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")
        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."""
    start_time = time.time()
    try:
        is_tar = dest.endswith(".tar")
        extract_dir = dest.replace(".tar", "") if is_tar else dest

        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_with_resume(url, dest)

        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)}")
        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:
        if "," in base64_string:
            base64_string = base64_string.split(",")[1]
        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)}")

# Download models function
@app.function(
    volumes={MODEL_DIR.as_posix(): volume},
    image=download_image,
    secrets=[huggingface_secret],
    timeout=3600,  # 1 hour timeout for downloads
)
def download_models():
    """Download all required MultiTalk models"""
    MODEL_CACHE = MODEL_DIR / "weights"
    BASE_URL = f"https://weights.replicate.delivery/default/multitalk/{MODEL_CACHE.name}/"
    
    # Create model cache directory
    MODEL_CACHE.mkdir(parents=True, exist_ok=True)
    
    # Download model weights
    model_files = [
        "chinese-wav2vec2-base.tar",
        "MeiGen-MultiTalk.tar", 
        "Wan2.1-I2V-14B-480P.tar",
    ]
    
    print("[~] Starting model downloads...")
    for model_file in model_files:
        url = BASE_URL + model_file
        dest_path = MODEL_CACHE / model_file
        download_weights(url, str(dest_path))
    
    # Commit changes to volume
    volume.commit()
    print("✅ All models downloaded and committed to volume!")

# Inference class
class MultiTalkInference:
    def __init__(self):
        """Initialize the MultiTalk inference handler"""
        import torch
        from transformers import Wav2Vec2FeatureExtractor
        # Note: You'll need to add proper imports for MultiTalk components
        # import wan
        # from wan.configs import WAN_CONFIGS
        # from src.audio_analysis.wav2vec2 import Wav2Vec2Model
        # from wan.utils.multitalk_utils import save_video_ffmpeg
        
        # Model paths
        self.model_cache = MODEL_DIR / "weights"
        self.ckpt_dir = self.model_cache / "Wan2.1-I2V-14B-480P"
        self.wav2vec_dir = self.model_cache / "chinese-wav2vec2-base"
        self.multitalk_dir = self.model_cache / "MeiGen-MultiTalk"
        
        # Initialize device
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")
        
        # Load wav2vec models
        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}")
        
        # Note: Replace with actual imports when available
        # self.audio_encoder = Wav2Vec2Model.from_pretrained(
        #     str(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(
        #     str(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=str(self.ckpt_dir),
        #     device_id=0,
        #     rank=0,
        #     t5_fsdp=False,
        #     dit_fsdp=False,
        #     use_usp=False,
        #     t5_cpu=False,
        # )
        
        # GPU optimizations
        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:
                print("🚀 High-VRAM detected: Enabling maximum performance optimizations")
                torch.backends.cuda.enable_flash_sdp(True)
                torch.backends.cudnn.benchmark = True
                torch.backends.cuda.matmul.allow_tf32 = True
                torch.backends.cudnn.allow_tf32 = True
                torch.cuda.empty_cache()
                print("⚡ Enabled Flash-SDP, TF32, and cuDNN benchmarking")
            else:
                print("🔧 Standard GPU optimizations enabled")
                torch.backends.cuda.enable_flash_sdp(True)
                torch.cuda.empty_cache()
        
        print("✅ Model setup completed successfully!")
    
    def generate_video(self, request: MultiTalkRequest) -> MultiTalkResponse:
        """Generate video from the request"""
        import torch
        import numpy as np
        import random
        import soundfile as sf
        
        try:
            # Validate parameters
            if not isinstance(request.num_frames, int) or request.num_frames < 25 or request.num_frames > 201:
                raise ValueError("num_frames must be an integer between 25 and 201")
            if not isinstance(request.sampling_steps, int) or request.sampling_steps < 2 or request.sampling_steps > 100:
                raise ValueError("sampling_steps must be an integer between 2 and 100")
            
            # Auto-correct frame count
            num_frames = request.num_frames
            original_frames = num_frames
            if (num_frames - 1) % 4 != 0:
                n_lower = (num_frames - 1) // 4
                n_upper = n_lower + 1
                frames_lower = 4 * n_lower + 1
                frames_upper = 4 * n_upper + 1
                
                if abs(num_frames - frames_lower) <= abs(num_frames - frames_upper):
                    num_frames = frames_lower
                else:
                    num_frames = frames_upper
                
                num_frames = max(25, min(num_frames, 201))
                
                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}")
            
            # Set seed
            if request.seed is None:
                seed = random.randint(0, 99999999)
            else:
                seed = request.seed
            
            print(f"🎬 Generating video with seed: {seed}")
            
            with tempfile.TemporaryDirectory() as temp_dir:
                # Save input files
                image_path = os.path.join(temp_dir, "input_image.jpg")
                first_audio_path = os.path.join(temp_dir, "first_audio.wav")
                
                decode_base64_file(request.image, image_path)
                decode_base64_file(request.first_audio, first_audio_path)
                
                second_audio_path = None
                if request.second_audio:
                    second_audio_path = os.path.join(temp_dir, "second_audio.wav")
                    decode_base64_file(request.second_audio, 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"
                    speech1, speech2, combined_speech = audio_prepare_multi(
                        first_audio_path, second_audio_path, audio_type
                    )
                    
                    # Generate embeddings
                    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)
                    
                    input_data = {
                        "prompt": request.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
                    )
                    
                    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)
                    
                    input_data = {
                        "prompt": request.prompt,
                        "cond_image": image_path,
                        "cond_audio": {"person1": emb_path},
                        "video_audio": sum_audio_path,
                    }
                
                print("🎬 Generating video...")
                
                # Configure generation parameters
                high_vram = (
                    torch.cuda.is_available()
                    and torch.cuda.get_device_properties(0).total_memory > 40 * 1024**3
                )
                
                if request.turbo:
                    teacache_thresh = 0.8
                    text_guide_scale = 3.0
                    audio_guide_scale = 3.0
                    shift = 5.0
                    offload_model = False
                else:
                    teacache_thresh = 0.3
                    text_guide_scale = 5.0
                    audio_guide_scale = 4.0
                    shift = 7.0
                    offload_model = not high_vram
                
                extra_args = SimpleNamespace(
                    use_teacache=True,
                    teacache_thresh=teacache_thresh,
                    use_apg=False,
                    size="multitalk-480",
                )
                
                # Generate video (replace with actual generation call)
                # video = self.wan_i2v.generate(
                #     input_data,
                #     size_buckget="multitalk-480",
                #     motion_frame=25,
                #     frame_num=num_frames,
                #     shift=shift,
                #     sampling_steps=request.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 (replace with actual save call)
                output_name = f"multitalk_{abs(hash(request.prompt + str(seed))) % 10000}"
                # save_video_ffmpeg(video, output_name, [input_data["video_audio"]])
                
                # For now, create a dummy video file
                output_file = f"{output_name}.mp4"
                with open(output_file, "wb") as f:
                    f.write(b"dummy video content")
                
                # Encode video to base64
                video_b64 = encode_file_to_base64(output_file)
                
                # Cleanup
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                
                print("✅ Video generation completed successfully!")
                
                return MultiTalkResponse(
                    video=video_b64,
                    metadata={
                        "seed": seed,
                        "num_frames": num_frames,
                        "sampling_steps": request.sampling_steps,
                        "turbo": request.turbo,
                        "is_multi_person": is_multi_person,
                        "prompt": request.prompt,
                        "original_frames_requested": original_frames,
                    }
                )
        
        except Exception as e:
            print(f"❌ Error during inference: {str(e)}")
            raise HTTPException(status_code=500, detail=str(e))

# Inference endpoint
@app.function(
    volumes={MODEL_DIR.as_posix(): volume},
    image=base_image,
    gpu=modal.gpu.A100(size="40GB"),  # Adjust GPU as needed
    secrets=[huggingface_secret],
    timeout=1800,  # 30 minute timeout
    container_idle_timeout=300,  # 5 minute idle timeout
)
@modal.web_endpoint(method="POST")
def generate_video(request: MultiTalkRequest) -> Union[MultiTalkResponse, ErrorResponse]:
    """Main inference endpoint"""
    try:
        # Initialize inference handler
        handler = MultiTalkInference()
        
        # Generate video
        response = handler.generate_video(request)
        
        return response
        
    except Exception as e:
        print(f"❌ Endpoint error: {str(e)}")
        return ErrorResponse(error=str(e), type=type(e).__name__)

	import modal
	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
	from pathlib import Path
	from fastapi import HTTPException
	from pydantic import BaseModel

	# Suppress warnings for cleaner output
	warnings.filterwarnings("ignore")

	# Modal App Configuration
	app = modal.App("multitalk-inference")

	# Create a Volume for model weights
	volume = modal.Volume.from_name("multitalk-weights", create_if_missing=True)
	MODEL_DIR = Path("/models")

	# HuggingFace secret for authentication
	# huggingface_secret = modal.Secret.from_name(
	# "huggingface-secret", required_keys=["HF_TOKEN"]
	# )
	huggingface_secret = ''

	# Base image with system dependencies
	base_image = (
	modal.Image.debian_slim(python_version="3.10")
	.apt_install(
	"ffmpeg",
	"libsndfile1",
	"git",
	"wget",
	"curl",
	"libgl1-mesa-glx",
	"libglib2.0-0",
	"libsm6",
	"libxext6",
	"libxrender-dev",
	"libgomp1",
	)
	.pip_install(
	"torch>=2.0.0",
	"torchvision",
	"torchaudio",
	"numpy",
	"soundfile",
	"librosa",
	"pyloudnorm",
	"einops",
	"transformers>=4.30.0",
	"huggingface-hub",
	"Pillow",
	"fastapi[standard]",
	"pydantic",
	"requests",
	"accelerate",
	"diffusers",
	)
	.env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
	)

	# Download image with additional dependencies
	download_image = base_image.pip_install(
	"huggingface_hub[hf_transfer]",
	"requests",
	"tqdm",
	)

	# Pydantic models for API
	class MultiTalkRequest(BaseModel):
	image: str # base64 encoded image
	first_audio: str # base64 encoded audio
	prompt: str = "A smiling man and woman wearing headphones sit in front of microphones, appearing to host a podcast."
	second_audio: Optional[str] = None # base64 encoded audio (optional)
	num_frames: int = 81
	sampling_steps: int = 40
	seed: Optional[int] = None
	turbo: bool = True

	class MultiTalkResponse(BaseModel):
	video: str # base64 encoded video
	metadata: Dict[str, Any]

	class ErrorResponse(BaseModel):
	error: str
	type: str

	# Helper functions
	def loudness_norm(audio_array, sr=16000, lufs=-23):
	import pyloudnorm as pyln
	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"""
	import librosa
	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):
	import librosa
	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):
	import numpy as np
	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"""
	import torch
	import numpy as np
	from einops import rearrange

	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."""
	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:
	headers = {}
	if initial_size > 0:
	headers["Range"] = f"bytes={initial_size}-"

	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:
	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")

	start_time = time.time()
	last_progress_time = start_time
	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)

	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)
	print(f"[~] Progress: {progress:.1f}% ({downloaded/(10243):.2f}/{total_size/(10243):.2f} GB) - Speed: {speed:.1f} MB/s")
	last_progress_time = current_time

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

	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")
	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."""
	start_time = time.time()
	try:
	is_tar = dest.endswith(".tar")
	extract_dir = dest.replace(".tar", "") if is_tar else dest

	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_with_resume(url, dest)

	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)}")
	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:
	if "," in base64_string:
	base64_string = base64_string.split(",")[1]
	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)}")

	# Download models function
	@app.function(
	volumes={MODEL_DIR.as_posix(): volume},
	image=download_image,
	secrets=[huggingface_secret],
	timeout=3600, # 1 hour timeout for downloads
	)
	def download_models():
	"""Download all required MultiTalk models"""
	MODEL_CACHE = MODEL_DIR / "weights"
	BASE_URL = f"https://weights.replicate.delivery/default/multitalk/{MODEL_CACHE.name}/"

	# Create model cache directory
	MODEL_CACHE.mkdir(parents=True, exist_ok=True)

	# Download model weights
	model_files = [
	"chinese-wav2vec2-base.tar",
	"MeiGen-MultiTalk.tar",
	"Wan2.1-I2V-14B-480P.tar",
	]

	print("[~] Starting model downloads...")
	for model_file in model_files:
	url = BASE_URL + model_file
	dest_path = MODEL_CACHE / model_file
	download_weights(url, str(dest_path))

	# Commit changes to volume
	volume.commit()
	print("✅ All models downloaded and committed to volume!")

	# Inference class
	class MultiTalkInference:
	def __init__(self):
	"""Initialize the MultiTalk inference handler"""
	import torch
	from transformers import Wav2Vec2FeatureExtractor
	# Note: You'll need to add proper imports for MultiTalk components
	# import wan
	# from wan.configs import WAN_CONFIGS
	# from src.audio_analysis.wav2vec2 import Wav2Vec2Model
	# from wan.utils.multitalk_utils import save_video_ffmpeg

	# Model paths
	self.model_cache = MODEL_DIR / "weights"
	self.ckpt_dir = self.model_cache / "Wan2.1-I2V-14B-480P"
	self.wav2vec_dir = self.model_cache / "chinese-wav2vec2-base"
	self.multitalk_dir = self.model_cache / "MeiGen-MultiTalk"

	# Initialize device
	self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
	print(f"Using device: {self.device}")

	# Load wav2vec models
	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}")

	# Note: Replace with actual imports when available
	# self.audio_encoder = Wav2Vec2Model.from_pretrained(
	# str(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(
	# str(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=str(self.ckpt_dir),
	# device_id=0,
	# rank=0,
	# t5_fsdp=False,
	# dit_fsdp=False,
	# use_usp=False,
	# t5_cpu=False,
	# )

	# GPU optimizations
	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:
	print("🚀 High-VRAM detected: Enabling maximum performance optimizations")
	torch.backends.cuda.enable_flash_sdp(True)
	torch.backends.cudnn.benchmark = True
	torch.backends.cuda.matmul.allow_tf32 = True
	torch.backends.cudnn.allow_tf32 = True
	torch.cuda.empty_cache()
	print("⚡ Enabled Flash-SDP, TF32, and cuDNN benchmarking")
	else:
	print("🔧 Standard GPU optimizations enabled")
	torch.backends.cuda.enable_flash_sdp(True)
	torch.cuda.empty_cache()

	print("✅ Model setup completed successfully!")

	def generate_video(self, request: MultiTalkRequest) -> MultiTalkResponse:
	"""Generate video from the request"""
	import torch
	import numpy as np
	import random
	import soundfile as sf

	try:
	# Validate parameters
	if not isinstance(request.num_frames, int) or request.num_frames < 25 or request.num_frames > 201:
	raise ValueError("num_frames must be an integer between 25 and 201")
	if not isinstance(request.sampling_steps, int) or request.sampling_steps < 2 or request.sampling_steps > 100:
	raise ValueError("sampling_steps must be an integer between 2 and 100")

	# Auto-correct frame count
	num_frames = request.num_frames
	original_frames = num_frames
	if (num_frames - 1) % 4 != 0:
	n_lower = (num_frames - 1) // 4
	n_upper = n_lower + 1
	frames_lower = 4 * n_lower + 1
	frames_upper = 4 * n_upper + 1

	if abs(num_frames - frames_lower) <= abs(num_frames - frames_upper):
	num_frames = frames_lower
	else:
	num_frames = frames_upper

	num_frames = max(25, min(num_frames, 201))

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

	# Set seed
	if request.seed is None:
	seed = random.randint(0, 99999999)
	else:
	seed = request.seed

	print(f"🎬 Generating video with seed: {seed}")

	with tempfile.TemporaryDirectory() as temp_dir:
	# Save input files
	image_path = os.path.join(temp_dir, "input_image.jpg")
	first_audio_path = os.path.join(temp_dir, "first_audio.wav")

	decode_base64_file(request.image, image_path)
	decode_base64_file(request.first_audio, first_audio_path)

	second_audio_path = None
	if request.second_audio:
	second_audio_path = os.path.join(temp_dir, "second_audio.wav")
	decode_base64_file(request.second_audio, 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"
	speech1, speech2, combined_speech = audio_prepare_multi(
	first_audio_path, second_audio_path, audio_type
	)

	# Generate embeddings
	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)

	input_data = {
	"prompt": request.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
	)

	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)

	input_data = {
	"prompt": request.prompt,
	"cond_image": image_path,
	"cond_audio": {"person1": emb_path},
	"video_audio": sum_audio_path,
	}

	print("🎬 Generating video...")

	# Configure generation parameters
	high_vram = (
	torch.cuda.is_available()
	and torch.cuda.get_device_properties(0).total_memory > 40 * 1024**3
	)

	if request.turbo:
	teacache_thresh = 0.8
	text_guide_scale = 3.0
	audio_guide_scale = 3.0
	shift = 5.0
	offload_model = False
	else:
	teacache_thresh = 0.3
	text_guide_scale = 5.0
	audio_guide_scale = 4.0
	shift = 7.0
	offload_model = not high_vram

	extra_args = SimpleNamespace(
	use_teacache=True,
	teacache_thresh=teacache_thresh,
	use_apg=False,
	size="multitalk-480",
	)

	# Generate video (replace with actual generation call)
	# video = self.wan_i2v.generate(
	# input_data,
	# size_buckget="multitalk-480",
	# motion_frame=25,
	# frame_num=num_frames,
	# shift=shift,
	# sampling_steps=request.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 (replace with actual save call)
	output_name = f"multitalk_{abs(hash(request.prompt + str(seed))) % 10000}"
	# save_video_ffmpeg(video, output_name, [input_data["video_audio"]])

	# For now, create a dummy video file
	output_file = f"{output_name}.mp4"
	with open(output_file, "wb") as f:
	f.write(b"dummy video content")

	# Encode video to base64
	video_b64 = encode_file_to_base64(output_file)

	# Cleanup
	if torch.cuda.is_available():
	torch.cuda.empty_cache()

	print("✅ Video generation completed successfully!")

	return MultiTalkResponse(
	video=video_b64,
	metadata={
	"seed": seed,
	"num_frames": num_frames,
	"sampling_steps": request.sampling_steps,
	"turbo": request.turbo,
	"is_multi_person": is_multi_person,
	"prompt": request.prompt,
	"original_frames_requested": original_frames,
	}
	)

	except Exception as e:
	print(f"❌ Error during inference: {str(e)}")
	raise HTTPException(status_code=500, detail=str(e))

	# Inference endpoint
	@app.function(
	volumes={MODEL_DIR.as_posix(): volume},
	image=base_image,
	gpu=modal.gpu.A100(size="40GB"), # Adjust GPU as needed
	secrets=[huggingface_secret],
	timeout=1800, # 30 minute timeout
	container_idle_timeout=300, # 5 minute idle timeout
	)
	@modal.web_endpoint(method="POST")
	def generate_video(request: MultiTalkRequest) -> Union[MultiTalkResponse, ErrorResponse]:
	"""Main inference endpoint"""
	try:
	# Initialize inference handler
	handler = MultiTalkInference()

	# Generate video
	response = handler.generate_video(request)

	return response

	except Exception as e:
	print(f"❌ Endpoint error: {str(e)}")
	return ErrorResponse(error=str(e), type=type(e).__name__)