path: root/app/backends/sglomni/catalog.py

"""The sglang-omni model catalog: what can be installed, hosted, and how.

One server process hosts one model (`sgl-omni serve --model-path <hf-repo>
--config <yaml>`), so each entry is one launchable unit — unlike audio.cpp's
multi-entry server.json. Every entry states how its voice is supplied
(`capability`), whether a reference clip is mandatory for narration
(`requires_reference`), and which model-companion packages (`extras`) must
be pip-installed into the backend venv before its server will start.

The catalog is static knowledge about sglang-omni's supported TTS models
(v0.1.4-era), not a live query: the server's /v1/models answers only which
model is currently hosted. Entries carry the HuggingFace repo id verbatim —
it is both the `--model-path` value and the model name /v1/models reports,
so a running server is matched back to its entry by that id. ZONOS2 is the
one entry without a config file (its GPU-conditional bf16 fallback lives in
`bf16_config`); every other entry launches with a vendored config.

Voice capabilities (mirroring the audio.cpp client's vocabulary):
  speaker  the voice is a named preset shipped with the model (Qwen3-TTS
           CustomVoice speaker table; Voxtral's preset voices)
  clone    the voice comes from a reference clip (per-request ref_audio +
           ref_text, transcribed with Whisper when not provided)
  design   the voice is described by instructions (Qwen3-TTS VoiceDesign)
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

# Voice capabilities (same words the audio.cpp client uses).
CAPABILITY_SPEAKER = "speaker"
CAPABILITY_CLONE = "clone"
CAPABILITY_DESIGN = "design"

# A model-companion pip package: (requirement spec, no_deps). The no-deps
# flag mirrors upstream install instructions that must not drag in their
# own (conflicting) transformer pins — the Qwen3-TTS companion pins
# Transformers 4 against sglang-omni's Transformers 5 stack.
Extra = Tuple[str, bool]


@dataclass(frozen=True)
class ModelEntry:
    """One installable/hostable sglang-omni TTS model."""
    key: str                  # catalog id used by --model and the forms
    label: str                # human-readable name (menus, status lines)
    repo: str                 # HuggingFace repo id == --model-path value
    config: Optional[str]     # vendored config yaml file name, None = none
    capability: str           # speaker / clone / design
    requires_reference: bool  # clone models: narration needs ref audio
    extras: Tuple[Extra, ...] = field(default=())
    system_dep: Optional[str] = None   # system binary the extras need
    system_hint: Optional[str] = None  # remediation when the binary is absent
    speakers: Optional[Tuple[str, ...]] = None  # preset voices (speaker)
    supports_seed: bool = False  # request-scoped seed accepted (Qwen3-TTS Base)
    notes: str = ""           # one-line description (documentation)
    # The model's DEFAULT pipeline dynamically quantizes its MoE experts to
    # FP8 at load time (sglang-omni's zonos2 config hardcodes it) — a Triton
    # fp8e4nv kernel that only compiles on compute capability 8.9+ (RTX
    # 4090/5090, Hopper). On older GPUs the server dies mid-boot; BF16_CONFIG
    # is the vendored config that turns FP8 off so the model runs in bf16
    # (~2x the MoE VRAM) there instead.
    fp8_moe: bool = False
    fp8_min_compute_capability: Optional[Tuple[int, int]] = None
    bf16_config: Optional[str] = None
    # max_new_tokens sent with every /v1/audio/speech request (None = send
    # nothing and leave the server default in charge). Some engines cap one
    # request below what a full CHUNK_SIZE sub-chunk can narrate — ZONOS2's
    # AR engine defaults to 1024 audio frames (44100 Hz / 512 hop =
    # 86.13 fps, ~12 s) and Higgs's to 2048 frames (75 fps, ~27 s, and
    # per-request values are clamped to the engine cap, so its vendored
    # config raises the cap too) — and silently truncate longer text.
    # 12288 frames covers a full 250-word CHUNK_SIZE sub-chunk on ZONOS2
    # (whose KV pools are >= 28083 tokens on every GPU that can host it)
    # but NOT on Higgs: its thinker engine pins context_length at 4096,
    # and the scheduler rejects any request whose prompt tokens plus
    # max_new_tokens exceed that window (kv_capacity=4095, on every GPU —
    # the pool side is never the binding constraint there).
    max_new_tokens: Optional[int] = None
    # The largest sub-request (words) the model can narrate within its
    # admission window, for models whose engine cannot cover a full
    # CHUNK_SIZE sub-chunk (None = no cap; config.CHUNK_SIZE stands).
    # Higgs: 3000 frames ≈ 40 s at 75 fps ≈ 80 words of narration, and
    # the run pre-flight offers to clamp CHUNK_SIZE for the run (the
    # client's adaptive retry still rescues requests the window rejects).
    chunk_words: Optional[int] = None


# The Qwen3-TTS CustomVoice speaker table — the same built-in speakers the
# qwen-tts demo and audio.cpp's CustomVoice entry expose (kept as a local
# copy so this catalog stays importable without the converter package; the
# test suite asserts it matches converter.clients.speakers).
QWEN_CUSTOMVOICE_SPEAKERS: Tuple[str, ...] = (
    "Vivian", "Serena", "Uncle_Fu", "Dylan", "Eric",
    "Ryan", "Aiden", "Ono_Anna", "Sohee")

# The Qwen3-TTS companions follow the upstream TTS guide exactly: the
# qwen-tts demo package installs WITHOUT dependencies (its Transformers 4
# pin would replace sglang-omni's pinned 5.12 stack; sglang-omni shims the
# API differences), as do sox/einops (a normal sox resolve pulls numpy past
# the ceiling numba imposes). The sox *binary* is a system package the
# wizard detects and guides (it is required at synthesis time).
_QWEN_EXTRAS: Tuple[Extra, ...] = (
    ("sox", True), ("einops", True), ("qwen-tts==0.1.1", True))
_SOX_HINT = ("install the sox system package (e.g. sudo pacman -S sox, "
             "sudo apt install sox, brew install sox)")
# The Fish Audio and ZONOS2 pipelines use the Descript DAC codec, which
# upstream installs WITH dependencies (nothing conflicts).
_DAC_EXTRAS: Tuple[Extra, ...] = (
    ("descript-audiotools==0.7.2", False),
    ("descript-audio-codec==1.0.0", False))

# Companion distributions whose top-level import name differs from the pip
# name's plain dash-to-underscore normalization (verified against their
# top_level.txt). Anything absent here normalizes: qwen-tts -> qwen_tts.
_EXTRA_IMPORT_OVERRIDES = {
    "descript-audiotools": "audiotools",
    "descript-audio-codec": "dac",
}


def extra_import_name(spec: str) -> str:
    """The Python module an extras requirement SPEC provides.

    Takes the distribution name portion of the pip requirement (so
    ``qwen-tts==0.1.1`` -> ``qwen_tts``) — the name a venv probe must
    import to prove the companion is installed."""
    base = spec.split("=")[0].split("<")[0].split(">")[0].strip()
    if base in _EXTRA_IMPORT_OVERRIDES:
        return _EXTRA_IMPORT_OVERRIDES[base]
    return base.replace("-", "_")

ENTRIES: Tuple[ModelEntry, ...] = (
    ModelEntry(
        key="qwen3_tts_0_6b_customvoice",
        label="Qwen3-TTS 0.6B CustomVoice",
        repo="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
        config="qwen3_tts_0_6b_customvoice.yaml",
        capability=CAPABILITY_SPEAKER,
        requires_reference=False,
        extras=_QWEN_EXTRAS, system_dep="sox", system_hint=_SOX_HINT,
        speakers=QWEN_CUSTOMVOICE_SPEAKERS,
        notes="built-in speakers, lightest model",
    ),
    ModelEntry(
        key="qwen3_tts_0_6b_base",
        label="Qwen3-TTS 0.6B Base",
        repo="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
        config="qwen3_tts_0_6b.yaml",
        capability=CAPABILITY_CLONE,
        requires_reference=True,
        extras=_QWEN_EXTRAS, system_dep="sox", system_hint=_SOX_HINT,
        supports_seed=True,
        notes="voice cloning from a reference clip",
    ),
    ModelEntry(
        key="qwen3_tts_1_7b_base",
        label="Qwen3-TTS 1.7B Base",
        repo="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
        config="qwen3_tts_1_7b.yaml",
        capability=CAPABILITY_CLONE,
        requires_reference=True,
        extras=_QWEN_EXTRAS, system_dep="sox", system_hint=_SOX_HINT,
        supports_seed=True,
        notes="voice cloning, higher quality",
    ),
    ModelEntry(
        key="qwen3_tts_1_7b_voicedesign",
        label="Qwen3-TTS 1.7B VoiceDesign",
        repo="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
        config="qwen3_tts_1_7b_voicedesign.yaml",
        capability=CAPABILITY_DESIGN,
        requires_reference=False,
        extras=_QWEN_EXTRAS, system_dep="sox", system_hint=_SOX_HINT,
        notes="voice described by instructions",
    ),
    ModelEntry(
        key="higgs_audio_v3_tts",
        label="Higgs Audio v3 TTS",
        repo="bosonai/higgs-audio-v3-tts-4b",
        # The vendored config trims the engine's VRAM budget (the upstream
        # pipeline budgets 0.98 of the card across its colocated stages,
        # which OOMs on 24 GB cards) and raises the engine's 2048-frame
        # generation cap (see the yaml header).
        config="higgs_audio_v3_tts.yaml",
        capability=CAPABILITY_CLONE,
        requires_reference=False,
        # The engine's 2048-frame default is ~27 s of speech at the codec's
        # 75 fps; requests are clamped to the engine cap server-side, so the
        # yaml raises the cap and every request carries 3000 frames (~40 s)
        # — the most the admission window allows: upstream pins the thinker
        # engine's context_length at 4096, and the scheduler rejects any
        # request whose prompt (including the reference-audio tokens) plus
        # max_new_tokens exceeds it. 3000 frames leaves ~1095 tokens of
        # prompt headroom (an 80-word chunk with a 20.5 s reference
        # measured 684). Sub-requests cap at 80 words so the text fits the
        # window too — the pre-flight offers to clamp CHUNK_SIZE for the
        # run, and the client refits rejected requests to whatever the
        # server reports as its capacity.
        max_new_tokens=3000,
        chunk_words=80,
        notes="zero-shot narration, cloning from a reference clip",
    ),
    ModelEntry(
        key="moss_tts",
        label="MOSS-TTS v1.5",
        repo="OpenMOSS-Team/MOSS-TTS-v1.5",
        config="moss_tts.yaml",
        capability=CAPABILITY_CLONE,
        requires_reference=True,
        notes="voice cloning from a reference clip",
    ),
    ModelEntry(
        key="moss_tts_local",
        label="MOSS-TTS Local v1.5",
        repo="OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5",
        config="moss_tts_local.yaml",
        capability=CAPABILITY_CLONE,
        requires_reference=False,
        notes="48 kHz, narration without a reference or cloning",
    ),
    ModelEntry(
        key="voxtral_tts",
        label="Voxtral TTS 4B",
        repo="mistralai/Voxtral-4B-TTS-2603",
        config="voxtral_tts.yaml",
        capability=CAPABILITY_SPEAKER,
        requires_reference=False,
        notes="preset voices, text-only requests",
    ),
    ModelEntry(
        key="dots_tts_mf",
        label="dots.tts (MeanFlow)",
        repo="dots-studio/dots.tts-mf",
        config="dots_tts.yaml",
        capability=CAPABILITY_CLONE,
        requires_reference=True,
        notes="voice cloning, requires reference audio + transcript",
    ),
    ModelEntry(
        key="fish_s2_pro",
        label="Fish Speech S2-Pro",
        repo="fishaudio/s2-pro",
        config="s2pro_tts.yaml",
        capability=CAPABILITY_CLONE,
        requires_reference=False,
        extras=_DAC_EXTRAS,
        notes="zero-shot narration or cloning from a reference clip",
    ),
    ModelEntry(
        key="zonos2",
        label="ZONOS2",
        repo="Zyphra/zonos2",
        config=None,
        capability=CAPABILITY_CLONE,
        requires_reference=True,
        extras=_DAC_EXTRAS,
        notes="voice cloning, 44.1 kHz DAC vocoder",
        fp8_moe=True,
        fp8_min_compute_capability=(8, 9),
        bf16_config="zonos2_bf16.yaml",
        max_new_tokens=12288,
    ),
)

_BY_KEY: Dict[str, ModelEntry] = {entry.key: entry for entry in ENTRIES}
_BY_REPO: Dict[str, ModelEntry] = {entry.repo: entry for entry in ENTRIES}


def entry_by_key(key: str) -> Optional[ModelEntry]:
    """The catalog entry for a catalog KEY, or None."""
    return _BY_KEY.get(key)


def entry_by_repo(repo: str) -> Optional[ModelEntry]:
    """The catalog entry hosting REPO (a /v1/models id), or None."""
    return _BY_REPO.get(repo)


def entries_by_keys(keys) -> List[ModelEntry]:
    """The ENTRIES for KEYS, in catalog order (unknown keys dropped)."""
    wanted = set(keys)
    return [entry for entry in ENTRIES if entry.key in wanted]


def config_path(entry: ModelEntry):
    """The vendored config yaml path for ENTRY, or None when it runs on
    --model-path alone (ZONOS2). A declared-but-missing file means
    a broken install — callers treat that like a missing entry."""
    return _config_file(entry.config)


def fallback_config_path(entry: ModelEntry):
    """The vendored bf16 config yaml for ENTRY (None when it has none).

    The config a GPU below the entry's fp8_min_compute_capability launches
    with instead of the model's default FP8-quantized pipeline."""
    return _config_file(entry.bf16_config)


