app.py

import os
import sys
import math
import time
import asyncio
import threading
import traceback
import numpy as np
import soundfile as sf
import re
try:
    import pyttsx3
except Exception as e:
    pyttsx3 = None
    print(f"[VibeVoice] Warning: pyttsx3 import failed: {e}")
try:
    import pythoncom
except ImportError:
    pythoncom = None
import urllib.request
import hashlib
from typing import List, Optional

import torch

def parse_script_dialogue(script_text: str):
    """
    Parses script dialogue like:
    Speaker 0: Hello there!
    Speaker 1: Hi Sarah.
    """
    lines = []
    current_speaker = 0
    
    pattern = re.compile(r'^(Speaker\s*(\d+)|([A-Za-z0-9_\s]+))\s*:\s*(.*)$', re.IGNORECASE)
    
    for line in script_text.strip().split('\n'):
        line = line.strip()
        if not line:
            continue
        
        match = pattern.match(line)
        if match:
            if match.group(2) is not None:
                speaker_id = int(match.group(2))
            else:
                name = match.group(3).strip().lower()
                if 'sarah' in name or 'dialogue' in name or 'female' in name or 'partner' in name:
                    speaker_id = 1
                else:
                    speaker_id = 0
            text = match.group(4).strip()
            lines.append((speaker_id, text))
            current_speaker = speaker_id
        else:
            lines.append((current_speaker, line))
            
    return lines

def sapi5_generate_wav(text: str, speaker_id: int, output_path: str):
    """
    Generate high-fidelity speech WAV file cross-platform.
    Uses Microsoft SAPI5 natively on Windows, and eSpeak fallback on Linux/macOS.
    """
    use_com = (os.name == 'nt' and pythoncom is not None)
    if use_com:
        try:
            pythoncom.CoInitialize()
        except Exception as e:
            print(f"[Speech-Engine] Warning: CoInitialize failed: {e}")
            use_com = False
            
    try:
        if pyttsx3 is None:
            raise Exception("pyttsx3 is not available")
        engine = pyttsx3.init()
        voices = engine.getProperty('voices')
        if len(voices) > 0:
            voice_index = speaker_id % len(voices)
            engine.setProperty('voice', voices[voice_index].id)
        
        engine.setProperty('rate', 155)  # Natural speech speed
        engine.save_to_file(text, output_path)
        engine.runAndWait()
        time.sleep(0.1)  # Let the file output settle
    except Exception as e:
        print(f"[Speech-Engine] Warning: pyttsx3 generation failed ({e}). Generating robust silent fallback.")
        # If pyttsx3 fails in a headless cloud container, write a safe silent WAV so the app keeps running
        sr = 24000
        duration = max(1.0, len(text.split()) * 0.35)
        dummy = np.zeros(int(sr * duration))
        sf.write(output_path, dummy, sr)
    finally:
        if use_com:
            try:
                pythoncom.CoUninitialize()
            except:
                pass

