esh-pfi-infrastructure/stacks/fish-cpp/server.py

"""fish-cpp — HTTP shim wrapping the s2.cpp CLI in Fish's /v1/tts API contract.

s2.cpp is a pure C++/GGML inference engine for Fish s2-pro GGUFs. It's
CLI-only (alpha), so this shim turns each HTTP POST into a subprocess
invocation that:

  1. Writes the optional reference WAV to /tmp/<id>.wav
  2. Runs:  s2 -m <model.gguf> -t <tokenizer.json>
              [-pa /tmp/ref.wav -pt "<transcript>"]
              -text "..." -o /tmp/<id>.out.wav
  3. Streams the resulting WAV back to the client.

The contract matches Fish's wrapper close enough that bench harnesses
written for fish-s2 work against fish-cpp without changes — same
POST /v1/tts, same {text, references[{audio,text}], format} request body.

What's NOT supported (vs Fish HF):
  * `streaming: true` — s2.cpp writes a complete WAV before returning.
    The shim accepts the field but ignores it. TTFB ≈ total wall time.
    If/when s2.cpp grows incremental output, wire it through here.
  * Paralinguistic tags — s2.cpp uses the same model weights, so tags
    SHOULD work, but quality may differ from the HF wrapper.
"""

from __future__ import annotations

import base64
import os
import subprocess
import tempfile
import uuid
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel

S2_BINARY = os.environ.get("S2_BINARY", "/usr/local/bin/s2")
WEIGHTS_DIR = Path(os.environ.get("WEIGHTS_DIR", "/weights"))
MODEL_FILE = os.environ.get("FISH_CPP_MODEL", "s2-pro-q6_k.gguf")
TOKENIZER_FILE = os.environ.get("FISH_CPP_TOKENIZER", "tokenizer.json")
DEVICE_INDEX = os.environ.get("FISH_CPP_DEVICE", "0")  # CUDA device id

MODEL_PATH = WEIGHTS_DIR / MODEL_FILE
TOKENIZER_PATH = WEIGHTS_DIR / TOKENIZER_FILE

app = FastAPI(title="fish-cpp")


class ReferenceAudio(BaseModel):
    audio: str  # base64-encoded WAV bytes (matches Fish HF schema)
    text: str   # transcript of the reference clip


class TTSRequest(BaseModel):
    text: str
    references: list[ReferenceAudio] = []
    format: str = "wav"
    streaming: bool = False  # accepted for contract parity; ignored


@app.get("/v1/health")
def health() -> dict:
    if not MODEL_PATH.exists():
        raise HTTPException(503, f"model missing: {MODEL_PATH}")
    if not TOKENIZER_PATH.exists():
        raise HTTPException(503, f"tokenizer missing: {TOKENIZER_PATH}")
    return {"status": "ok", "model": MODEL_FILE}


@app.post("/v1/tts")
def tts(req: TTSRequest) -> FileResponse:
    if req.format not in ("wav",):
        raise HTTPException(400, f"only wav format supported, got {req.format}")

    job = uuid.uuid4().hex
    out_path = Path(tempfile.gettempdir()) / f"fish-cpp-{job}.wav"
    ref_path: Optional[Path] = None

    cmd: list[str] = [
        S2_BINARY,
        "-m", str(MODEL_PATH),
        "-t", str(TOKENIZER_PATH),
        "-text", req.text,
        "-o", str(out_path),
        # -v selects voice mode; 0 = default (no preset). Preset voices
        # would need s2.cpp to ship a voice library; clone via -pa/-pt.
        "-v", "0",
    ]

    if req.references:
        # s2.cpp accepts a single reference at a time via -pa/-pt. If the
        # client sent multiple, use the first; the upstream API doesn't
        # appear to support multi-reference cloning.
        ref = req.references[0]
        ref_path = Path(tempfile.gettempdir()) / f"fish-cpp-{job}-ref.wav"
        ref_path.write_bytes(base64.b64decode(ref.audio))
        cmd.extend(["-pa", str(ref_path), "-pt", ref.text])

    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = DEVICE_INDEX

    try:
        result = subprocess.run(
            cmd, env=env, capture_output=True, text=True, timeout=180,
        )
    except subprocess.TimeoutExpired:
        raise HTTPException(504, "s2 binary timed out (180s)")
    finally:
        if ref_path and ref_path.exists():
            ref_path.unlink()

    if result.returncode != 0:
        raise HTTPException(
            500,
            f"s2 binary failed (rc={result.returncode}): {result.stderr[-500:]}",
        )

    if not out_path.exists() or out_path.stat().st_size == 0:
        raise HTTPException(500, "s2 binary produced no output")

    # FileResponse streams the file + cleans up after sending. Caller
    # gets WAV bytes immediately; we reap the temp file via background
    # task once the response finishes.
    from fastapi import BackgroundTasks
    bg = BackgroundTasks()
    bg.add_task(out_path.unlink, missing_ok=True)
    return FileResponse(
        out_path, media_type="audio/wav", filename=f"{job}.wav", background=bg,
    )