server.py

import sys
import types as _types

# ── _lzma shim for Python builds missing the C extension ──────────────────
if "_lzma" not in sys.modules:
    try:
        import _lzma as _real_lzma  # noqa: F401
    except ImportError:
        _fake = _types.ModuleType("_lzma")
        _fake.FILTER_ARM = 0x3010207
        _fake.FILTER_ARMTHUMB = 0x3010208
        _fake.FILTER_DELTA = 0x30301
        _fake.FILTER_IA64 = 0x3010206
        _fake.FILTER_LZMA1 = 0x201
        _fake.FILTER_LZMA2 = 0x21001
        _fake.FILTER_POWERPC = 0x3010205
        _fake.FILTER_SPARC = 0x301020A
        _fake.FILTER_X86 = 0x3010101
        _fake.FORMAT_ALONE = 1
        _fake.FORMAT_AUTO = 0
        _fake.FORMAT_RAW = 3
        _fake.FORMAT_XZ = 2
        _fake.CHECK_NONE = 0
        _fake.CHECK_CRC32 = 1
        _fake.CHECK_CRC64 = 4
        _fake.CHECK_SHA256 = 10
        _fake.CHECK_ID_MAX = 255
        _fake.CHECK_UNKNOWN = 255
        _fake.PRESET_DEFAULT = 60
        _fake.PRESET_EXTREME = 90 | (1 << 31)

        class _LZMACompressor:
            def __init__(self, *a, **kw): pass
            def compress(self, data): return data
            def flush(self): return b""

        class _LZMADecompressor:
            def __init__(self, *a, **kw): pass
            def decompress(self, data, max_length=-1): return data

        class _LZMAError(Exception): pass

        _fake.LZMACompressor = _LZMACompressor
        _fake.LZMADecompressor = _LZMADecompressor
        _fake.LZMAError = _LZMAError
        _fake._encode_filter_properties = lambda *a: b""
        _fake._decode_filter_properties = lambda *a: ()
        sys.modules["_lzma"] = _fake
        print("[VoxCPM2] Warning: _lzma C extension missing, using pure-Python fallback (lzma compression disabled)")

import io
import os
import uuid
import tempfile
import traceback
from pathlib import Path

import numpy as np
import soundfile as sf
import uvicorn
from fastapi import FastAPI, File, Form, UploadFile, Request
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse

# ── Config ───────────────────────────────────────────────────────────────
MODEL_PATH = os.environ.get(
    "VOXCPM_MODEL_PATH",
    "/Users/pangu/.omlx/models/VoxCPM2",
)
PORT = int(os.environ.get("VOXCPM_PORT", "8808"))

# ── App ──────────────────────────────────────────────────────────────────
app = FastAPI(title="VoxCPM2 Web UI")


@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    tb = traceback.format_exc()
    print(f"[VoxCPM2] ERROR {request.method} {request.url.path}:\n{tb}")
    return JSONResponse(
        status_code=500,
        content={"error": str(exc), "detail": tb},
    )

INDEX_HTML = Path(__file__).parent / "index.html"
APP_JS = Path(__file__).parent / "app.js"

OUTPUT_DIR = Path(tempfile.mkdtemp(prefix="voxcpm2_out_"))

# ── Model (loaded lazily on first request) ──────────────────────────────
_model = None


def get_model():
    global _model
    if _model is None:
        from voxcpm import VoxCPM

        print(f"[VoxCPM2] Loading model from {MODEL_PATH} ...")
        _model = VoxCPM.from_pretrained(MODEL_PATH, load_denoiser=False)
        print(f"[VoxCPM2] Model loaded on device: {_model.tts_model.device}")
    return _model


# ── Helpers ──────────────────────────────────────────────────────────────

def wav_bytes(wav_array: np.ndarray, sample_rate: int) -> bytes:
    buf = io.BytesIO()
    sf.write(buf, wav_array, sample_rate, format="WAV")
    buf.seek(0)
    return buf.read()