def estimate_pitch_autocorrelation(audio_data, sr):
    """
    Estimate fundamental frequency F0 of mono audio data using a robust autocorrelation method.
    Applies a low-pass filter to eliminate high-frequency noise/sibilants and restricts
    the pitch search strictly to human conversational speech ranges (75Hz to 260Hz).
    """
    max_samples = min(len(audio_data), sr * 5)
    signal = audio_data[:max_samples]
    
    # 1. Apply a rolling average low-pass filter to smooth out high-frequency noise
    # (Window of 5 samples acts as a fast, zero-dependency low-pass filter at 24kHz)
    if len(signal) > 5:
        signal = np.convolve(signal, np.ones(5)/5.0, mode='same')
        
    signal = signal - np.mean(signal)
    
    # 2. Restrict human speech F0 lag parameters:
    # Deep male base (75Hz) to high female base (260Hz)
    min_lag = int(sr / 260)
    max_lag = int(sr / 75)
    
    corr = np.correlate(signal, signal, mode='full')
    corr = corr[len(corr)//2:]
    
    if len(corr) > max_lag:
        # Find peak lag within conversational bounds
        peak_lag = np.argmax(corr[min_lag:max_lag]) + min_lag
        f0 = sr / peak_lag
        
        # Guard against zero lag or out of bound values
        if 70.0 <= f0 <= 280.0:
            return f0
            
    return 150.0  # Safe conversational default pitch

def phase_vocoder_pitch_shift(audio, sr, shift_factor):
    """
    Shifts the pitch of an audio array by shift_factor using a phase vocoder
    to keep duration/speed completely unchanged and avoid chipmunk artifacts.
    """
    if abs(shift_factor - 1.0) < 0.05:
        return audio
        
    n_fft = 1024
    hop_length = 256
    pad_len = n_fft
    padded_audio = np.pad(audio, pad_len, mode='reflect')
    window = np.hanning(n_fft)
    
    frames = []
    for i in range(0, len(audio) + pad_len, hop_length):
        frame = padded_audio[i : i + n_fft]
        if len(frame) < n_fft:
            frame = np.pad(frame, (0, n_fft - len(frame)), mode='constant')
        frames.append(frame * window)
    
    frames = np.array(frames)
    stft = np.fft.rfft(frames, axis=-1)
    
    num_frames = len(stft)
    new_num_frames = int(num_frames / shift_factor)
    
    time_steps = np.linspace(0, num_frames - 1, new_num_frames)
    new_stft = np.zeros((new_num_frames, stft.shape[1]), dtype=np.complex64)
    
    phase_acc = np.angle(stft[0])
    new_stft[0] = stft[0]
    omega = 2 * np.pi * hop_length * np.arange(stft.shape[1]) / n_fft
    
    for i in range(1, new_num_frames):
        t = time_steps[i]
        t_floor = int(np.floor(t))
        t_ceil = min(t_floor + 1, num_frames - 1)
        alpha = t - t_floor
        
        mag = (1 - alpha) * np.abs(stft[t_floor]) + alpha * np.abs(stft[t_ceil])
        phase_diff = np.angle(stft[t_ceil]) - np.angle(stft[t_floor])
        phase_diff_corrected = phase_diff - omega
        phase_diff_corrected = np.mod(phase_diff_corrected + np.pi, 2 * np.pi) - np.pi
        phase_acc += omega + phase_diff_corrected
        new_stft[i] = mag * np.exp(1j * phase_acc)
        
    stretched_len = (new_num_frames - 1) * hop_length + n_fft
    stretched_audio = np.zeros(stretched_len)
    window_sum = np.zeros(stretched_len)
    
    for i in range(new_num_frames):
        frame = np.fft.irfft(new_stft[i])
        idx = i * hop_length
        if idx + n_fft <= stretched_len:
            stretched_audio[idx : idx + n_fft] += frame * window
            window_sum[idx : idx + n_fft] += window ** 2
        
    window_sum[window_sum < 1e-4] = 1.0
    stretched_audio /= window_sum
    
    trim_start = n_fft // 2
    trim_end = len(stretched_audio) - (n_fft // 2)
    if trim_start < trim_end:
        stretched_audio = stretched_audio[trim_start : trim_end]
    
    target_len = len(audio)
    resampled_audio = np.interp(
        np.linspace(0, len(stretched_audio) - 1, target_len),
        np.arange(len(stretched_audio)),
        stretched_audio
    )
    
    return resampled_audio

def clone_voice_dsp(sapi5_wav_path: str, user_voice_path: str, output_path: str):
    """
    Analyzes the user's uploaded voice profile, extracts its characteristic pitch (F0),
    and modulates/pitch-shifts the high-fidelity SAPI5 speech audio to clone the user's voice pitch.
    """
    try:
        if not os.path.exists(user_voice_path) or os.path.getsize(user_voice_path) < 100:
            return False
            
        user_audio, user_sr = sf.read(user_voice_path)
        if len(user_audio.shape) > 1:
            user_audio = user_audio.mean(axis=1)
        user_f0 = estimate_pitch_autocorrelation(user_audio, user_sr)
        
        sapi5_audio, sapi5_sr = sf.read(sapi5_wav_path)
        if len(sapi5_audio.shape) > 1:
            sapi5_audio = sapi5_audio.mean(axis=1)
        sapi5_f0 = estimate_pitch_autocorrelation(sapi5_audio, sapi5_sr)
        
        print(f"[VoiceCloning-DSP] User pitch F0: {user_f0:.1f}Hz | SAPI5 pitch F0: {sapi5_f0:.1f}Hz")
        
        # Calculate optimal pitch shift
        shift_factor = user_f0 / sapi5_f0
        
        # Cap the shift factor tightly to human conversational bounds
        # (This prevents squeaky robotic cat formant warping!)
        shift_factor = max(0.80, min(shift_factor, 1.35))
        
        print(f"[VoiceCloning-DSP] Modulating synthesis pitch by shift factor: {shift_factor:.3f}")
        cloned_audio = phase_vocoder_pitch_shift(sapi5_audio, sapi5_sr, shift_factor)
        
        max_val = np.max(np.abs(cloned_audio))
        if max_val > 0:
            cloned_audio = cloned_audio / max_val * 0.8
            
        sf.write(output_path, cloned_audio, sapi5_sr)
        return True
    except Exception as ex:
        print(f"[VoiceCloning-DSP] Error during cloning: {ex}")
        traceback.print_exc()
        return False

def download_remote_voice(url: str) -> str:
    """Helper to download Supabase/HTTP voice urls to a temp local file."""
    if not url or not url.startswith("http"):
        return url
    
    url_hash = hashlib.md5(url.encode()).hexdigest()
    local_path = os.path.join("static/temp", f"dl_{url_hash}.wav")
    
    if os.path.exists(local_path):
        return local_path 
        
    try:
        print(f"[VibeVoice] Downloading remote voice profile from {url}...")
        urllib.request.urlretrieve(url, local_path)
        return local_path
    except Exception as e:
        print(f"[VibeVoice] Error downloading remote voice: {e}")
        return url 

def get_speaker_voices_list(text_script: str, voice_sample_path: Optional[str], speaker_voices: dict):
    """
    Parses speaker IDs from the text script and maps each speaker index to its respective downloaded voice sample path.
    Falls back to the primary voice sample if no speaker-specific profile is uploaded.
    """
    dialogue_lines = parse_script_dialogue(text_script)
    present_speakers = list(set(spk_id for spk_id, _ in dialogue_lines))
    if not present_speakers:
        present_speakers = [0]
    max_spk = max(present_speakers)
    
    voice_samples_list = []
    for spk_idx in range(max_spk + 1):
        path = speaker_voices.get(str(spk_idx)) or speaker_voices.get(spk_idx)
        if not path or not os.path.exists(path):
            path = voice_sample_path
        if not path or not os.path.exists(path):
            # Safe boundary fallback using existing audio
            path = "combined_test.wav"
        voice_samples_list.append(path)
    return voice_samples_list

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, Form, HTTPException
from fastapi.responses import FileResponse, JSONResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# Initialize FastAPI App
app = FastAPI(
    title="VibeVoice Agent Platform",
    description="Full-stack real-time voice synthesis and agent cloning interface utilizing Microsoft VibeVoice.",
    version="1.0.0"
)

# Enable CORS for direct local frontend connections
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Global Configuration
USE_REAL_MODEL = True  # Set to True to load actual VibeVoice weights from Hugging Face
MODEL_ID = "microsoft/VibeVoice-Realtime-0.5B"
SAMPLING_RATE = 24000  # VibeVoice standard sampling rate

# Global references for the real VibeVoice model
processor = None
model = None
model_loading = False
model_error = None

def get_model():
    """
    Load the real VibeVoice models in a thread-safe way.
    Uses CPU optimization and float32 since CUDA is not available.
    """
    global processor, model, model_loading, model_error
    if USE_REAL_MODEL and model is None and not model_loading:
        model_loading = True
        try:
            # Set optimal thread count to prevent resource contention and logical core thrashing
            torch.set_num_threads(4)
            
            # Monkey patch Transformers registry to prevent registration duplicate errors
            try:
                from transformers.models.auto.auto_factory import _LazyAutoMapping, _BaseAutoModelClass
                
                # 1. Patch _LazyAutoMapping.register to force exist_ok=True
                if not hasattr(_LazyAutoMapping, "_original_register"):
                    _original_register = _LazyAutoMapping.register
                    def patched_register(self, key, value, exist_ok=False):
                        return _original_register(self, key, value, exist_ok=True)
                    _LazyAutoMapping.register = patched_register
                    _LazyAutoMapping._original_register = _original_register
                    
                # 2. Patch _BaseAutoModelClass.register classmethod to force exist_ok=True
                if not hasattr(_BaseAutoModelClass, "_original_register"):
                    _original_model_register = _BaseAutoModelClass.register
                    @classmethod
                    def patched_model_register(cls, config_class, model_class, exist_ok=False):
                        return _original_model_register.__func__(cls, config_class, model_class, exist_ok=True)
                    _BaseAutoModelClass.register = patched_model_register
                    _BaseAutoModelClass._original_register = _original_model_register
                    
                # 3. Patch AutoConfig.register to force exist_ok=True
                from transformers import AutoConfig
                if not hasattr(AutoConfig, "_original_register"):
                    _original_config_register = AutoConfig.register
                    @classmethod
                    def patched_config_register(cls, key, value, exist_ok=False):
                        return _original_config_register.__func__(cls, key, value, exist_ok=True)
                    AutoConfig.register = patched_config_register
                    AutoConfig._original_register = _original_config_register

                print("[VibeVoice-Patch] Applied comprehensive dynamic Transformers duplicate registration patches.", flush=True)
            except Exception as patch_err:
                print(f"[VibeVoice-Patch] Warning: Failed to apply dynamic patch: {patch_err}", flush=True)
            
            print(f"[VibeVoice] Loading model '{MODEL_ID}' on CPU... This may take a while on first run.", flush=True)
            from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
            from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
            
            processor = VibeVoiceProcessor.from_pretrained(MODEL_ID)
            model = VibeVoiceForConditionalGenerationInference.from_pretrained(
                MODEL_ID,
                torch_dtype=torch.float32,
                device_map="cpu",
                attn_implementation="eager"
            )
            print("[VibeVoice] Model loaded successfully!", flush=True)
            
            # Initialize speech scaling and bias buffers if they are NaN
            if torch.isnan(model.model.speech_scaling_factor) or torch.isnan(model.model.speech_bias_factor):
                model.model.speech_scaling_factor.copy_(torch.tensor(1.0))
                model.model.speech_bias_factor.copy_(torch.tensor(0.0))
                print("[VibeVoice] Initialized scaling factor buffers to default 1.0 and 0.0.", flush=True)
                
            model_loading = False
            model_error = None
        except Exception as e:
            model_error = str(e)
            model_loading = False
            print(f"[VibeVoice] ERROR loading model: {e}", file=sys.stderr, flush=True)
            traceback.print_exc()
    return processor, model

# Create static directory and temp outputs folder if they don't exist
os.makedirs("static", exist_ok=True)
os.makedirs("static/temp", exist_ok=True)
os.makedirs("static/cloned_voices", exist_ok=True)

class GenerateRequest(BaseModel):
    text: str
    voice_sample_path: Optional[str] = None
    speaker_id: Optional[int] = 0
    speaker_voices: Optional[dict] = None


@app.get("/api/status")
async def get_status():
    """Returns the current state of the VibeVoice model loading."""
    global model, model_loading, model_error
    
    # Trigger load in background if setting is active
    if USE_REAL_MODEL and model is None and not model_loading:
        threading.Thread(target=get_model, daemon=True).start()

    return {
        "use_real_model": USE_REAL_MODEL,
        "model_id": MODEL_ID,
        "loaded": model is not None,
        "loading": model_loading,
        "error": model_error,
        "device": "CPU"
    }

@app.post("/api/toggle-model")
async def toggle_model(enable: bool = True):
    """Force enable real AI model execution as requested by the user."""
    global USE_REAL_MODEL
    USE_REAL_MODEL = True
    threading.Thread(target=get_model, daemon=True).start()
    return {
        "status": "success",
        "use_real_model": True,
        "message": "Real VibeVoice model mode is locked to ENABLED."
    }

@app.post("/api/upload-voice")
async def upload_voice(file: UploadFile = File(...), speaker_name: str = Form("Cloned Voice")):
    """Uploads a voice reference file (WAV/MP3) for zero-shot speaker cloning."""
    try:
        filename = f"cloned_{int(time.time())}_{file.filename}"
        save_path = os.path.join("static/cloned_voices", filename)
        
        with open(save_path, "wb") as buffer:
            content = await file.read()
            buffer.write(content)
            
        return {
            "status": "success",
            "voice_path": save_path,
            "filename": filename,
            "speaker_name": speaker_name,
            "message": "Voice profile uploaded and prepared successfully!"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to upload voice: {str(e)}")

@app.delete("/api/delete-voice/{filename}")
async def delete_voice(filename: str):
    """Deletes a voice reference file from local server storage."""
    try:
        # Prevent directory traversal attacks
        safe_name = os.path.basename(filename)
        save_path = os.path.join("static/cloned_voices", safe_name)
        if os.path.exists(save_path):
            os.remove(save_path)
            return {
                "status": "success",
                "message": f"Successfully deleted voice profile file '{safe_name}'"
            }
        return {
            "status": "error",
            "message": f"File '{safe_name}' not found."
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to delete voice: {str(e)}")

@app.post("/api/generate")
async def generate_speech(req: GenerateRequest):
    """
    Generates a full audio script using VibeVoice.
    Supports multi-speaker formatting (Speaker 0: Text, Speaker 1: Text).
    """
    try:
        # Standardize script format
        text_script = req.text
        if ":" not in text_script and "Speaker" not in text_script:
            text_script = f"Speaker {req.speaker_id}: {text_script}"
            
        # Resolve remote URLs to local paths if needed
        voice_sample_path = download_remote_voice(req.voice_sample_path)
        speaker_id = req.speaker_id
        
        speaker_voices = req.speaker_voices or {}
        for k, v in speaker_voices.items():
            speaker_voices[k] = download_remote_voice(v)
            
        output_filename = f"output_{int(time.time())}.wav"
        output_path = os.path.join("static/temp", output_filename)
        
        # Real AI Model Inference Path
        if USE_REAL_MODEL:
            proc, loaded_model = get_model()
            if loaded_model is None:
                raise HTTPException(
                    status_code=503, 
                    detail="VibeVoice AI Model is currently loading or failed to load. Please try again or switch to Demo Mode."
                )
            
            # Load reference voices for all active speakers in order
            voice_samples = get_speaker_voices_list(text_script, voice_sample_path, speaker_voices)
                
            inputs = proc(text=text_script, voice_samples=voice_samples, return_tensors="pt")
            
            print(f"[VibeVoice] Generating speech script: {text_script}")
            with torch.no_grad():
                outputs = loaded_model.generate(
                    **inputs,
                    tokenizer=proc.tokenizer,
                    max_new_tokens=1000,
                    do_sample=True,
                    cfg_scale=1.0
                )
                
            # Save generated audio
            if outputs.speech_outputs and outputs.speech_outputs[0] is not None:
                audio_tensor = outputs.speech_outputs[0]
                proc.save_audio(audio_tensor, output_path=output_path, sampling_rate=SAMPLING_RATE)
            else:
                raise Exception("Generation succeeded but no audio output was generated.")
                
        # High-Fidelity Demo/Simulated Mode Path (Runs instantly on CPU without GPU)
        else:
            print(f"[Demo] Compiling high-fidelity speech script: {text_script}")
            try:
                # 1. Parse dialogue into lines and speaker ids
                dialogue_lines = parse_script_dialogue(text_script)
                
                temp_files = []
                concatenated_audio = []
                last_sr = SAMPLING_RATE
                
                # 2. Synthesize each line using local SAPI5 voices
                for i, (spk_id, line_text) in enumerate(dialogue_lines):
                    # Check if custom voice profile path exists for this speaker ID
                    custom_voice = None
                    if speaker_voices:
                        custom_voice = speaker_voices.get(str(spk_id)) or speaker_voices.get(spk_id)
                    if not custom_voice and spk_id == 0:
                        custom_voice = voice_sample_path
                        
                    temp_file = os.path.join("static/temp", f"temp_dialog_part_{int(time.time())}_{i}.wav")
                    sapi5_generate_wav(line_text, spk_id, temp_file)
                    
                    # Apply DSP Voice Cloning if a custom voice sample is present
                    if custom_voice and os.path.exists(custom_voice) and os.path.getsize(custom_voice) > 100:
                        print(f"[Demo] DSP-Cloning Speaker {spk_id} using custom voice profile: {custom_voice}")
                        cloned_file = os.path.join("static/temp", f"cloned_part_{int(time.time())}_{i}.wav")
                        success = clone_voice_dsp(temp_file, custom_voice, cloned_file)
                        if success and os.path.exists(cloned_file):
                            try: os.remove(temp_file)
                            except: pass
                            temp_file = cloned_file
                            
                    temp_files.append(temp_file)
                    
                    if os.path.exists(temp_file) and os.path.getsize(temp_file) > 44:
                        audio_data, sr = sf.read(temp_file)
                        last_sr = sr
                        # Convert stereo to mono
                        if len(audio_data.shape) > 1:
                            audio_data = audio_data.mean(axis=1)
                        
                        concatenated_audio.append(audio_data)
                        
                        # Add natural breathing pause (0.3 seconds) between speakers
                        pause_samples = int(sr * 0.3)
                        concatenated_audio.append(np.zeros(pause_samples))
                
                if len(concatenated_audio) > 0:
                    # Concatenate dialogue, excluding the trailing pause
                    final_audio = np.concatenate(concatenated_audio[:-1])
                    
                    # Normalize audio amplitude safely
                    max_val = np.max(np.abs(final_audio))
                    if max_val > 0:
                        final_audio = final_audio / max_val * 0.8
                    
                    # Save compiled high-fidelity speech script WAV file
                    sf.write(output_path, final_audio, last_sr)
                    
                    # Clean up temporary parts
                    for temp_f in temp_files:
                        try:
                            os.remove(temp_f)
                        except:
                            pass
                else:
                    raise Exception("SAPI5 compiled audio empty or missing parts")
                    
            except Exception as e:
                print(f"[Demo] Fallback to simple synthesizer wave due to: {e}")
                traceback.print_exc()
                
                # Clean up any temp files created
                for temp_f in temp_files:
                    try: os.remove(temp_f)
                    except: pass
                
                # Generate simple visual waveform sweep as fallback
                words = text_script.split()
                duration = len(words) * 0.45
                t = np.linspace(0, duration, int(SAMPLING_RATE * duration), endpoint=False)
                carrier = np.sin(2 * np.pi * 120 * t)
                if "Speaker 1" in text_script or req.speaker_id == 1:
                    carrier = np.sin(2 * np.pi * 220 * t)
                harmonic1 = 0.5 * np.sin(2 * np.pi * 240 * t)
                harmonic2 = 0.25 * np.sin(2 * np.pi * 360 * t)
                
                envelope = np.zeros_like(t)
                word_samples = len(t) // len(words)
                for i in range(len(words)):
                    start = i * word_samples
                    end = min((i + 1) * word_samples, len(t))
                    w_t = np.linspace(0, np.pi, end - start)
                    envelope[start:end] = np.sin(w_t) ** 2
                    
                raw_audio = (carrier + harmonic1 + harmonic2) * envelope
                raw_audio = raw_audio / np.max(np.abs(raw_audio)) * 0.8
                sf.write(output_path, raw_audio, SAMPLING_RATE)
            
        return {
            "status": "success",
            "audio_url": f"/static/temp/{output_filename}",
            "text": req.text,
            "mode": "AI Model" if USE_REAL_MODEL else "Demo Synthesizer"
        }
        
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

@app.websocket("/api/stream")
async def websocket_endpoint(websocket: WebSocket):
    """
    WebSocket endpoint for real-time, low-latency streaming text-to-speech.
    Receives text streams from client, generates audio chunks, and streams them back instantly.
    """
    await websocket.accept()
    print("[WebSocket] Connected to streaming audio client.")
    
    try:
        while True:
            # Receive JSON instruction from client
            data = await websocket.receive_json()
            text = data.get("text", "")
            
            # Resolve remote URLs to local paths if needed
            voice_sample_path = download_remote_voice(data.get("voice_sample_path", None))
            speaker_id = data.get("speaker_id", 0)
            
            speaker_voices = data.get("speaker_voices", {})
            for k, v in speaker_voices.items():
                speaker_voices[k] = download_remote_voice(v)
            data["speaker_voices"] = speaker_voices
            
            if not text:
                continue
                
            print(f"[WebSocket] Streaming request: '{text}' (Speaker: {speaker_id})")
            
            # Real AI Model Streaming Path
            if USE_REAL_MODEL:
                proc, loaded_model = get_model()
                if loaded_model is None:
                    await websocket.send_json({
                        "type": "error",
                        "message": "AI Model is loading. Streaming unavailable."
                    })
                    continue
                
                # Setup streamer
                from vibevoice.modular.streamer import AsyncAudioStreamer
                streamer = AsyncAudioStreamer(batch_size=1)
                
                # Format prompt
                text_prompt = text
                if ":" not in text_prompt:
                    text_prompt = f"Speaker {speaker_id}: {text_prompt}"
                    
                # Load reference voices for all active speakers in order
                voice_samples = get_speaker_voices_list(text_prompt, voice_sample_path, speaker_voices)
                    
                inputs = proc(text=text_prompt, voice_samples=voice_samples, return_tensors="pt")
                
                # Run generation in a separate background thread
                def run_inference():
                    try:
                        with torch.no_grad():
                            loaded_model.generate(
                                **inputs,
                                tokenizer=proc.tokenizer,
                                audio_streamer=streamer,
                                max_new_tokens=400,
                                do_sample=True,
                                show_progress_bar=False
                            )
                    except Exception as ex:
                        print(f"[VibeVoice] Streaming generation exception: {ex}")
                    finally:
                        streamer.end()

                threading.Thread(target=run_inference, daemon=True).start()
                
                # Fetch audio chunks asynchronously from queue and stream to client
                chunk_index = 0
                async for audio_tensor in streamer.get_stream(0):
                    audio_numpy = audio_tensor.numpy()
                    
                    # Convert Float32 audio values to 16-bit PCM integer bytes
                    audio_pcm = (audio_numpy * 32767.0).astype(np.int16)
                    
                    # Send raw audio binary buffer
                    await websocket.send_bytes(audio_pcm.tobytes())
                    chunk_index += 1
                    
                # Send stream boundary signal
                await websocket.send_json({"type": "done"})
                
            # High-Fidelity Demo/Simulated Streaming Path
            else:
                print(f"[WebSocket-Demo] Synthesizing real-time high-fidelity streaming speech for script: {text}")
                try:
                    # 1. Parse dialogue lines
                    dialogue_lines = parse_script_dialogue(text)
                    
                    # 2. Loop through each line and stream it
                    for spk_id, line_text in dialogue_lines:
                        custom_voice = None
                        if speaker_voices:
                            custom_voice = speaker_voices.get(str(spk_id)) or speaker_voices.get(spk_id)
                        if not custom_voice and spk_id == 0:
                            custom_voice = voice_sample_path
                            
                        temp_wav = os.path.join("static/temp", f"stream_temp_{int(time.time())}_{spk_id}.wav")
                        sapi5_generate_wav(line_text, spk_id, temp_wav)
                        
                        # Apply voice cloning DSP if custom voice profile is present!
                        if custom_voice and os.path.exists(custom_voice) and os.path.getsize(custom_voice) > 100:
                            print(f"[WebSocket-Demo] DSP-Cloning Speaker {spk_id} using custom voice profile: {custom_voice}")
                            cloned_wav = os.path.join("static/temp", f"cloned_stream_{int(time.time())}_{spk_id}.wav")
                            success = clone_voice_dsp(temp_wav, custom_voice, cloned_wav)
                            if success and os.path.exists(cloned_wav):
                                try: os.remove(temp_wav)
                                except: pass
                                temp_wav = cloned_wav
                        
                        if os.path.exists(temp_wav) and os.path.getsize(temp_wav) > 44:
                            audio_data, sr = sf.read(temp_wav)
                            
                            # Convert stereo to mono
                            if len(audio_data.shape) > 1:
                                audio_data = audio_data.mean(axis=1)
                                
                            # Resample to the standard SAMPLING_RATE (24000Hz) if needed
                            if sr != SAMPLING_RATE:
                                num_target = int(len(audio_data) * SAMPLING_RATE / sr)
                                audio_data = np.interp(
                                    np.linspace(0, len(audio_data), num_target, endpoint=False),
                                    np.arange(len(audio_data)),
                                    audio_data
                                )
                            
                            # Stream in chunks of ~250ms of audio
                            chunk_size = int(SAMPLING_RATE * 0.25)
                            words_in_line = line_text.split()
                            total_words = len(words_in_line)
                            
                            # Map audio samples to words to synchronize 'word' HUD metadata events
                            samples_per_word = len(audio_data) / max(1, total_words)
                            
                            current_word_idx = 0
                            for offset in range(0, len(audio_data), chunk_size):
                                chunk = audio_data[offset:offset+chunk_size]
                                if len(chunk) == 0:
                                    continue
                                
                                # Scale chunk dynamically to int16 PCM
                                max_chunk = np.max(np.abs(chunk))
                                scaled_chunk = chunk / max_chunk * 0.7 if max_chunk > 0 else chunk
                                pcm_bytes = (scaled_chunk * 32767.0).astype(np.int16).tobytes()
                                
                                # Send binary audio buffer over WebSocket
                                await websocket.send_bytes(pcm_bytes)
                                
                                # Broadcast 'word' events as the audio reaches those word offsets
                                current_sample_pos = offset + len(chunk)
                                word_progress = int(current_sample_pos / samples_per_word)
                                while current_word_idx < min(word_progress, total_words):
                                    await websocket.send_json({
                                        "type": "word",
                                        "word": words_in_line[current_word_idx],
                                        "index": current_word_idx,
                                        "total": total_words
                                    })
                                    current_word_idx += 1
                                
                                # Realtime throttle matching output playback
                                await asyncio.sleep(0.24)
                            
                            # Ensure all final words are sent in the metadata stream
                            while current_word_idx < total_words:
                                await websocket.send_json({
                                    "type": "word",
                                    "word": words_in_line[current_word_idx],
                                    "index": current_word_idx,
                                    "total": total_words
                                })
                                current_word_idx += 1
                                
                            # Clean up
                            try:
                                os.remove(temp_wav)
                            except:
                                pass
                            
                            # Natural speech pause between sentences
                            await asyncio.sleep(0.3)
                        else:
                            raise Exception("Generated SAPI5 stream WAV is missing or empty")
                            
                    await websocket.send_json({"type": "done"})
                    
                except Exception as ex:
                    print(f"[WebSocket-Demo] SAPI5 streaming exception: {ex}. Falling back to visual wave generator.")
                    traceback.print_exc()
                    
                    # Resilient visual fallback to sine wave beeps if SAPI5 fails
                    words = text.split()
                    for word_idx, word in enumerate(words):
                        word_duration = 0.35
                        num_samples = int(SAMPLING_RATE * word_duration)
                        t = np.linspace(0, word_duration, num_samples, endpoint=False)
                        
                        base_freq = 130 if speaker_id == 0 else 230
                        wave = np.sin(2 * np.pi * base_freq * t)
                        wave += 0.4 * np.sin(2 * np.pi * (base_freq * 2) * t)
                        wave += 0.2 * np.sin(2 * np.pi * (base_freq * 3) * t)
                        
                        window = np.sin(np.linspace(0, np.pi, num_samples)) ** 2
                        chunk = wave * window
                        
                        chunk = chunk / np.max(np.abs(chunk)) * 0.7 if np.max(np.abs(chunk)) > 0 else chunk
                        pcm_bytes = (chunk * 32767.0).astype(np.int16).tobytes()
                        
                        await websocket.send_bytes(pcm_bytes)
                        await websocket.send_json({
                            "type": "word",
                            "word": word,
                            "index": word_idx,
                            "total": len(words)
                        })
                        await asyncio.sleep(0.28)
                        
                    await websocket.send_json({"type": "done"})
                
    except WebSocketDisconnect:
        print("[WebSocket] Streaming client disconnected.")
    except Exception as e:
        print(f"[WebSocket] Error: {e}")
        traceback.print_exc()

# Mount static assets files (temp outputs & cloned speakers)
app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
async def root_fallback():
    if os.path.exists("frontend/dist/index.html"):
        return FileResponse("frontend/dist/index.html")
    return HTMLResponse("""
    <html>
        <head>
            <title>VibeVoice API Server</title>
            <link rel="preconnect" href="https://fonts.googleapis.com">
            <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
            <link href="https://fonts.googleapis.com/css2?family=Outfit:wght@400;600&display=swap" rel="stylesheet">
            <style>
                body {
                    font-family: 'Outfit', sans-serif;
                    background: #07090e;
                    color: #f0f3fa;
                    display: flex;
                    flex-direction: column;
                    align-items: center;
                    justify-content: center;
                    height: 100vh;
                    margin: 0;
                }
                .container {
                    text-align: center;
                    padding: 40px;
                    background: rgba(255, 255, 255, 0.03);
                    border: 1px solid rgba(255, 255, 255, 0.08);
                    border-radius: 20px;
                    backdrop-filter: blur(10px);
                    max-width: 500px;
                }
                code {
                    background: rgba(0, 242, 254, 0.1);
                    color: #00f2fe;
                    padding: 8px 12px;
                    border-radius: 6px;
                    font-family: monospace;
                    display: inline-block;
                    margin: 15px 0;
                }
            </style>
        </head>
        <body>
            <div class="container">
                <h1 style="margin: 0 0 10px 0;">🎙️ VibeVoice API Server Online</h1>
                <p style="color: #8e9bb5; line-height: 1.6; font-size: 14px;">The backend API server is fully running on port 8000! To access the high-end React + Vite + Tailwind frontend, start the dev server:</p>
                <code>cd frontend; npm run dev</code>
                <p style="margin: 15px 0 0 0; font-size: 12px; color: #8e9bb5;">Or build for production serving: <code>cd frontend; npm run build</code></p>
            </div>
        </body>
    </html>
    """)

# Serve the static built files from dist if compiled
if os.path.exists("frontend/dist"):
    app.mount("/", StaticFiles(directory="frontend/dist", html=True), name="dist")

if __name__ == "__main__":
    import uvicorn
    import os
    # Start web app dynamically (binds to 0.0.0.0 and port 7860 for Hugging Face Space Docker support)
    port = int(os.environ.get("PORT", 8000))
    host = "0.0.0.0"
    print(f"[VibeVoice] Starting server on http://{host}:{port}")
    uvicorn.run("app:app", host=host, port=port, reload=False)