app.py

"""
OmniVoice TTS — fully local FastAPI server for HuggingFace Spaces (T4 / dedicated GPU).

Loads `k2-fsa/OmniVoice` once at startup and serves four endpoints:
    GET  /            → service descriptor (JSON)
    GET  /health      → liveness probe
    POST /tts         → one-shot WAV (audio/wav, 24 kHz mono int16)
    POST /tts/stream  → chunked WAV (header + raw PCM tail) — same payload
    WS   /ws/tts      → real-time PCM frames + JSON status messages

Voice-design (no reference audio) and voice-clone (with reference audio) modes
are both supported. Attributes (gender / age / pitch / accent / dialect / style)
are forwarded to the OmniVoice model via the `instruct` parameter, mirroring
the upstream Gradio demo's behaviour.

This file is self-contained: it does NOT call out to the public Gradio Space.
The model runs in-process on the local GPU.
"""

from __future__ import annotations

import asyncio
import io
import json
import logging
import os
import struct
import tempfile
import threading
from collections import OrderedDict
from typing import Any, AsyncIterator, Optional

import numpy as np
import torch
from fastapi import FastAPI, HTTPException, UploadFile, File, Form, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, Response, StreamingResponse
from pydantic import BaseModel, Field

from omnivoice import OmniVoice, OmniVoiceGenerationConfig

