path: root/app/converter/clients/sglomni.py

"""Client for the SGLang-Omni OpenAI-compatible TTS server.

SGLang-Omni (``sgl-omni serve --model-path <hf-repo>``) hosts one TTS
model per process behind the OpenAI-style ``/v1/audio/speech`` endpoint.
This client speaks that endpoint for every catalog model, resolving the
request shape from the model's voice capability (``backends.sglomni.
catalog``):

  speaker  the voice names a preset shipped with the model (Qwen3-TTS
           CustomVoice speakers; Voxtral preset voices)
  clone    the voice comes from a reference clip sent per request as
           ``ref_audio`` + ``ref_text``. The reference is transcribed
           with a local Whisper backend when no transcript is given (the
           qwen backend's flow). On a loopback server the clip travels
           as a local path the server reads directly; anywhere else it
           is inlined as a base64 data URL, so ``--api-url`` remote
           servers work without any server-side file setup.
  design   the voice is described by instructions (``task_type=
           "VoiceDesign"`` + ``instructions``, Qwen3-TTS VoiceDesign).

Clone-capable models without a reference synthesize their built-in
default voice ("default") unless the catalog marks a reference as
mandatory (Qwen3-TTS Base, dots.tts, ZONOS2 — those refuse at connect).
"""

import base64
import json
import logging
import re
import tempfile
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import List, Optional

from .. import config
from ..audio import concat_audio_files
from ..chunking import split_into_chunks
from .base import (BaseTTSClient, ConversionCancelled,
                   NonRetryableTTSError, resolve_request_seed)
from .languages import normalize_language

logger = logging.getLogger(__name__)

# Response formats the endpoint offers; complete WAV files need no
# sample-rate handling client-side (the header carries it, and models
# differ: 24 kHz Voxtral/Higgs, 44.1 kHz ZONOS2, 48 kHz MOSS Local).
RESPONSE_FORMAT = "wav"

# The voice name the server synthesizes with when the request does not
# pick a preset or clone a reference.
DEFAULT_VOICE = "default"

# Mimetypes for inlined reference audio (data URLs), by file suffix.
_MIME_BY_SUFFIX = {
    ".wav": "audio/wav", ".mp3": "audio/mpeg", ".flac": "audio/flac",
    ".ogg": "audio/ogg", ".aac": "audio/aac", ".m4a": "audio/mp4",
    ".webm": "audio/webm", ".mp4": "audio/mp4",
}

# Error-envelope types the server returns for deterministic request
# problems (bad voice, missing reference, unknown model): the identical
# request fails on every retry, so the chunk loop gives up immediately.
_NON_RETRYABLE_TYPES = ("BadRequestError", "InvalidRequestError",
                        "NotFoundError", "PermissionDeniedError")

# The scheduler's KV-window admission error ("Request requires more tokens
# than the thinker KV cache can hold (input_tokens=684, max_new_tokens=
# 12288, required_tokens=12972, kv_capacity=4095)..."): the server names
# the numbers a refit needs, and upstream classifies the message as a
# deterministic bad request — the identical request fails on every retry,
# so the only useful response is to send a smaller one.
_KV_ADMISSION_MARKER = "thinker KV cache can hold"

# Frames kept below the capacity the server reported, and the smallest
# refitted cap worth generating with (~14 s of speech at 75 fps): below
# the floor the request would truncate almost immediately, so the run
# surfaces guidance instead of near-empty audio.
_KV_FIT_MARGIN = 64
_KV_FIT_FLOOR = 1024


def _is_loopback(url: str) -> bool:
    """True when URL's host is this machine (the server can read local
    reference files by path)."""
    try:
        host = urllib.parse.urlsplit(url).hostname or "127.0.0.1"
    except ValueError:
        return False
    return host in ("127.0.0.1", "localhost", "::1")


def _data_url(path: Path) -> str:
    """PATH's audio bytes as a base64 data URL (for remote servers)."""
    mime = _MIME_BY_SUFFIX.get(path.suffix.lower(), "audio/wav")
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:{mime};base64,{encoded}"


