VoiceStudio/backend/engines/voxcpm2_subprocess/main.py

"""voxcpm2 sidecar: VoxCPM2 in the engine's own venv (one-click install).

Launched as ``<engine venv python> main.py`` by VoxCPM2SubprocessBackend. It
imports nothing from the app: the venv holds only ``voxcpm`` and what it
depends on (torch, torchaudio, numpy), so this file must stay importable with
the standard library plus those. The parent prepares the reference clip and
trims the output's silent tail, exactly as the in-process engine does; this
process only loads the model and synthesizes.

Wire protocol: length-prefixed JSON over stdio, identical to the other
sidecars (engines/pockettts/main.py). A ``ready`` frame comes first, then one
``audio`` (or ``error``) frame per ``synthesize``, with ``progress`` frames
while a cold load runs so the parent's watchdog stays armed.
"""
from __future__ import annotations

import base64
import json
import os
import re
import struct
import sys
import threading
import time
import traceback

# Mirrors services/subprocess_backend.py::MAX_FRAME_BYTES.
MAX_FRAME_BYTES = 64 * 1024 * 1024
#: VoxCPM2's studio output rate; the in-process engine assumes the same.
VOXCPM2_SAMPLE_RATE = 48_000
#: Emit a progress frame at least this often during a cold load (a multi-GB
#: first download) so the parent's recv watchdog doesn't kill a healthy sidecar.
_HEARTBEAT_S = 5.0
#: ref_audio must be a local file path, not a URL (local-first; no SSRF).
_URL_RE = re.compile(r"^[a-z][a-z0-9+.\-]*://", re.IGNORECASE)
#: A download failure worth retrying (the HF cache resumes, so a retry
#: continues rather than restarts). Anything else propagates at once.
_TRANSIENT_MARKERS = (
    "connection", "timed out", "timeout", "peer closed", "incomplete",
    "remoteprotocolerror", "temporarily unavailable",
)

_MODEL = None

# -- wire protocol -----------------------------------------------------------

#: Serializes _send across threads (the cold-load heartbeat + the main loop) so
#: concurrent length+body writes can't interleave and corrupt the framing.
_send_lock = threading.Lock()


def _send(stream, obj: dict) -> None:
    body = json.dumps(obj, separators=(",", ":")).encode("utf-8")
    with _send_lock:
        stream.write(struct.pack("!I", len(body)))
        stream.write(body)
        stream.flush()


def _recv(stream):
    header = stream.read(4)
    if len(header) < 4:
        return None  # EOF
    (n,) = struct.unpack("!I", header)
    if n > MAX_FRAME_BYTES:
        raise IOError(f"frame too large: {n}")
    body = bytearray()
    while len(body) < n:
        chunk = stream.read(n - len(body))
        if not chunk:
            raise IOError("short read")
        body.extend(chunk)
    return json.loads(bytes(body).decode("utf-8"))


def _measure_vram_mb() -> float:
    try:
        import torch  # noqa: PLC0415

        if torch.cuda.is_available():
            return float(torch.cuda.memory_allocated()) / (1024 * 1024)
    except Exception:  # noqa: BLE001 — a probe, never fatal
        pass
    return 0.0


# -- model loading (lazy, on the first synthesize) ---------------------------


def _with_retries(load):
    """Run ``load``, retrying a transient download failure with a short
    backoff, the way the app's own loader does for in-process engines."""
    try:
        attempts = max(1, int(os.environ.get("OMNIVOICE_MODEL_LOAD_RETRIES", "3")))
    except ValueError:
        attempts = 3
    for attempt in range(1, attempts + 1):
        try:
            return load()
        except Exception as exc:  # noqa: BLE001 — classified below
            text = f"{type(exc).__name__}: {exc}".lower()
            if attempt == attempts or not any(m in text for m in _TRANSIENT_MARKERS):
                raise
            time.sleep(2.0 * attempt)
    raise AssertionError("unreachable")


def _load_model(stdout):
    global _MODEL
    if _MODEL is not None:
        return _MODEL
    _send(stdout, {"op": "progress", "stage": "loading_model", "percent": 0})
    stop = threading.Event()

    def _heartbeat() -> None:
        pct = 1
        while not stop.wait(_HEARTBEAT_S):
            pct = min(pct + 1, 99)
            _send(stdout, {"op": "progress", "stage": "loading_model", "percent": pct})

    hb = threading.Thread(target=_heartbeat, daemon=True)
    hb.start()
    try:
        from voxcpm import VoxCPM  # type: ignore[import-not-found]  # noqa: PLC0415

        checkpoint = os.environ.get("OMNIVOICE_VOXCPM_MODEL", "openbmb/VoxCPM2")
        _MODEL = _with_retries(
            lambda: VoxCPM.from_pretrained(checkpoint, load_denoiser=False)
        )
    finally:
        stop.set()
        hb.join(timeout=_HEARTBEAT_S + 1)
    _send(stdout, {"op": "progress", "stage": "loading_model", "percent": 100})
    return _MODEL