# ── Routes ───────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index():
    content = INDEX_HTML.read_text(encoding="utf-8")
    return HTMLResponse(
        content=content,
        headers={
            "Cache-Control": "no-cache, no-store, must-revalidate",
            "Pragma": "no-cache",
            "Expires": "0",
        },
    )


@app.get("/app.js")
async def app_js():
    content = APP_JS.read_text(encoding="utf-8")
    return HTMLResponse(
        content=content,
        media_type="application/javascript",
        headers={
            "Cache-Control": "no-cache, no-store, must-revalidate",
            "Pragma": "no-cache",
            "Expires": "0",
        },
    )


@app.get("/api/health")
async def health():
    return {
        "status": "ok",
        "model_loaded": _model is not None,
        "model_path": MODEL_PATH,
    }


@app.post("/api/tts")
async def tts(
    text: str = Form(...),
    cfg_value: float = Form(2.0),
    inference_timesteps: int = Form(10),
):
    """Basic text-to-speech."""
    model = get_model()
    wav = model.generate(
        text=text,
        cfg_value=cfg_value,
        inference_timesteps=inference_timesteps,
    )
    audio = wav_bytes(wav, model.tts_model.sample_rate)
    return StreamingResponse(io.BytesIO(audio), media_type="audio/wav")


@app.post("/api/design")
async def design(
    text: str = Form(...),
    voice_description: str = Form(...),
    cfg_value: float = Form(2.0),
    inference_timesteps: int = Form(10),
):
    """Voice design: synthesize text with a described voice style."""
    model = get_model()
    prompt = f"({voice_description}){text}"
    wav = model.generate(
        text=prompt,
        cfg_value=cfg_value,
        inference_timesteps=inference_timesteps,
    )
    audio = wav_bytes(wav, model.tts_model.sample_rate)
    return StreamingResponse(io.BytesIO(audio), media_type="audio/wav")


@app.post("/api/clone")
async def clone(
    text: str = Form(...),
    reference_audio: UploadFile = File(...),
    cfg_value: float = Form(2.0),
    inference_timesteps: int = Form(10),
    style_hint: str = Form(""),
):
    """Voice cloning from reference audio."""
    model = get_model()

    ref_data = await reference_audio.read()
    ref_path = OUTPUT_DIR / f"ref_{uuid.uuid4().hex[:8]}.wav"
    ref_path.write_bytes(ref_data)

    prompt_text = f"({style_hint}){text}" if style_hint else text
    wav = model.generate(
        text=prompt_text,
        reference_wav_path=str(ref_path),
        cfg_value=cfg_value,
        inference_timesteps=inference_timesteps,
    )
    audio = wav_bytes(wav, model.tts_model.sample_rate)
    ref_path.unlink(missing_ok=True)

    return StreamingResponse(io.BytesIO(audio), media_type="audio/wav")


@app.post("/api/clone-ultimate")
async def clone_ultimate(
    text: str = Form(...),
    reference_audio: UploadFile = File(...),
    prompt_text: str = Form(...),
    cfg_value: float = Form(2.0),
    inference_timesteps: int = Form(10),
):
    """Ultimate (Hi-Fi) cloning: reference audio + transcript."""
    model = get_model()

    ref_data = await reference_audio.read()
    ref_path = OUTPUT_DIR / f"ref_{uuid.uuid4().hex[:8]}.wav"
    ref_path.write_bytes(ref_data)

    wav = model.generate(
        text=text,
        prompt_wav_path=str(ref_path),
        prompt_text=prompt_text,
        reference_wav_path=str(ref_path),
        cfg_value=cfg_value,
        inference_timesteps=inference_timesteps,
    )
    audio = wav_bytes(wav, model.tts_model.sample_rate)
    ref_path.unlink(missing_ok=True)

    return StreamingResponse(io.BytesIO(audio), media_type="audio/wav")


# ── Main ─────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    print(f"[VoxCPM2] Starting server on http://localhost:{PORT}")
    print(f"[VoxCPM2] Model path: {MODEL_PATH}")
    print(f"[VoxCPM2] Model will be loaded on first request...")
    uvicorn.run(app, host="0.0.0.0", port=PORT)