def _http_error_detail(exc: urllib.error.HTTPError) -> str:
    """The error response body as text (empty when it cannot be read)."""
    try:
        return exc.read().decode("utf-8", errors="replace")
    except Exception:
        return ""


def _kv_error_number(detail: str, name: str) -> Optional[int]:
    """The integer NAME=... reports in a KV-window admission message."""
    match = re.search(rf"\b{name}=(\d+)", detail)
    return int(match.group(1)) if match else None


class SgOmniTTSClient(BaseTTSClient):
    """Generates audio chunks through an SGLang-Omni server."""

    def __init__(self, chunks_dir: Path,
                 model: Optional[str] = None,
                 voice: Optional[str] = None,
                 ref_audio: Optional[str] = None,
                 ref_text: Optional[str] = None,
                 skip_transcription: bool = False,
                 instructions: Optional[str] = None,
                 language: Optional[str] = None,
                 api_url: Optional[str] = None,
                 chunk_size: Optional[int] = None,
                 quiet: bool = False, cancel=None):
        super().__init__(chunks_dir, quiet=quiet, cancel=cancel)
        # The catalog entry this run targets (the backend package validates
        # the key; only its repo id and capability are client business).
        from backends.sglomni.catalog import entry_by_key
        from backends.sglomni.constants import DEFAULT_PORT, SERVER_NAME
        from backends.common import port_of
        self.entry = entry_by_key((model or "").strip())
        if self.entry is None:
            raise RuntimeError(
                f"Unknown SGLang-Omni model {model!r} — pick a catalog key "
                "(see Configure Backends → SGLang-Omni or the backend docs).")
        self.api_url = ((api_url or config.SGLOMNI_API_URL).strip()
                        .rstrip("/"))
        self.port = port_of(self.api_url, DEFAULT_PORT)
        self.voice = (voice or "").strip() or None
        self.ref_audio = (ref_audio or "").strip() or None
        self.ref_text = (ref_text or "").strip()
        self.skip_transcription = skip_transcription
        self.instructions = (instructions or "").strip()
        # Per-run sub-request word cap (the pre-flight chunk popup's "set
        # chunk" answer); None follows config.CHUNK_SIZE.
        self.chunk_size = chunk_size
        # A KV-window capacity the server taught us via an admission
        # rejection (None = none learned): later requests keep their
        # max_new_tokens under it. See _kv_admission_fit.
        self._kv_fit = None
        # Seed sent with every request: config.SEED as-is, or (with
        # CONSTANT_SEED and SEED < 0) one random value drawn per run and
        # reused for every chunk so the voice stays consistent across
        # chunk boundaries. Only sent to models that accept a
        # request-scoped seed (Voxtral rejects it outright), and only
        # when a concrete seed is in play (a negative one means "re-sample
        # every generation", so there is nothing to send).
        seed = resolve_request_seed() if self.entry.supports_seed else None
        self._seed = seed if (seed is not None and seed >= 0) else None
        if language is None:
            language = config.LANGUAGE
        self.language = normalize_language(language)
        self._check_connect_inputs()
        self._connect()

    # ------------------------------------------------------------------
    # Connection
    # ------------------------------------------------------------------

    def _check_connect_inputs(self) -> None:
        """Validate the voice inputs against the model's capability."""
        entry = self.entry
        if entry.capability == "design" and not self.instructions:
            raise RuntimeError(
                f"{entry.label} designs the voice from an instruction: "
                'pass --instructions "..." describing the voice.')
        if entry.capability == "clone" and entry.requires_reference \
                and not self.ref_audio:
            raise RuntimeError(
                f"{entry.label} requires reference audio to narrate: "
                "pass --clone PATH (a .wav reference clip), or pick a "
                "model that synthesizes without one.")
        if entry.capability == "speaker" and self.ref_audio:
            self._report(f"[WARNING] --clone is ignored with {entry.label}: "
                         "it voices text with its built-in presets.")
            self.ref_audio = None
        elif self.ref_audio and not Path(self.ref_audio).is_file():
            raise RuntimeError(
                f"Reference audio not found: {self.ref_audio}")
        if entry.speakers and self.voice \
                and self.voice not in entry.speakers:
            self._report(
                f"[WARNING] Voice {self.voice!r} is not one of "
                f"{entry.label}'s presets ({', '.join(entry.speakers)}); "
                "the server will reject it if it does not know the name.")
    def _connect(self) -> None:
        """Verify the server is up, healthy, and hosting the expected model.

        The managed-server lifecycle (managed.ensure_running / the run
        view's autostart) normally boots exactly the selected model; a
        foreign server hosting something else — or a remote one the form
        could not classify — fails here with both model names instead of
        producing per-chunk failures later.
        """
        entry, url = self.entry, self.api_url
        try:
            payload = self._fetch_json("/health", timeout=10)
        except Exception as exc:
            raise RuntimeError(
                f"SGLang-Omni server not reachable at {url}: {exc}. Start "
                "the sgl-omni server first (the CLI and the hub start the "
                "managed instance automatically when the backend is "
                "installed), or point --api-url at a running server."
            ) from exc
        if not isinstance(payload, dict) \
                or payload.get("status") != "healthy":
            raise RuntimeError(
                f"The SGLang-Omni server at {url} is not healthy yet "
                f"(health: {payload}). Wait for it to finish booting and "
                "retry.")
        served = self._served_model()
        if served is not None and served != entry.repo:
            raise RuntimeError(
                f"The SGLang-Omni server at {url} hosts {served}, but "
                f"this run selected {entry.repo}. Restart it with that "
                "model (the managed server restarts automatically), or "
                "pick the hosted model for this run.")
        self._resolve_reference_text()
        mode = {"speaker": "built-in presets",
                "clone": "voice cloning",
                "design": "voice design"}[entry.capability]
        self._report(f"[OK] Connected to SGLang-Omni at {url} "
                     f"({entry.label}, {mode})")

    def _served_model(self) -> Optional[str]:
        """The repo id the server hosts (None when it cannot be read)."""
        try:
            payload = self._get_json("/v1/models", timeout=10)
        except Exception:
            return None
        entries = (payload or {}).get("data")
        if isinstance(entries, list) and entries \
                and isinstance(entries[0], dict):
            return entries[0].get("id")
        return None

    def _resolve_reference_text(self) -> None:
        """Resolve the clone reference transcript: explicit text, then a
        local Whisper transcription."""
        if self.entry.capability != "clone" or not self.ref_audio:
            return
        if not self.ref_text and not self.skip_transcription:
            self._report("[INFO] Transcribing reference audio for voice "
                         "cloning...")
            from .transcribe import transcribe_reference_audio
            self.ref_text = transcribe_reference_audio(self.ref_audio) or ""
        if self.ref_text:
            self._report(f"[OK] Reference text: {self.ref_text}")
        else:
            self._report("[WARNING] No reference transcript: cloning runs "
                         "without ref_text, which lowers quality for "
                         "models that use it. Pass --transcription \"...\" "
                         "for best results.")

    # ------------------------------------------------------------------
    # HTTP requests
    # ------------------------------------------------------------------

    def _fetch_json(self, path: str, timeout: int = 10) -> dict:
        """GET PATH and parse the JSON body, raising on connection errors.

        Unlike _get_json this surfaces unreachable servers to the caller —
        the connect flow needs to tell "nothing is listening" (start the
        server) apart from "listening but still booting" (wait).
        """
        url = f"{self.api_url}{path}"
        with urllib.request.urlopen(url, timeout=timeout) as response:
            payload = json.loads(response.read().decode("utf-8"))
        return payload if isinstance(payload, dict) else {}

    def _get_json(self, path: str, timeout: int = 10) -> Optional[dict]:
        """GET PATH and parse a JSON object, or None on any error."""
        try:
            return self._fetch_json(path, timeout=timeout)
        except (OSError, ValueError):
            return None

    def _ref_audio_value(self) -> str:
        """The ref_audio request value: a local path on a loopback server
        (the server reads the file directly), else a base64 data URL."""
        path = Path(self.ref_audio)
        if not path.is_file():
            raise RuntimeError(
                f"Reference audio not found: {self.ref_audio}")
        if _is_loopback(self.api_url):
            return str(path.resolve())
        return _data_url(path)

    def _request_payload(self, text: str) -> dict:
        """The /v1/audio/speech JSON body for one sub-chunk."""
        entry = self.entry
        payload = {
            "model": entry.repo,
            "voice": self.voice or DEFAULT_VOICE,
            "input": text,
            "response_format": RESPONSE_FORMAT,
            "language": self.language,
        }
        if self._seed is not None:
            payload["seed"] = self._seed
        if entry.max_new_tokens is not None:
            # Models whose engine caps a request below what a full
            # sub-chunk can narrate (Zonos2's 1024-frame default is ~12 s):
            # raise the ceiling per request. Generation still stops at
            # natural EOS, so an unused margin costs nothing. A capacity
            # learned from an admission rejection (Higgs pins the window)
            # keeps later requests under it too.
            cap = entry.max_new_tokens
            if self._kv_fit is not None:
                cap = min(cap, self._kv_fit)
            payload["max_new_tokens"] = cap
        if entry.capability == "design":
            payload["task_type"] = "VoiceDesign"
            payload["instructions"] = self.instructions
        elif entry.capability == "clone" and self.ref_audio:
            payload["ref_audio"] = self._ref_audio_value()
            if self.ref_text:
                payload["ref_text"] = self.ref_text
        return payload

    def _request_wav(self, text: str) -> bytes:
        """POST one sub-chunk and return the complete WAV bytes."""
        payload = self._request_payload(text)
        try:
            return self._post_speech(payload)
        except urllib.error.HTTPError as exc:
            # A KV-window rejection is deterministic (upstream maps it to a
            # bad request): refit the generation cap to the capacity the
            # server reported and resend once before surfacing anything.
            detail = _http_error_detail(exc)
            fitted = (self._kv_admission_fit(detail)
                      if "max_new_tokens" in payload else None)
            if fitted is not None:
                try:
                    return self._post_speech(
                        dict(payload, max_new_tokens=fitted))
                except urllib.error.HTTPError as retry_exc:
                    exc = retry_exc
                    detail = _http_error_detail(exc)
            raise self._request_error(exc.code, detail) from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(
                f"SGLang-Omni request failed: {exc.reason}") from exc

    def _post_speech(self, payload: dict) -> bytes:
        """POST PAYLOAD to /v1/audio/speech; HTTPErrors propagate raw."""
        url = f"{self.api_url}/v1/audio/speech"
        request = urllib.request.Request(
            url, data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"}, method="POST")
        try:
            with urllib.request.urlopen(request,
                                        timeout=config.API_TIMEOUT) as response:
                wav = response.read()
        except urllib.error.HTTPError:
            # Re-raise raw (HTTPError subclasses URLError): the caller maps
            # it — and refits KV-window rejections — from the status code.
            raise
        except urllib.error.URLError as exc:
            raise RuntimeError(
                f"SGLang-Omni request failed: {exc.reason}") from exc
        if not wav:
            raise RuntimeError("SGLang-Omni server returned empty audio")
        return wav

    def _kv_admission_fit(self, detail: str) -> Optional[int]:
        """A refitted max_new_tokens for a KV-window rejection, or None.

        DETAIL is the error response body. The server's message names the
        request's prompt length and the KV window it must fit; the refit
        keeps a small margin below the window, is remembered for this
        client's remaining sub-requests, and the refitted request carries
        it. When the window leaves less than a useful minimum after the
        prompt (a very long reference clip), the run fails with guidance
        instead of near-empty audio.
        """
        if _KV_ADMISSION_MARKER not in detail:
            return None
        input_tokens = _kv_error_number(detail, "input_tokens")
        kv_capacity = _kv_error_number(detail, "kv_capacity")
        if input_tokens is None or kv_capacity is None:
            return None
        fitted = kv_capacity - input_tokens - _KV_FIT_MARGIN
        if fitted < _KV_FIT_FLOOR:
            raise NonRetryableTTSError(
                f"SGLang-Omni rejected the request: the model's KV window "
                f"({kv_capacity} tokens) leaves {fitted} frames after this "
                f"request's prompt ({input_tokens} tokens) — too little to "
                "narrate anything useful. Use a shorter reference clip or "
                "a smaller Chunk Size setting; the server caps prompt plus "
                "generation at that window for every request.")
        if self._kv_fit is not None:
            fitted = min(fitted, self._kv_fit)
        self._kv_fit = fitted
        return fitted

    def _request_error(self, status: int, detail: str) -> Exception:
        """Map the OpenAI-style error envelope to the retry decision.

        A 4xx envelope (BadRequestError et al.) is deterministic — the
        identical request fails identically on every attempt — so it
        surfaces as NonRetryableTTSError and the chunk loop aborts with
        the server's message; anything else stays retryable.
        """
        message = detail[:500] or f"HTTP {status}"
        kind = None
        try:
            envelope = json.loads(detail)
            error = envelope.get("error")
            if isinstance(error, dict):
                message = str(error.get("message") or message)
                kind = error.get("type")
        except ValueError:
            pass
        if 400 <= status < 500 and (kind is None
                                    or kind in _NON_RETRYABLE_TYPES):
            return NonRetryableTTSError(
                f"SGLang-Omni rejected the request (HTTP {status}): "
                f"{message}")
        return RuntimeError(
            f"SGLang-Omni server returned HTTP {status}: {message}")

    # ------------------------------------------------------------------
    # Chunk generation
    # ------------------------------------------------------------------

    def generate_chunk(self, text: str, chunk_num: int) -> Optional[str]:
        """Generate one audio chunk; returns its path in the chunks folder.

        The text is split into sub-requests of at most ``chunk_size`` words
        each — the per-run cap the pre-flight chunk popup sets for models
        whose engine cannot narrate a full CHUNK_SIZE sub-chunk (Higgs),
        else ``config.CHUNK_SIZE`` (the book-level chunker normally
        guarantees this already; the split is defense in depth against
        pathological input such as a punctuation-free run of text) — and
        the returned WAV files are concatenated into one chunk file.
        """
        try:
            sub_chunks = split_into_chunks(
                text, max_words=self.chunk_size or config.CHUNK_SIZE)
            if not sub_chunks:
                raise RuntimeError("No text to synthesize")

            with self._chunk_heartbeat(chunk_num):
                wav_parts: List[bytes] = [
                    self._request_wav(sub_text) for sub_text in sub_chunks]

            output_path = self._chunk_path(chunk_num, ".wav")
            if len(wav_parts) == 1:
                output_path.write_bytes(wav_parts[0])
            else:
                # Several sub-request WAVs: concatenate through the shared
                # ffmpeg path (each part is a complete file with headers).
                with tempfile.TemporaryDirectory(
                        prefix="sglomni_parts_") as parts_dir:
                    part_paths: List[Path] = []
                    for index, wav in enumerate(wav_parts, 1):
                        part = Path(parts_dir) / f"part_{index:02d}.wav"
                        part.write_bytes(wav)
                        part_paths.append(part)
                    concat_audio_files(part_paths, output_path)

            logger.debug("Chunk %d generated (%d sub-request(s))",
                         chunk_num, len(wav_parts))
            return str(output_path)

        except ConversionCancelled:
            raise
        except Exception as exc:
            logger.error("SGLang-Omni chunk processing failed for chunk "
                         "%d: %s", chunk_num, exc)
            return None