def _config_file(name: Optional[str]):
    """CONFIGS_DIR/NAME, or None when NAME is None (no config for ENTRY)."""
    if name is None:
        return None
    from .constants import CONFIGS_DIR
    return CONFIGS_DIR / name


def install_tree_families(available: List[ModelEntry]) -> List[dict]:
    """The checkbox-tree shape for the model picker, grouped by upstream.

    One family per model origin with one option per model, so the picker
    reads like the audio.cpp one (a collapsed tree of related packages).
    No ``detail`` is set: the tree's status line under the buttons would
    only repeat the catalog keys under the cursor. AVAILABLE filters what
    is shown (already-installed models stay listed so they can be
    re-checked or repaired).
    """
    groups: List[Tuple[str, List[ModelEntry]]] = [
        ("Qwen3-TTS (Qwen)", [e for e in available if e.repo.startswith("Qwen/")]),
        ("Boson AI", [e for e in available if e.repo.startswith("bosonai/")]),
        ("OpenMOSS", [e for e in available if e.repo.startswith("OpenMOSS-Team/")]),
        ("Mistral AI", [e for e in available if e.repo.startswith("mistralai/")]),
        ("dots.studio", [e for e in available if e.repo.startswith("dots-studio/")]),
        ("Fish Audio", [e for e in available if e.repo.startswith("fishaudio/")]),
        ("Zyphra", [e for e in available if e.repo.startswith("Zyphra/")]),
    ]
    families = []
    for label, entries in groups:
        if not entries:
            continue
        families.append({
            "label": label,
            "options": [
                {"key": entry.key, "label": entry.label,
                 "recommended": False}
                for entry in entries
            ],
        })
    return families