logging.basicConfig(
    level=os.environ.get("LOG_LEVEL", "INFO"),
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
logger = logging.getLogger("omnivoice-tts")

# ─── Model loading ────────────────────────────────────────────────────────────
CHECKPOINT = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice")
DEVICE = os.environ.get("OMNIVOICE_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
# Whisper ASR is only needed to auto-transcribe a reference clip when ref_text
# is omitted. Our workflow always supplies ref_text, so default to NOT loading it
# (~1 GB less to download/load → faster cold starts). Set OMNIVOICE_LOAD_ASR=1 to
# re-enable auto-transcription.
LOAD_ASR = os.environ.get("OMNIVOICE_LOAD_ASR", "0") == "1"

logger.info("Loading OmniVoice from %s onto %s (dtype=%s, asr=%s)",
            CHECKPOINT, DEVICE, DTYPE, LOAD_ASR)
model = OmniVoice.from_pretrained(
    CHECKPOINT,
    device_map=DEVICE,
    dtype=DTYPE,
    load_asr=LOAD_ASR,
)
SAMPLE_RATE: int = int(model.sampling_rate)
logger.info("Model loaded — sample_rate=%d Hz", SAMPLE_RATE)


# ─── Voice registry (clone once, reuse many times) ────────────────────────────
# A client registers a reference voice once via POST /voices under a chosen name;
# we compute the OmniVoice voice-clone prompt a single time and cache it here.
# Subsequent /tts (and /tts/stream, /ws/tts) requests pass `voice_id` to reuse it,
# skipping all reference re-encoding (and the Whisper ASR pass when ref_text was
# supplied). Prompts are persisted to VOICES_DIR so they survive restarts.
MAX_VOICES = int(os.environ.get("OMNIVOICE_MAX_VOICES", "50"))
VOICES_DIR = os.environ.get("OMNIVOICE_VOICES_DIR", "/data/voices")
_VOICE_CACHE: "OrderedDict[str, Any]" = OrderedDict()
_VOICE_LOCK = threading.Lock()

os.makedirs(VOICES_DIR, exist_ok=True)


def _voice_path(name: str) -> str:
    """Return the .pt path for a given voice name (safe filename)."""
    safe = name.replace("/", "_").replace("\\", "_")
    return os.path.join(VOICES_DIR, f"{safe}.pt")


def _persist_voice(name: str, prompt: Any) -> None:
    """Save a voice-clone prompt to disk using torch.save."""
    try:
        torch.save(prompt, _voice_path(name))
    except Exception:
        logger.warning("Could not persist voice %r to disk", name, exc_info=True)


def _delete_persisted_voice(name: str) -> None:
    """Remove the on-disk .pt file for a voice (best-effort)."""
    try:
        os.unlink(_voice_path(name))
    except Exception:
        pass


def _load_persisted_voices() -> None:
    """At startup, reload all saved voice prompts from VOICES_DIR into the cache."""
    try:
        files = [f for f in os.listdir(VOICES_DIR) if f.endswith(".pt")]
    except Exception:
        return
    loaded = 0
    for fname in files:
        name = fname[:-3]  # strip .pt
        try:
            prompt = torch.load(os.path.join(VOICES_DIR, fname), map_location=DEVICE, weights_only=False)
            _VOICE_CACHE[name] = prompt
            loaded += 1
        except Exception:
            logger.warning("Failed to reload voice %r from disk", name, exc_info=True)
    if loaded:
        logger.info("Reloaded %d persisted voice(s) from %s", loaded, VOICES_DIR)


_load_persisted_voices()


def _register_voice(name: str, ref_audio_path: str, ref_text: Optional[str]) -> int:
    """Build and cache a voice-clone prompt under `name` (LRU-evicting). Returns cache size."""
    prompt = model.create_voice_clone_prompt(ref_audio=ref_audio_path, ref_text=ref_text)
    _persist_voice(name, prompt)
    with _VOICE_LOCK:
        _VOICE_CACHE[name] = prompt          # overwrite if name already exists
        _VOICE_CACHE.move_to_end(name)       # mark most-recently-used
        while len(_VOICE_CACHE) > MAX_VOICES:
            evicted, _ = _VOICE_CACHE.popitem(last=False)  # drop least-recently-used
            _delete_persisted_voice(evicted)
            logger.info("Evicted LRU voice %r (cache full at %d)", evicted, MAX_VOICES)
        return len(_VOICE_CACHE)


def _get_voice(name: str) -> Optional[Any]:
    """Return the cached prompt for `name`, marking it most-recently-used, or None."""
    with _VOICE_LOCK:
        prompt = _VOICE_CACHE.get(name)
        if prompt is not None:
            _VOICE_CACHE.move_to_end(name)
        return prompt


# ─── Voice-design attribute vocabulary (matches upstream Gradio demo) ─────────
# Each entry is "English / 中文". For accents we pass the English half; for
# dialects we pass the Chinese half — mirrors omnivoice/cli/demo.py logic.
GENDERS = {"Auto", "Male / 男", "Female / 女"}
AGES = {
    "Auto",
    "Child / 儿童",
    "Teenager / 少年",
    "Young Adult / 青年",
    "Middle-aged / 中年",
    "Elderly / 老年",
}
PITCHES = {
    "Auto",
    "Very Low Pitch / 极低音调",
    "Low Pitch / 低音调",
    "Moderate Pitch / 中音调",
    "High Pitch / 高音调",
    "Very High Pitch / 极高音调",
}
STYLES = {"Auto", "Whisper / 耳语"}
ACCENTS = {
    "Auto",
    "American Accent / 美式口音",
    "Australian Accent / 澳大利亚口音",
    "British Accent / 英国口音",
    "Chinese Accent / 中国口音",
    "Canadian Accent / 加拿大口音",
    "Indian Accent / 印度口音",
    "Korean Accent / 韩国口音",
    "Portuguese Accent / 葡萄牙口音",
    "Russian Accent / 俄罗斯口音",
    "Japanese Accent / 日本口音",
}
DIALECTS = {
    "Auto",
    "Henan Dialect / 河南话",
    "Shaanxi Dialect / 陕西话",
    "Sichuan Dialect / 四川话",
    "Guizhou Dialect / 贵州话",
    "Yunnan Dialect / 云南话",
    "Guilin Dialect / 桂林话",
    "Jinan Dialect / 济南话",
    "Shijiazhuang Dialect / 石家庄话",
    "Gansu Dialect / 甘肃话",
    "Ningxia Dialect / 宁夏话",
    "Qingdao Dialect / 青岛话",
    "Northeast Dialect / 东北话",
}

LANGUAGES_ADVERTISED = ["Auto", "English", "Hindi", "Panjabi", "Urdu", "Western Panjabi"]


def _attr_part(value: Optional[str], is_dialect: bool = False) -> Optional[str]:
    """Pick the right half of a bilingual attribute label, or None for Auto."""
    if not value or value == "Auto":
        return None
    if " / " in value:
        en, zh = value.split(" / ", 1)
        return zh.strip() if is_dialect else en.strip()
    return value.strip()


def build_instruct(
    *,
    gender: Optional[str] = None,
    age: Optional[str] = None,
    pitch: Optional[str] = None,
    style: Optional[str] = None,
    accent: Optional[str] = None,
    dialect: Optional[str] = None,
    extra: Optional[str] = None,
) -> Optional[str]:
    """Compose the OmniVoice `instruct` string from voice-design attributes."""
    parts = []
    for v, is_dialect in (
        (gender, False), (age, False), (pitch, False), (style, False),
        (accent, False), (dialect, True),
    ):
        p = _attr_part(v, is_dialect=is_dialect)
        if p:
            parts.append(p)
    if extra and extra.strip():
        parts.append(extra.strip())
    return ", ".join(parts) if parts else None


# ─── Synthesis core ───────────────────────────────────────────────────────────
class TTSRequest(BaseModel):
    text: str = Field(..., min_length=1, description="Text to synthesize.")
    language: str = Field("Auto", description="Language name (English label, e.g. 'Urdu', 'Panjabi') or 'Auto'.")
    # Voice-design attributes — accept either the canonical bilingual label
    # ("Female / 女") or the friendly English half ("Female"). Auto = no constraint.
    gender: str = "Auto"
    age: str = "Auto"
    pitch: str = "Auto"
    style: str = "Auto"
    accent: str = "Auto"
    dialect: str = "Auto"
    instruct: Optional[str] = None  # raw override; takes precedence if set
    voice_id: Optional[str] = None  # reuse a voice registered via POST /voices
    # Generation config
    speed: float = 1.0
    duration: Optional[float] = None
    nfe_steps: int = Field(32, ge=4, le=64)
    guidance: float = Field(2.0, ge=0.0, le=4.0)
    # Sampling temperatures (model defaults: class=0.0 greedy, position=5.0).
    # class_temperature 0.0 = deterministic/greedy → can sound flat/robotic;
    # raise toward ~0.3–0.7 for more natural variation.
    class_temperature: float = Field(0.3, ge=0.0, le=2.0)
    position_temperature: float = Field(5.0, ge=0.0, le=10.0)
    denoise: bool = True
    preprocess_prompt: bool = True
    postprocess_output: bool = True


def _normalize_label(value: str, valid: set) -> str:
    """Accept friendly aliases ('female', 'young', 'low') in addition to bilingual labels."""
    if not value or value == "Auto":
        return "Auto"
    if value in valid:
        return value
    lc = value.strip().lower()
    aliases = {
        "male": "Male / 男", "m": "Male / 男",
        "female": "Female / 女", "f": "Female / 女",
        "child": "Child / 儿童",
        "teen": "Teenager / 少年", "teenager": "Teenager / 少年",
        "young": "Young Adult / 青年", "young_adult": "Young Adult / 青年",
        "middle": "Middle-aged / 中年", "middle_aged": "Middle-aged / 中年",
        "elderly": "Elderly / 老年", "old": "Elderly / 老年",
        "very_low": "Very Low Pitch / 极低音调", "very low": "Very Low Pitch / 极低音调",
        "low": "Low Pitch / 低音调",
        "moderate": "Moderate Pitch / 中音调", "medium": "Moderate Pitch / 中音调",
        "high": "High Pitch / 高音调",
        "very_high": "Very High Pitch / 极高音调", "very high": "Very High Pitch / 极高音调",
        "whisper": "Whisper / 耳语",
    }
    return aliases.get(lc, value if value in valid else "Auto")


def _synthesize(
    req: TTSRequest,
    *,
    ref_audio_path: Optional[str] = None,
    ref_text: Optional[str] = None,
    voice_clone_prompt: Optional[Any] = None,
) -> np.ndarray:
    """Run OmniVoice inference and return int16 PCM samples (1-D ndarray)."""
    # Normalize friendly aliases to canonical bilingual labels
    gender = _normalize_label(req.gender, GENDERS)
    age = _normalize_label(req.age, AGES)
    pitch = _normalize_label(req.pitch, PITCHES)
    style = _normalize_label(req.style, STYLES)
    accent = req.accent if req.accent in ACCENTS else "Auto"
    dialect = req.dialect if req.dialect in DIALECTS else "Auto"

    instruct = req.instruct or build_instruct(
        gender=gender, age=age, pitch=pitch, style=style,
        accent=accent, dialect=dialect,
    )
    lang = req.language if req.language and req.language != "Auto" else None

    gen_config = OmniVoiceGenerationConfig(
        num_step=req.nfe_steps,
        guidance_scale=req.guidance,
        class_temperature=req.class_temperature,
        position_temperature=req.position_temperature,
        denoise=req.denoise,
        preprocess_prompt=req.preprocess_prompt,
        postprocess_output=req.postprocess_output,
    )

    kw: dict[str, Any] = dict(
        text=req.text.strip(),
        language=lang,
        generation_config=gen_config,
    )
    if req.speed and req.speed != 1.0:
        kw["speed"] = float(req.speed)
    if req.duration and req.duration > 0:
        kw["duration"] = float(req.duration)
    if instruct:
        kw["instruct"] = instruct
    if voice_clone_prompt is not None:
        kw["voice_clone_prompt"] = voice_clone_prompt
    elif ref_audio_path:
        kw["voice_clone_prompt"] = model.create_voice_clone_prompt(
            ref_audio=ref_audio_path,
            ref_text=ref_text,
        )

    cloning = voice_clone_prompt is not None or bool(ref_audio_path)
    logger.info("Synthesize: lang=%s len=%d instruct=%r clone=%s voice_id=%s",
                lang, len(req.text), instruct, cloning, req.voice_id)

    with torch.inference_mode():
        audio = model.generate(**kw)

    # audio[0] is float32 in [-1, 1]; convert to int16 PCM.
    pcm = np.clip(audio[0], -1.0, 1.0)
    return (pcm * 32767).astype(np.int16)


# ─── WAV helpers ──────────────────────────────────────────────────────────────
def _wav_header(num_samples: int, sample_rate: int = SAMPLE_RATE, channels: int = 1) -> bytes:
    """Build a 44-byte WAV header for int16 mono PCM."""
    sample_width = 2
    data_size = num_samples * channels * sample_width
    return (
        b"RIFF"
        + struct.pack("<I", 36 + data_size)
        + b"WAVE"
        + b"fmt "
        + struct.pack("<I", 16)
        + struct.pack("<H", 1)
        + struct.pack("<H", channels)
        + struct.pack("<I", sample_rate)
        + struct.pack("<I", sample_rate * channels * sample_width)
        + struct.pack("<H", channels * sample_width)
        + struct.pack("<H", sample_width * 8)
        + b"data"
        + struct.pack("<I", data_size)
    )


def _wrap_wav(pcm_int16: np.ndarray) -> bytes:
    return _wav_header(len(pcm_int16)) + pcm_int16.tobytes()


# ─── FastAPI app ──────────────────────────────────────────────────────────────
app = FastAPI(
    title="OmniVoice TTS",
    description="In-process OmniVoice TTS server (no upstream proxy).",
    version="1.0.0",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

_DESCRIPTOR = {
    "service": "OmniVoice TTS",
    "model": CHECKPOINT,
    "device": DEVICE,
    "sample_rate": SAMPLE_RATE,
    "endpoints": {
        "GET /health": "Health probe",
        "POST /tts": "One-shot WAV (audio/wav)",
        "POST /tts/stream": "Chunked WAV (header + raw PCM tail)",
        "WS /ws/tts": "Real-time binary PCM chunks + JSON status",
        "POST /tts/clone": "Ad-hoc voice cloning (multipart form: ref_audio + text)",
        "POST /voices": "Register a voice once (multipart: name + ref_audio + ref_text?)",
        "GET /voices": "List cached voice ids",
        "DELETE /voices/{name}": "Free a registered voice",
    },
    "languages": LANGUAGES_ADVERTISED,
    "genders": sorted(GENDERS),
    "ages": sorted(AGES),
    "pitches": sorted(PITCHES),
    "styles": sorted(STYLES),
    "accents": sorted(ACCENTS),
    "dialects": sorted(DIALECTS),
}


@app.get("/")
async def root():
    return JSONResponse(_DESCRIPTOR)


@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": True, "device": DEVICE}


def _resolve_voice(voice_id: Optional[str]) -> Optional[Any]:
    """Look up a registered voice prompt, or raise 404. Returns None when no voice_id given."""
    if not voice_id:
        return None
    prompt = _get_voice(voice_id)
    if prompt is None:
        raise HTTPException(
            status_code=404,
            detail=f"voice_id {voice_id!r} not registered; POST /voices first",
        )
    return prompt


@app.post("/tts")
async def tts_one_shot(req: TTSRequest):
    """Return the full WAV in one response."""
    voice_prompt = _resolve_voice(req.voice_id)
    try:
        # Run inference in a thread so the event loop stays responsive.
        pcm = await asyncio.to_thread(_synthesize, req, voice_clone_prompt=voice_prompt)
    except Exception as e:
        logger.exception("synth failed")
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
    if pcm.size == 0:
        raise HTTPException(status_code=500, detail="empty audio")
    return Response(content=_wrap_wav(pcm), media_type="audio/wav")


@app.post("/tts/stream")
async def tts_stream(req: TTSRequest):
    """
    Chunked HTTP: emit a WAV header up-front (with a placeholder data size)
    followed by the PCM bytes. Most decoders that read sequentially handle
    this fine. Lower TTFB than /tts because the header flushes immediately.
    """
    voice_prompt = _resolve_voice(req.voice_id)
    try:
        pcm = await asyncio.to_thread(_synthesize, req, voice_clone_prompt=voice_prompt)
    except Exception as e:
        logger.exception("stream synth failed")
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")

    async def gen() -> AsyncIterator[bytes]:
        yield _wav_header(len(pcm))
        # Yield ~0.2 s frames so clients see audio progressively.
        frame = SAMPLE_RATE // 5
        raw = pcm.tobytes()
        bytes_per_frame = frame * 2  # int16
        for i in range(0, len(raw), bytes_per_frame):
            yield raw[i: i + bytes_per_frame]

    return StreamingResponse(gen(), media_type="audio/wav")


@app.post("/tts/clone")
async def tts_clone(
    text: str = Form(...),
    ref_audio: UploadFile = File(...),
    language: str = Form("Auto"),
    ref_text: Optional[str] = Form(None),
    speed: float = Form(1.0),
    nfe_steps: int = Form(32),
    guidance: float = Form(2.0),
    class_temperature: float = Form(0.3),
    position_temperature: float = Form(5.0),
    denoise: bool = Form(True),
):
    """Voice cloning with an uploaded reference audio (ad-hoc, one-off)."""
    suffix = os.path.splitext(ref_audio.filename or "ref.wav")[1] or ".wav"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(await ref_audio.read())
        tmp_path = tmp.name
    try:
        req = TTSRequest(
            text=text, language=language,
            speed=speed, nfe_steps=nfe_steps,
            guidance=guidance, denoise=denoise,
            class_temperature=class_temperature,
            position_temperature=position_temperature,
        )
        pcm = await asyncio.to_thread(
            _synthesize, req,
            ref_audio_path=tmp_path, ref_text=ref_text,
        )
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass
    return Response(content=_wrap_wav(pcm), media_type="audio/wav")


# ─── Voice registry endpoints (clone once, reuse via voice_id) ────────────────
@app.post("/voices")
async def register_voice(
    name: str = Form(..., min_length=1),
    ref_audio: UploadFile = File(...),
    ref_text: Optional[str] = Form(None),
):
    """
    Register a reference voice ONCE under `name`. The voice-clone prompt is
    computed here and cached in memory; reuse it on /tts, /tts/stream and
    /ws/tts by passing "voice_id": "<name>". Re-registering the same name
    overwrites it. Omitting ref_text triggers Whisper ASR (needs OMNIVOICE_LOAD_ASR=1).
    """
    suffix = os.path.splitext(ref_audio.filename or "ref.wav")[1] or ".wav"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(await ref_audio.read())
        tmp_path = tmp.name
    try:
        cached = await asyncio.to_thread(_register_voice, name, tmp_path, ref_text)
    except Exception as e:
        logger.exception("voice registration failed")
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass
    return {"voice_id": name, "voices_cached": cached, "max_voices": MAX_VOICES}


@app.get("/voices")
async def list_voices():
    """List currently cached voice ids (most-recently-used last)."""
    with _VOICE_LOCK:
        return {"voices": list(_VOICE_CACHE.keys()), "max_voices": MAX_VOICES}


@app.get("/debug/voices-dir")
async def debug_voices_dir():
    """Show what .pt files are actually on disk in VOICES_DIR."""
    result: dict = {"voices_dir": VOICES_DIR, "data_writable": os.access("/data", os.W_OK)}
    try:
        files = os.listdir(VOICES_DIR)
        result["files"] = {
            f: os.path.getsize(os.path.join(VOICES_DIR, f))
            for f in sorted(files)
        }
    except Exception as e:
        result["error"] = str(e)
    return result


@app.get("/debug/gpu")
async def debug_gpu():
    """Report live VRAM / RAM usage so we can right-size the GPU."""
    out: dict = {"device": DEVICE, "dtype": str(DTYPE), "load_asr": LOAD_ASR}
    try:
        if torch.cuda.is_available():
            free, total = torch.cuda.mem_get_info()
            out["gpu"] = {
                "name": torch.cuda.get_device_name(0),
                "vram_allocated_gb": round(torch.cuda.memory_allocated() / 1e9, 2),
                "vram_reserved_gb": round(torch.cuda.memory_reserved() / 1e9, 2),
                "vram_used_total_gb": round((total - free) / 1e9, 2),  # incl. other procs/context
                "vram_total_gb": round(total / 1e9, 2),
            }
    except Exception as e:
        out["gpu_error"] = str(e)
    try:
        import psutil
        p = psutil.Process(os.getpid())
        vm = psutil.virtual_memory()
        out["ram_rss_gb"] = round(p.memory_info().rss / 1e9, 2)
        out["ram_used_total_gb"] = round((vm.total - vm.available) / 1e9, 2)
        out["ram_total_gb"] = round(vm.total / 1e9, 2)
    except Exception as e:
        out["ram_error"] = str(e)
    return out


@app.delete("/voices/{name}")
async def delete_voice(name: str):
    """Free one registered voice. 404 if it isn't cached."""
    with _VOICE_LOCK:
        if _VOICE_CACHE.pop(name, None) is None:
            raise HTTPException(status_code=404, detail=f"voice_id {name!r} not registered")
        remaining = len(_VOICE_CACHE)
    _delete_persisted_voice(name)
    return {"deleted": name, "voices_cached": remaining}


@app.websocket("/ws/tts")
async def ws_tts(ws: WebSocket):
    """
    Bidirectional WebSocket:
      Client sends JSON request (TTSRequest schema).
      Server replies:
        {"type":"started","sample_rate":24000}
        <binary frame 1>
        <binary frame 2>
        ...
        {"type":"complete","duration_s":1.23}
      On error: {"type":"error","message":"..."}
    Connection closes after one synthesis (matches the Parler pattern).
    """
    await ws.accept()
    try:
        raw = await ws.receive_text()
        try:
            payload = json.loads(raw)
        except json.JSONDecodeError as e:
            await ws.send_json({"type": "error", "message": f"bad json: {e}"})
            return
        try:
            req = TTSRequest(**payload)
        except Exception as e:
            await ws.send_json({"type": "error", "message": f"bad request: {e}"})
            return

        voice_prompt = _get_voice(req.voice_id) if req.voice_id else None
        if req.voice_id and voice_prompt is None:
            await ws.send_json({
                "type": "error",
                "message": f"voice_id {req.voice_id!r} not registered; POST /voices first",
            })
            return

        await ws.send_json({"type": "started", "sample_rate": SAMPLE_RATE})

        try:
            pcm = await asyncio.to_thread(_synthesize, req, voice_clone_prompt=voice_prompt)
        except Exception as e:
            logger.exception("ws synth failed")
            await ws.send_json({"type": "error", "message": f"{type(e).__name__}: {e}"})
            return

        # Stream as ~0.5 s binary chunks so the client can begin playback early.
        frame_bytes = SAMPLE_RATE * 2 // 2  # 0.5 s of int16 mono
        raw_pcm = pcm.tobytes()
        for i in range(0, len(raw_pcm), frame_bytes):
            await ws.send_bytes(raw_pcm[i: i + frame_bytes])

        await ws.send_json({
            "type": "complete",
            "duration_s": round(len(pcm) / SAMPLE_RATE, 3),
        })
    except WebSocketDisconnect:
        logger.info("ws client disconnected")
    except Exception as e:
        logger.exception("ws unexpected error")
        try:
            await ws.send_json({"type": "error", "message": str(e)})
        except Exception:
            pass
    finally:
        try:
            await ws.close()
        except Exception:
            pass


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))