def _sample_rate(model) -> int:
    for owner in (model, getattr(model, "tts_model", None)):
        sr = getattr(owner, "sample_rate", None)
        if isinstance(sr, int) and sr > 0:
            return sr
    return VOXCPM2_SAMPLE_RATE


def _at_engine_rate(wav, sample_rate: int):
    """The waveform at VOXCPM2_SAMPLE_RATE. The parent reads the PCM at that
    fixed rate (it trims the tail and labels the audio with it), so a model
    reporting another rate is resampled here rather than mislabelled."""
    if sample_rate == VOXCPM2_SAMPLE_RATE:
        return wav
    import torch  # noqa: PLC0415
    import torchaudio  # noqa: PLC0415

    tensor = torch.as_tensor(wav.detach().cpu() if hasattr(wav, "detach") else wav,
                             dtype=torch.float32).reshape(-1)
    return torchaudio.functional.resample(tensor, sample_rate, VOXCPM2_SAMPLE_RATE)


def _to_pcm_b64(wav) -> tuple[str, int]:
    """A float waveform in [-1, 1] (numpy or torch) as base64 int16 PCM."""
    import numpy as np  # noqa: PLC0415

    if hasattr(wav, "detach"):
        wav = wav.detach().float().cpu().numpy()
    arr = np.asarray(wav, dtype=np.float32).squeeze()
    if arr.ndim > 1:
        raise ValueError(f"expected mono audio (1-D after squeeze), got shape {arr.shape}")
    arr = np.clip(arr, -1.0, 1.0)
    pcm = (arr * 32767.0).astype(np.int16).tobytes()
    return base64.b64encode(pcm).decode("ascii"), int(arr.shape[-1])


def _handle_synthesize(msg: dict, stdout) -> None:
    """One synthesize request. The mapping mirrors VoxCPM2Backend.generate."""
    text = msg.get("text")
    if not text or not isinstance(text, str):
        raise ValueError("synthesize: missing or non-string 'text'")
    ref_audio = msg.get("ref_audio") or None
    if ref_audio and _URL_RE.match(str(ref_audio)):
        raise ValueError(
            "ref_audio must be a local file path; URLs are not accepted (local-first)."
        )
    model = _load_model(stdout)
    description = msg.get("description")
    cfg_value = msg.get("guidance_scale", 2.0)
    timesteps = msg.get("num_step", 10)
    if description and not ref_audio:
        # Voice design: a voice from a text description, no reference clip.
        wav = model.generate(
            text=text,
            voice_description=description,
            cfg_value=cfg_value,
            inference_timesteps=timesteps,
        )
    else:
        instruct = msg.get("instruct")
        ref_text = msg.get("ref_text")
        wav = model.generate(
            text=f"({instruct}){text}" if instruct else text,
            cfg_value=cfg_value,
            inference_timesteps=timesteps,
            reference_wav_path=ref_audio,
            prompt_wav_path=ref_audio if ref_text else None,
            prompt_text=ref_text,
        )
    pcm_b64, n_samples = _to_pcm_b64(_at_engine_rate(wav, _sample_rate(model)))
    _send(stdout, {
        "op": "audio",
        "audio_pcm_b64": pcm_b64,
        "sample_rate": VOXCPM2_SAMPLE_RATE,
        "n_samples": n_samples,
    })


# -- main loop ---------------------------------------------------------------


def main() -> int:
    stdin = sys.stdin.buffer
    # Frames go down a PRIVATE fd, and fd 1 is pointed at stderr (#1428): the
    # libraries this loads print to fd 1 (tqdm, native torch output), and those
    # bytes would otherwise interleave with the length-prefixed frames.
    _frame_fd = os.dup(1)
    os.dup2(2, 1)
    stdout = os.fdopen(_frame_fd, "wb")

    # Ready handshake fires BEFORE any heavy import.
    _send(stdout, {"op": "ready", "engine": "voxcpm2", "sample_rate": VOXCPM2_SAMPLE_RATE})

    while True:
        try:
            msg = _recv(stdin)
        except Exception as exc:  # noqa: BLE001
            _send(stdout, {
                "op": "error",
                "stage": "recv",
                "message": f"{type(exc).__name__}: {exc}",
                "traceback": traceback.format_exc(),
            })
            return 1
        if msg is None:
            return 0
        op = msg.get("op") if isinstance(msg, dict) else None
        try:
            if op == "ping":
                _send(stdout, {"op": "pong", "vram_mb": _measure_vram_mb()})
            elif op == "synthesize":
                _handle_synthesize(msg, stdout)
            elif op == "shutdown":
                return 0
            else:
                _send(stdout, {"op": "error", "stage": "dispatch", "message": f"unknown op: {op!r}"})
        except Exception as exc:  # noqa: BLE001
            _send(stdout, {
                "op": "error",
                "stage": op or "unknown",
                "message": f"{type(exc).__name__}: {exc}",
                "traceback": traceback.format_exc(),
            })


if __name__ == "__main__":
    sys.exit(main())