diff options
Diffstat (limited to 'app/backends')
25 files changed, 1938 insertions, 63 deletions
diff --git a/app/backends/__init__.py b/app/backends/__init__.py index 89d95e2..5a1a602 100644 --- a/app/backends/__init__.py +++ b/app/backends/__init__.py @@ -1,12 +1,12 @@ """Registry of the TTS backends the audiobook generator can talk to. -Each backend (audio.cpp, qwen, faster) lives in its own module and owns -its setup wizard, its status detection, its uninstaller, and the launch -command it prints once configured. This package aggregates them into a single -registry so ``audiobook.py``'s TUI hub and future tools can iterate backends -without hardcoding their names: ``backends.detect_all()`` reports which are set -up (and whether their server is currently running), and the registry -drives the hub's "Configure Backends" menu. +Each backend (audio.cpp, qwen, faster, SGLang-Omni) lives in its own module +and owns its setup wizard, its status detection, its uninstaller, and the +launch command it prints once configured. This package aggregates them into +a single registry so ``audiobook.py``'s TUI hub and future tools can iterate +backends without hardcoding their names: ``backends.detect_all()`` reports +which are set up (and whether their server is currently running), and the +registry drives the hub's "Configure Backends" menu. The registry is built lazily on the first call to ``get``/``detect_all``/ ``detect`` (not at package import time), because the backend modules pull @@ -59,12 +59,19 @@ class ServerSpec: waits for the server to answer HTTP with that identity — not merely to accept TCP connections — so "listening but still starting" servers are caught. None keeps the plain TCP-connect readiness check. + + START_TIMEOUT overrides the shared ``servers.SERVER_START_TIMEOUT`` for + this spec (seconds): the sglang-omni pipeline (preprocessing, TTS + generation, vocoder stages) boots far slower than the single-process + servers, and its first start may also pull companion weights from the + HuggingFace cache. None uses the shared default. """ name: str url: str argv: List[str] cwd: Optional[Path] = None identity: Optional[str] = None + start_timeout: Optional[int] = None @dataclass @@ -189,7 +196,7 @@ def _build_registry() -> None: """Import the backend modules and wire up REGISTRY (once).""" if REGISTRY: return - from . import audiocpp, faster, qwen + from . import audiocpp, faster, qwen, sglomni REGISTRY.append(BackendInfo( key="audiocpp", @@ -216,6 +223,15 @@ def _build_registry() -> None: uninstall=faster.uninstall, update=faster.update, )) + REGISTRY.append(BackendInfo( + key="sglomni", + label="SGLang-Omni", + detect=sglomni.detect, + setup_screen=sglomni.setup_screen, + uninstall=sglomni.uninstall, + update=sglomni.update, + configure_screen=sglomni.models_screen, + )) for info in REGISTRY: _BY_KEY[info.key] = info diff --git a/app/backends/common.py b/app/backends/common.py index 24746b6..811f13a 100644 --- a/app/backends/common.py +++ b/app/backends/common.py @@ -371,7 +371,8 @@ def write_prompt_text(wav_dir: Path, def run_console_subprocess(argv: List[str], cwd: Optional[Path] = None, *, emit=None, cancel=None, on_cancel=None, - stall_timeout: Optional[float] = None) -> int: + stall_timeout: Optional[float] = None, + env: Optional[Dict[str, str]] = None) -> int: """Run a subprocess, streaming output to the console or to EMIT. With EMIT None the child inherits the real terminal and its output @@ -393,13 +394,17 @@ def run_console_subprocess(argv: List[str], cwd: Optional[Path] = None, returned so callers can report a stall distinctly from a plain failure. None (the default) waits forever, as before. + ENV, when given, replaces the child's environment wholesale (e.g. + provisioning helpers pointing uv at a project-local interpreter + install dir); None inherits the parent's. + Returns the process exit code. """ import subprocess if emit is None: try: - result = subprocess.run(argv, - cwd=str(cwd) if cwd is not None else None) + result = subprocess.run( + argv, cwd=str(cwd) if cwd is not None else None, env=env) except OSError as exc: print(f"[ERROR] Could not run {' '.join(argv)}: {exc}") return 1 @@ -408,6 +413,8 @@ def run_console_subprocess(argv: List[str], cwd: Optional[Path] = None, popen_kwargs = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT} if cwd is not None: popen_kwargs["cwd"] = str(cwd) + if env is not None: + popen_kwargs["env"] = env if sys.platform == "win32": popen_kwargs["creationflags"] = \ subprocess.CREATE_NEW_PROCESS_GROUP # type: ignore[attr-defined] @@ -673,23 +680,28 @@ def run_console_subprocess_quiet(argv: List[str], def pip_install(packages: List[str], *, emit=None, cancel=None, env_dir: Optional[Path] = None, - upgrade: bool = False) -> int: + upgrade: bool = False, + extra_args: Optional[List[str]] = None, + interpreter: Optional[Path] = None) -> int: """pip install PACKAGES into a managed venv. Returns exit code. Delegates to ``backends.envs.pip_install`` so backend TTS packages are installed into their dedicated tool-managed environments (``envs/tts`` - default; ``envs/qwen`` / ``envs/faster`` via ENV_DIR) rather than into - whatever interpreter happens to be running the wizard — and never two - conflicting stacks into the same env. With UPGRADE pip runs with - ``-U`` (the backend update action's freshness check: pip only installs - when a newer version resolves, else reports "already satisfied"). - With EMIT given (the in-TUI task view) pip runs with its output streamed - into EMIT; CANCEL aborts it. The import is local to avoid a circular - import (envs imports this module). + default; ``envs/qwen`` / ``envs/faster`` / ``envs/sglomni`` via ENV_DIR) + rather than into whatever interpreter happens to be running the wizard + — and never two conflicting stacks into the same env. With UPGRADE pip + runs with ``-U`` (the backend update action's freshness check: pip only + installs when a newer version resolves, else reports "Requirement + already satisfied"). EXTRA_ARGS pass through to pip verbatim (e.g. + ``--pre``, ``--no-deps``); INTERPRETER builds a missing env from that + Python. With EMIT given (the in-TUI task view) pip runs with its output + streamed into EMIT; CANCEL aborts it. The import is local to avoid a + circular import (envs imports this module). """ from backends import envs return envs.pip_install(packages, emit=emit, cancel=cancel, - env_dir=env_dir, upgrade=upgrade) + env_dir=env_dir, upgrade=upgrade, + extra_args=extra_args, interpreter=interpreter) def pip_uninstall(packages: List[str], *, emit=None, diff --git a/app/backends/envs.py b/app/backends/envs.py index beef23b..98f1cc1 100644 --- a/app/backends/envs.py +++ b/app/backends/envs.py @@ -28,10 +28,11 @@ import hashlib import json import os import re +import shutil import subprocess import sys from pathlib import Path -from typing import Dict, List, Optional, Tuple +from typing import Dict, Iterable, List, Optional, Tuple from backends import common @@ -41,9 +42,18 @@ TTS_ROOT = Path(__file__).resolve().parent.parent.parent # One app venv for the app requirements; one venv per pip-installed TTS # backend. The per-backend split keeps qwen-tts's transformers 4 pin away # from faster-qwen3-tts's transformers 5 requirement (and away from the app). +# The sglang-omni backend additionally pins the interpreter *version* +# (sglang-omni requires Python >=3.10,<3.13), so its venv may be created +# from a different interpreter than the launching one (see +# compatible_interpreter / provision_env_with_uv). ENV_DIR = TTS_ROOT / "app" / "envs" / "tts" QWEN_ENV_DIR = TTS_ROOT / "app" / "envs" / "qwen" FASTER_ENV_DIR = TTS_ROOT / "app" / "envs" / "faster" +SGLOMNI_ENV_DIR = TTS_ROOT / "app" / "envs" / "sglomni" +# uv-managed standalone CPython installs (downloaded on demand when no +# 3.10-3.12 interpreter exists on the system) live here, inside the +# project, rather than uv's default ~/.local/share/uv/python. +PYTHON_INSTALL_DIR = TTS_ROOT / "app" / "envs" / "pythons" REQUIREMENTS_PATH = TTS_ROOT / "requirements.txt" # requirements.txt lines whose comment starts with this tag are installed @@ -103,21 +113,26 @@ def is_managed_env() -> bool: return False -def create_env(env_dir: Optional[Path] = None) -> int: - """Create ENV (an env dir, defaulting to the app one) with the launching - interpreter (inherits its version). +def create_env(env_dir: Optional[Path] = None, + interpreter: Optional[Path] = None) -> int: + """Create ENV (an env dir, defaulting to the app one). + + With INTERPRETER the venv is built from that Python (an absolute path, + e.g. a 3.12 found on PATH for the sglang-omni backend); otherwise the + launching interpreter is used, so the env inherits its version. pip is bootstrapped inside the venv by ensurepip. Returns the ``python -m venv`` exit code; a non-zero result is reported with platform remediation. """ target = env_dir if env_dir is not None else ENV_DIR + launcher = Path(interpreter) if interpreter is not None else sys.executable print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━") print(" Setting up your environment for the first time...") print(" This may take a minute.") print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━") print(f"[INFO] creating managed environment at {target}...") rc = common.run_console_subprocess( - [sys.executable, "-m", "venv", str(target)]) + [str(launcher), "-m", "venv", str(target)]) if rc != 0: print(f"[ERROR] python -m venv failed (exit {rc}).") if _is_windows(): @@ -198,7 +213,9 @@ def install_requirements(skip_optional: bool = False) -> int: def pip_install(packages: List[str], *, emit=None, cancel=None, env_dir: Optional[Path] = None, - upgrade: bool = False) -> int: + upgrade: bool = False, + extra_args: Optional[List[str]] = None, + interpreter: Optional[Path] = None) -> int: """pip install PACKAGES into ENV (an env dir, default the app env), creating it first if needed. @@ -207,12 +224,15 @@ def pip_install(packages: List[str], *, emit=None, cancel=None, alongside each other or the app requirements. With UPGRADE the install runs with ``-U``: pip then resolves the latest version itself and reports "Requirement already satisfied" when the env already holds it — - the backend update action's cheap freshness check. Returns pip's exit - code. With EMIT given (the in-TUI task view) pip runs with - ``--progress-bar off`` so its output is clean status lines rather than - carriage-return progress spam. + the backend update action's cheap freshness check. EXTRA_ARGS are passed + through to pip verbatim (e.g. ``--pre`` for prerelease-resolution stacks + like sglang-omni, ``--no-deps`` for its model-companion packages); with + INTERPRETER a missing ENV is created from that Python instead of the + launching one. Returns pip's exit code. With EMIT given (the in-TUI task + view) pip runs with ``--progress-bar off`` so its output is clean status + lines rather than carriage-return progress spam. """ - if not env_exists(env_dir) and create_env(env_dir) != 0: + if not env_exists(env_dir) and create_env(env_dir, interpreter) != 0: return 1 target = env_dir if env_dir is not None else ENV_DIR print(f"[INFO] pip install {' '.join(packages)} into {target}...") @@ -222,6 +242,7 @@ def pip_install(packages: List[str], *, emit=None, cancel=None, if emit is not None: argv.append("--progress-bar") argv.append("off") + argv.extend(extra_args or []) argv.extend(packages) return common.run_console_subprocess(argv, emit=emit, cancel=cancel) @@ -265,6 +286,114 @@ def module_available(module: str, env_dir: Optional[Path] = None) -> bool: return result.returncode == 0 +# -- Interpreter selection (backends with Python-version requirements) -------- +# +# sglang-omni requires Python >=3.10,<3.13 while the app itself is +# version-agnostic: a user launching ``python audiobook.py`` with 3.13+ +# must still be able to install the sglang-omni backend. These helpers +# resolve a usable interpreter for such a backend's venv: the app env's +# own interpreter when it fits, else a versioned python on PATH, else a +# uv-managed standalone CPython downloaded on demand (uv itself is +# pip-installed into the app env — no system package or root required). + +# Interpreter path names probed on PATH, newest acceptable version first. +_VERSIONED_PYTHON_NAMES = ("python3.12", "python3.11", "python3.10") + + +def python_version(python: Path) -> Optional[Tuple[int, int]]: + """The (major, minor) version of the interpreter at PYTHON, or None. + + A subprocess probe (sys.version_info) rather than the filename: a + ``python3.12`` on PATH may be anything, and the launching interpreter's + version is not parseable from its path at all. + """ + try: + proc = subprocess.run( + [str(python), "-c", + "import sys; print(sys.version_info.major, sys.version_info.minor)"], + capture_output=True, text=True, timeout=30, check=False) + except (OSError, subprocess.TimeoutExpired): + return None + try: + major, minor = proc.stdout.split() + return int(major), int(minor) + except ValueError: + return None + + +def compatible_interpreter(versions=((3, 12), (3, 11), (3, 10)), + env_dir: Optional[Path] = None, + extra_names: Tuple[str, ...] = ()) -> Optional[Path]: + """An existing interpreter whose version is in VERSIONS, or None. + + Checked in order: the managed env's own interpreter (ENV_DIR's — it + inherits the launching interpreter's version, and exists whenever the + app runs), then ``python3.X`` names for each VERSION on PATH (newest + first), then any EXTRA_NAMES (a backend may probe e.g. ``python3``). + The app interpreter first keeps the common case (a 3.10-3.12 host) + download-free; the PATH scan covers hosts that keep several Pythons + side by side (Arch's python312 AUR builds, Ubuntu deadsnakes, Homebrew + versioned formulae). + """ + names = [f"python{major}.{minor}" for major, minor in versions] + candidates: List[Path] = [] + if env_exists(env_dir): + candidates.append(env_python(env_dir)) + for name in [*names, *extra_names]: + found = shutil.which(name) + if found: + candidates.append(Path(found)) + for candidate in candidates: + version = python_version(candidate) + if version is not None and tuple(version) in {tuple(v) for v in versions}: + return candidate + return None + + +def uv_script(env_dir: Optional[Path] = None) -> Path: + """Absolute path to the uv console script installed in ENV.""" + return env_script("uv", env_dir) + + +def ensure_uv(emit=None, cancel=None) -> int: + """Make uv available inside the app env (pip install it when missing). + + uv is only needed to provision a Python for version-pinned backend + venvs (sglang-omni), so it is installed lazily into the *app* env — + never into backend envs — and its wheel exists for every platform the + backends run on (manylinux, musllinux, macOS). Returns pip's exit + code; pip's own freshness check makes a repeat call a fast no-op. + """ + if uv_script().is_file(): + return 0 + return pip_install(["uv"], emit=emit, cancel=cancel) + + +def provision_env_with_uv(env_dir: Path, python_spec: str = "3.12", + *, emit=None, cancel=None) -> int: + """Create ENV from a uv-managed PYTHON_SPEC (e.g. "3.12"). + + uv discovers a matching interpreter (system first) and downloads a + checksummed standalone CPython into PYTHON_INSTALL_DIR when none + exists — the no-prerequisites path for hosts whose only Python is + outside the backend's required range. ``--seed`` installs pip into + the new venv so the regular ``python -m pip`` helpers keep working. + UV_PYTHON_INSTALL_DIR keeps the downloaded interpreters inside the + project (survives uv cache cleans, uninstallable with the backend). + Returns the uv exit code. + """ + uv = uv_script() + if not uv.is_file(): + print("[ERROR] uv is not installed in the app environment") + return 1 + env = dict(os.environ) + env["UV_PYTHON_INSTALL_DIR"] = str(PYTHON_INSTALL_DIR) + print(f"[INFO] creating {env_dir} with uv-managed Python {python_spec} " + f"(downloaded into {PYTHON_INSTALL_DIR} when needed)...") + argv = [str(uv), "venv", "--seed", "--python", python_spec, str(env_dir)] + return common.run_console_subprocess(argv, env=env, emit=emit, cancel=cancel) + + # Import names that differ from their requirements.txt distribution name. _IMPORT_NAMES = { "beautifulsoup4": "bs4", diff --git a/app/backends/managed.py b/app/backends/managed.py index 3b98aea..d7c8164 100644 --- a/app/backends/managed.py +++ b/app/backends/managed.py @@ -29,7 +29,7 @@ from dataclasses import dataclass from typing import Optional from backends import ServerSpec, common, probe, servers -from converter.clients import BACKEND_QWEN +from converter.clients import BACKEND_QWEN, BACKEND_SGLOMNI @dataclass @@ -56,20 +56,22 @@ class ManagedServer: servers.stop(self.spec.name) -def ensure_running(backend: str, voice_mode: str) -> Optional[ManagedServer]: +def ensure_running(backend: str, voice_mode: str, + model: Optional[str] = None) -> Optional[ManagedServer]: """Make the backend's managed server ready for a conversion run. Resolves the server spec for BACKEND (qwen: the demo hosting the - model VOICE_MODE needs; the others: their single configured spec), - then starts it when its port is free — waiting out the boot and - streaming ``servers``' console progress — or reuses the server - already answering there (qwen: restarting a managed server that hosts - another model, refusing a foreign one). Returns the run's - ``ManagedServer`` (call ``shutdown`` when the conversion is over), or - None when the backend is not installed here and nothing can be - started: the caller proceeds unmanaged, since a foreign server at the - configured endpoint may still answer and otherwise the conversion - fails with the converter's own unreachable-server message. + model VOICE_MODE needs; sglomni: the server hosting MODEL — one + model per process; the others: their single configured spec), then + starts it when its port is free — waiting out the boot and streaming + ``servers``' console progress — or reuses the server already answering + there (qwen/sglomni: restarting a managed server that hosts another + model, refusing a foreign one). Returns the run's ``ManagedServer`` + (call ``shutdown`` when the conversion is over), or None when the + backend is not installed here and nothing can be started: the caller + proceeds unmanaged, since a foreign server at the configured endpoint + may still answer and otherwise the conversion fails with the + converter's own unreachable-server message. Raises KeyboardInterrupt when the boot poll is interrupted (after stopping a server this call spawned, so nothing is left loading). @@ -83,12 +85,20 @@ def ensure_running(backend: str, voice_mode: str) -> Optional[ManagedServer]: "automatically; the conversion will use the configured " "endpoint (see the TUI's Configure Backends to install it).") return None - spec = _spec_for(status, backend, voice_mode) - return _boot(spec, backend, voice_mode) + spec = _spec_for(status, backend, voice_mode, model) + return _boot(spec, backend, voice_mode, model) -def _spec_for(status, backend: str, voice_mode: str) -> ServerSpec: +def _spec_for(status, backend: str, voice_mode: str, + model: Optional[str] = None) -> ServerSpec: """The server spec this run needs, from STATUS's detected servers.""" + if backend == BACKEND_SGLOMNI: + # sglomni hosts one model per process: aim the spec at the model + # this run selected (the converter resolves it again; resolving + # here too keeps the boot check and the conversion consistent). + from backends.sglomni import models as sg_models + from backends.sglomni import status as sg_status + return sg_status.build_spec(sg_models.resolve_model(model)) if backend != BACKEND_QWEN: return status.servers[0] # qwen hosts one model per process: aim the spec at the model this @@ -97,18 +107,50 @@ def _spec_for(status, backend: str, voice_mode: str) -> ServerSpec: return qwen.build_spec(qwen.model_for_voice_mode(voice_mode)) -def _boot(spec: ServerSpec, backend: str, voice_mode: str) -> ManagedServer: +def _boot(spec: ServerSpec, backend: str, voice_mode: str, + model: Optional[str] = None) -> ManagedServer: """Start or reuse the server SPEC describes, per the run's needs.""" wanted_model = None + wanted_repo = None if backend == BACKEND_QWEN: from backends import qwen wanted_model = qwen.model_for_voice_mode(voice_mode) + if backend == BACKEND_SGLOMNI: + from backends.sglomni import models as sg_models + entry = sg_models.resolve_model(model) + wanted_repo = entry.repo + # When this GPU cannot run the model's default FP8 pipeline, the + # spec launches the vendored bf16 config — say so before the boot. + from backends.sglomni import status as sg_status + note = sg_status.gpu_fallback_note(entry) + if note: + print(f"[WARNING] {note}") if common.server_running(spec.url): - if wanted_model is None: + if wanted_model is None and wanted_repo is None: print(f"[INFO] using the {spec.name} server already running " f"at {spec.url}") return ManagedServer(spec) + if wanted_repo is not None: + running_repo = probe.sglomni_served_model(spec.url) + if running_repo == wanted_repo: + print(f"[INFO] using the {spec.name} server already " + f"running at {spec.url} (hosting {wanted_repo})") + return ManagedServer(spec) + hosted = running_repo or "an unknown model" + if not servers.alive(spec.name): + print(f"[ERROR] a server this tool did not start is " + f"running at {spec.url} hosting {hosted} — this run " + f"needs {wanted_repo}. Stop that server first, or " + "convert with it by picking that model.") + return ManagedServer(spec, ok=False) + # Ours: stop it and boot the newly-selected model on the same + # port (the TUI's Generate form restarts a managed server the + # same way when the run's model selection changes). + print(f"[INFO] restarting the {spec.name} server to host " + f"{wanted_repo}...") + servers.stop(spec.name) + return _start(spec) running_model = qwen.model_for_identity( probe.identify_server(spec.url)) if running_model == wanted_model: diff --git a/app/backends/probe.py b/app/backends/probe.py index 818aadb..d54a9f3 100644 --- a/app/backends/probe.py +++ b/app/backends/probe.py @@ -26,7 +26,7 @@ it stays cheap to import alongside ``backends.common``. import json import urllib.parse import urllib.request -from typing import Optional +from typing import List, Optional from backends import common @@ -35,6 +35,7 @@ IDENTITY_FASTER = "faster" IDENTITY_QWEN_CUSTOM = "qwen-custom" IDENTITY_QWEN_CLONE = "qwen-clone" IDENTITY_QWEN_DESIGN = "qwen-design" +IDENTITY_SGLOMNI = "sglomni" # Endpoint names the converter resolves for each qwen demo server (see # converter.clients QwenTTSClient). Mirror them here so identification matches @@ -77,7 +78,7 @@ def _get_json(url: str, timeout: float) -> Optional[dict]: def _identify_health(base: str, timeout: float) -> Optional[str]: - """Identify audio.cpp / faster from their ``/health`` responses.""" + """Identify audio.cpp / faster / sglang-omni from their /health responses.""" payload = _get_json(f"{base}/health", timeout) if payload is None: return None @@ -93,6 +94,15 @@ def _identify_health(base: str, timeout: float) -> Optional[str]: if isinstance(entries, list) and entries \ and any(isinstance(e, dict) and e.get("id") for e in entries): return IDENTITY_AUDIOCPP + # sglang-omni's /health reports {"status": "healthy", "stages": [...]} + # (200 when serving, 503 with "unhealthy" while booting). A 503 body + # still parses as JSON here, so require the healthy word explicitly — + # an "unhealthy" sgl-omni must not count as usable. The pipeline + # "stages" list is confirmed as a secondary mark (present on every + # sgl-omni 0.1.x server) before trusting the generic-sounding status. + if payload.get("status") == "healthy" \ + and isinstance(payload.get("stages"), list): + return IDENTITY_SGLOMNI return None @@ -184,3 +194,48 @@ def faster_model_loaded(url: str, timeout: float = DEFAULT_TIMEOUT) -> bool: """ payload = health_payload(url, timeout) return bool(payload and payload.get("model_loaded")) + + +def sglomni_served_model(url: str, + timeout: float = DEFAULT_TIMEOUT) -> Optional[str]: + """The HuggingFace repo id a sglang-omni server at URL hosts, or None. + + ``GET /v1/models`` answers ``{"data": [{"id": <served repo>}, ...]}`` + with exactly one entry (one model per server process) — the same + model-identity role qwen's probe plays for its three demos. Used to + name the running model in statuses and to decide when a managed + server must be restarted to host the model a run selected. + """ + if not url: + return None + models = _get_json(f"{url.rstrip('/')}/v1/models", timeout) + if models is None: + return None + entries = models.get("data") + if isinstance(entries, list) and entries \ + and isinstance(entries[0], dict): + model_id = entries[0].get("id") + if isinstance(model_id, str) and model_id: + return model_id + return None + + +def sglomni_voice_names(url: str, + timeout: float = DEFAULT_TIMEOUT) -> Optional[List[str]]: + """The uploaded voice names registered on a sglang-omni server, or None. + + ``GET /v1/audio/voices?names_only=true`` answers + ``{"uploaded_voice_names": [...]}`` — the server-side voices a remote + Convert form can offer in its voice picker (uploaded clips persist + across server restarts). None when the URL does not answer. + """ + if not url: + return None + payload = _get_json(f"{url.rstrip('/')}/v1/audio/voices?names_only=true", + timeout) + if payload is None: + return None + names = payload.get("uploaded_voice_names") + if isinstance(names, list): + return [str(name) for name in names if isinstance(name, str) and name] + return None diff --git a/app/backends/servers.py b/app/backends/servers.py index 25ab472..7815f30 100644 --- a/app/backends/servers.py +++ b/app/backends/servers.py @@ -45,12 +45,23 @@ STOP_GRACE_SECONDS = 10 # How often the start poll re-checks readiness (seconds). POLL_INTERVAL = 1 +# Known crash signatures in a failed boot's log tail, each with a +# plain-language hint appended to the failure report (the raw tail alone +# is often a wall of framework traceback). +_BOOT_HINTS = ( + # sglang fused-MoE fp8e4nv kernel on pre-sm_89 GPUs (e.g. an FP8 + # checkpoint or a default FP8 pipeline on Ampere). + ("fp8e4nv not supported", + "the server crashed compiling an FP8 MoE kernel: FP8 needs compute " + "capability 8.9+ (RTX 4090/5090, Hopper) and cannot run on this GPU"), +) + # Progress callback: called with an event dict. KIND is one of: # "starting" {name, argv, cwd, log_path, pid} spawned, waiting for boot # "elapsed" {name, seconds} heartbeat while waiting # "ready" {name, url} server is up and answering -# "exited" {name, returncode, log_tail} process exited while booting -# "timeout" {name, seconds, log_tail} readiness deadline elapsed +# "exited" {name, returncode, log_tail, hint} process died while booting +# "timeout" {name, seconds, log_tail, hint} readiness deadline elapsed # "running" {name, url} already up (no spawn) # "cancelled" {name} boot aborted via cancel # "error" {message} could not spawn the executable @@ -77,10 +88,14 @@ def _console_progress(event: dict) -> None: print(f"[ERROR] {event['name']} server exited with code " f"{event['returncode']}") _print_tail(event.get("log_tail")) + if event.get("hint"): + print(f"[WARNING] hint: {event['hint']}") elif kind == "timeout": print(f"[ERROR] {event['name']} server did not start within " f"{int(event['seconds'])}s") _print_tail(event.get("log_tail")) + if event.get("hint"): + print(f"[WARNING] hint: {event['hint']}") elif kind == "error": print(f"[ERROR] {event['message']}") @@ -114,6 +129,15 @@ def _read_log_tail(name: str, lines: int = 20) -> List[str]: return text.splitlines()[-lines:] +def _boot_hint(log_tail: List[str]) -> Optional[str]: + """A plain-language hint for a known crash signature in LOG_TAIL.""" + text = "\n".join(log_tail) + for signature, hint in _BOOT_HINTS: + if signature in text: + return hint + return None + + def _print_tail(tail: List[str]) -> None: """Print a log-tail event payload (used by the console callback).""" if tail: @@ -282,14 +306,16 @@ def start(spec, progress: ProgressCallback = None, spec's CWD when it has one (audio.cpp discovers model_specs/ from its process working directory), records the pid, and polls readiness — ``_server_ready``, so an IDENTITY spec must actually answer HTTP — until - it is up or ``SERVER_START_TIMEOUT`` elapses. Returns True when the - server is up; on timeout or early exit reports the log tail and returns - False. A no-op (True) when the server is already running. + it is up or the spec's start timeout elapses (``ServerSpec.start_timeout`` + overrides SERVER_START_TIMEOUT; the sglang-omni pipeline needs the + longer budget). Returns True when the server is up; on timeout or early + exit reports the log tail and returns False. A no-op (True) when the + server is already running. PROGRESS, when given, receives each boot event (see ProgressCallback); - the default ``_console_progress`` prints them, preserving the old - console output. CANCEL (a threading.Event) aborts the boot: the spawned - process is terminated and False is reported (event kind "cancelled"). + the default ``_console_progress`` prints them, preserving the old console + output. CANCEL (a threading.Event) aborts the boot: the spawned process + is terminated and False is reported (event kind "cancelled"). """ report = progress if progress is not None else _console_progress argv: List[str] = list(spec.argv) @@ -303,6 +329,9 @@ def start(spec, progress: ProgressCallback = None, report({"kind": "running", "name": spec.name, "url": spec.url}) return True + start_timeout = getattr(spec, "start_timeout", None) \ + or SERVER_START_TIMEOUT + LOG_DIR.mkdir(parents=True, exist_ok=True) # Refuse to double-start: a live pid file means a previous start is # still booting (or its process is wedged). Spawning a second server @@ -381,7 +410,7 @@ def start(spec, progress: ProgressCallback = None, started = time.time() next_heartbeat = started + 15 - deadline = started + SERVER_START_TIMEOUT + deadline = started + start_timeout while time.time() < deadline: if cancel is not None and cancel.is_set(): # User cancelled while booting: kill what we spawned (the @@ -394,9 +423,11 @@ def start(spec, progress: ProgressCallback = None, report({"kind": "cancelled", "name": spec.name}) return False if proc.poll() is not None: + tail = _read_log_tail(spec.name) report({"kind": "exited", "name": spec.name, "returncode": proc.returncode, - "log_tail": _read_log_tail(spec.name)}) + "log_tail": tail, + "hint": _boot_hint(tail)}) try: pid_file.unlink() except OSError: @@ -410,9 +441,11 @@ def start(spec, progress: ProgressCallback = None, "seconds": time.time() - started}) next_heartbeat += 15 time.sleep(POLL_INTERVAL) + tail = _read_log_tail(spec.name) report({"kind": "timeout", "name": spec.name, - "seconds": SERVER_START_TIMEOUT, - "log_tail": _read_log_tail(spec.name)}) + "seconds": start_timeout, + "log_tail": tail, + "hint": _boot_hint(tail)}) # Leave the pid file in place so stop() can kill it (it may still load). return False diff --git a/app/backends/sglomni/__init__.py b/app/backends/sglomni/__init__.py new file mode 100644 index 0000000..688e488 --- /dev/null +++ b/app/backends/sglomni/__init__.py @@ -0,0 +1,98 @@ +"""sglang-omni backend — setup wizard, server specs, model management. + +A package of focused modules behind one import surface: everything public +is re-exported here, so callers (the hub, the CLI, tests) keep using +``backends.sglomni.<name>`` regardless of which module implements it. + +Modules: + constants shared constants (package name, port, timeouts, configs) + catalog the model catalog + the install tree shape + gpu the NVIDIA GPU facts (nvidia-smi) launch decisions use + pythonenv interpreter selection/provisioning for the backend venv + models model install state, (un)install, run-model resolution + status detect() for the hub's backend menu + the server spec + wizard the TUI wizard, configure screen, uninstall/update, CLI +""" + +from .constants import ( + CONFIGS_DIR, + DEFAULT_PORT, + PYTHON_SPEC, + PYTHON_VERSIONS, + SERVER_NAME, + SERVER_START_TIMEOUT, + SGLOMNI_PIP_PKG, + UV_PIP_PKG, +) +from .gpu import ( + compute_capability, + describe, +) +from .catalog import ( + CAPABILITY_CLONE, + CAPABILITY_DESIGN, + CAPABILITY_SPEAKER, + ENTRIES, + ModelEntry, + config_path, + entries_by_keys, + entry_by_key, + entry_by_repo, + install_tree_families, +) +from .pythonenv import ( + env_compatible, + env_version, + prepare_env, +) +from .models import ( + delete_model_weights, + install_model, + installed_entries, + installed_keys, + model_installed, + preset_voices, + repo_dir, + resolve_model, + uninstall_model, +) +from .status import ( + build_spec, + detect, + gpu_fallback_note, + launch_config_path, + needs_fp8_fallback, +) +from .wizard import ( + build_parser, + main, + models_screen, + run_tui, + setup_screen, + uninstall, + update, +) + +__all__ = [ + # constants + "CONFIGS_DIR", "DEFAULT_PORT", "PYTHON_SPEC", "PYTHON_VERSIONS", + "SERVER_NAME", "SERVER_START_TIMEOUT", "SGLOMNI_PIP_PKG", "UV_PIP_PKG", + # gpu + "compute_capability", "describe", + # catalog + "CAPABILITY_CLONE", "CAPABILITY_DESIGN", "CAPABILITY_SPEAKER", + "ENTRIES", "ModelEntry", "config_path", "entries_by_keys", + "entry_by_key", "entry_by_repo", "install_tree_families", + # pythonenv + "env_compatible", "env_version", "prepare_env", + # models + "delete_model_weights", "install_model", "installed_entries", + "installed_keys", "model_installed", "preset_voices", "repo_dir", + "resolve_model", "uninstall_model", + # status + "build_spec", "detect", "gpu_fallback_note", "launch_config_path", + "needs_fp8_fallback", + # wizard + "build_parser", "main", "models_screen", "run_tui", "setup_screen", + "uninstall", "update", +] diff --git a/app/backends/sglomni/__main__.py b/app/backends/sglomni/__main__.py new file mode 100644 index 0000000..faba61e --- /dev/null +++ b/app/backends/sglomni/__main__.py @@ -0,0 +1,8 @@ +"""Direct CLI execution: ``python -m backends.sglomni`` (from app/).""" + +import sys + +from backends.sglomni import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/app/backends/sglomni/catalog.py b/app/backends/sglomni/catalog.py new file mode 100644 index 0000000..c89e865 --- /dev/null +++ b/app/backends/sglomni/catalog.py @@ -0,0 +1,277 @@ +"""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. + +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 + + +# 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)) + +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", + config=None, + capability=CAPABILITY_CLONE, + requires_reference=False, + 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", + ), +) + +_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 (Higgs, 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 diff --git a/app/backends/sglomni/configs/dots_tts.yaml b/app/backends/sglomni/configs/dots_tts.yaml new file mode 100644 index 0000000..d2c42f4 --- /dev/null +++ b/app/backends/sglomni/configs/dots_tts.yaml @@ -0,0 +1,36 @@ +config_cls: DotsTTSPipelineConfig +model_path: dots-studio/dots.tts-mf + +# The solver settings feed two consumers: the latent AR engine solves with +# them, and preprocessing sizes its generation schedule from them. The +# selector writes one value into both stages; an explicit per-stage entry +# under stages: would override it. +shared: + - select: + stages: [preprocessing, latent_engine] + factory: + num_steps: 4 + max_generate_length: 500 + +stages: + reference_encode: + factory: + max_concurrency: 8 + max_batch_size: 1 + max_batch_wait_ms: 4 + latent_engine: + factory: + optimize: true + engine: + mem_fraction_static: 0.20 + max_running_requests: 16 + # note (luojiaxuan): backbone decode runs through the SGLang CUDA graph + # with the model-owned feedback buffer; delete these two lines (or set + # disable_cuda_graph: true) to fall back to eager backbone decode. + disable_cuda_graph: false + cuda_graph_max_bs: 16 + vocoder: + factory: + optimize: true + max_batch_size: 4 + max_batch_wait_ms: 2 diff --git a/app/backends/sglomni/configs/moss_tts.yaml b/app/backends/sglomni/configs/moss_tts.yaml new file mode 100644 index 0000000..ef8d8fa --- /dev/null +++ b/app/backends/sglomni/configs/moss_tts.yaml @@ -0,0 +1,2 @@ +config_cls: MossTTSPipelineConfig +model_path: OpenMOSS-Team/MOSS-TTS-v1.5 diff --git a/app/backends/sglomni/configs/moss_tts_local.yaml b/app/backends/sglomni/configs/moss_tts_local.yaml new file mode 100644 index 0000000..0b37f4e --- /dev/null +++ b/app/backends/sglomni/configs/moss_tts_local.yaml @@ -0,0 +1,2 @@ +config_cls: MossTTSLocalPipelineConfig +model_path: OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 diff --git a/app/backends/sglomni/configs/qwen3_tts_0_6b.yaml b/app/backends/sglomni/configs/qwen3_tts_0_6b.yaml new file mode 100644 index 0000000..a712ef9 --- /dev/null +++ b/app/backends/sglomni/configs/qwen3_tts_0_6b.yaml @@ -0,0 +1,2 @@ +config_cls: Qwen3TTSPipelineConfig +model_path: Qwen/Qwen3-TTS-12Hz-0.6B-Base diff --git a/app/backends/sglomni/configs/qwen3_tts_0_6b_customvoice.yaml b/app/backends/sglomni/configs/qwen3_tts_0_6b_customvoice.yaml new file mode 100644 index 0000000..6b284da --- /dev/null +++ b/app/backends/sglomni/configs/qwen3_tts_0_6b_customvoice.yaml @@ -0,0 +1,2 @@ +config_cls: Qwen3TTSPipelineConfig +model_path: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice diff --git a/app/backends/sglomni/configs/qwen3_tts_1_7b.yaml b/app/backends/sglomni/configs/qwen3_tts_1_7b.yaml new file mode 100644 index 0000000..4f7706d --- /dev/null +++ b/app/backends/sglomni/configs/qwen3_tts_1_7b.yaml @@ -0,0 +1,2 @@ +config_cls: Qwen3TTSPipelineConfig +model_path: Qwen/Qwen3-TTS-12Hz-1.7B-Base diff --git a/app/backends/sglomni/configs/qwen3_tts_1_7b_voicedesign.yaml b/app/backends/sglomni/configs/qwen3_tts_1_7b_voicedesign.yaml new file mode 100644 index 0000000..20ae0b8 --- /dev/null +++ b/app/backends/sglomni/configs/qwen3_tts_1_7b_voicedesign.yaml @@ -0,0 +1,2 @@ +config_cls: Qwen3TTSPipelineConfig +model_path: Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign diff --git a/app/backends/sglomni/configs/s2pro_tts.yaml b/app/backends/sglomni/configs/s2pro_tts.yaml new file mode 100644 index 0000000..0bc3eba --- /dev/null +++ b/app/backends/sglomni/configs/s2pro_tts.yaml @@ -0,0 +1,2 @@ +config_cls: S2ProPipelineConfig +model_path: fishaudio/s2-pro diff --git a/app/backends/sglomni/configs/voxtral_tts.yaml b/app/backends/sglomni/configs/voxtral_tts.yaml new file mode 100644 index 0000000..450cbce --- /dev/null +++ b/app/backends/sglomni/configs/voxtral_tts.yaml @@ -0,0 +1,2 @@ +config_cls: VoxtralTTSPipelineConfig +model_path: mistralai/Voxtral-4B-TTS-2603 diff --git a/app/backends/sglomni/configs/zonos2_bf16.yaml b/app/backends/sglomni/configs/zonos2_bf16.yaml new file mode 100644 index 0000000..79d5d26 --- /dev/null +++ b/app/backends/sglomni/configs/zonos2_bf16.yaml @@ -0,0 +1,28 @@ +# ZONOS2 without the FP8-quantized MoE pipeline. +# +# sglang-omni's default ZONOS2 config (Zonos2PipelineConfig) hardcodes +# `fp8: True` in the tts_engine stage factory: the MoE experts are +# dynamically quantized bf16 -> fp8 at load time, and sglang's fused-MoE +# Triton kernel for fp8e4nv only compiles on compute capability 8.9+ +# (RTX 4090/5090, Hopper). On older GPUs the server dies mid-boot with +# `type fp8e4nv not supported in this architecture`. +# +# This copy turns FP8 off (`factory.fp8: false` overrides the hardcoded +# kwarg — free-form factory keys pass through to the stage factory), so +# the model runs in bf16 and works on e.g. Ampere (RTX 30xx, A100) at +# roughly twice the MoE VRAM. The managed spec selects it automatically +# on GPUs the FP8 path cannot run (backends.sglomni.status.build_spec). +# +# The bf16 weights (~11.5 GB) also need a bigger static-pool budget than +# the builder's default 0.5: on a 24 GB card 0.5 leaves no room for the +# KV cache inside 12 GB, and the server aborts with "Loaded weights +# leave no GPU memory for the KV cache" (the profiler asks for >=0.64). +# 0.70 gives ~16.8 GB static -> ~1.4 GB KV, and leaves ~7 GB of the card +# for the colocated speaker-encode/vocoder stages. +config_cls: Zonos2PipelineConfig +model_path: Zyphra/zonos2 +stages: + tts_engine: + factory: + fp8: false + mem_fraction_static: 0.70 diff --git a/app/backends/sglomni/constants.py b/app/backends/sglomni/constants.py new file mode 100644 index 0000000..97a4f3d --- /dev/null +++ b/app/backends/sglomni/constants.py @@ -0,0 +1,37 @@ +"""Shared constants for the sglang-omni backend.""" + +from pathlib import Path + +# The pip package providing the `sgl-omni` server CLI. Installed unpinned +# (like qwen-tts and faster-qwen3-tts) so the update action can move with +# upstream releases; the stack this code was verified against is 0.1.4. +SGLOMNI_PIP_PKG = "sglang-omni" +# uv is provisioned into the app env only when a Python the sglang-omni +# stack accepts (>=3.10,<3.13) is not already available. +UV_PIP_PKG = "uv" + +# The dedicated venv is backends.envs.SGLOMNI_ENV_DIR (imported from there +# by the modules that need it) — its interpreter may differ from the +# launching Python's, see pythonenv.prepare_env. + +# The vendored copies of upstream examples/configs/*.yaml (the `sgl-omni +# serve --config` argument). Two catalog models run without one (Higgs, +# ZONOS2) and are launched from --model-path alone. +CONFIGS_DIR = Path(__file__).resolve().parent / "configs" + +# The managed server's port when the configured URL names none. +DEFAULT_PORT = 8100 + +# The ServerSpec name (pid/log files: app/logs/sglomni-server.*). +SERVER_NAME = "sglomni" + +# The managed server boots a multi-stage pipeline (preprocessing, TTS +# generation, vocoder) and may pull companion weights on first start, so +# its spec overrides the shared 600 s start timeout. +SERVER_START_TIMEOUT = 1200 + +# Python interpreters the sglang-omni stack accepts (requires-python +# ">=3.10,<3.13"), newest first — the search order for both the +# compatible-interpreter scan and the uv fallback. +PYTHON_VERSIONS = ((3, 12), (3, 11), (3, 10)) +PYTHON_SPEC = "3.12" diff --git a/app/backends/sglomni/gpu.py b/app/backends/sglomni/gpu.py new file mode 100644 index 0000000..f305098 --- /dev/null +++ b/app/backends/sglomni/gpu.py @@ -0,0 +1,62 @@ +"""NVIDIA GPU facts for launch decisions (best-effort, via nvidia-smi). + +sglang-omni model pipelines carry GPU-architecture constraints the app +must respect when it builds a server spec (e.g. ZONOS2's default pipeline +quantizes its MoE experts to FP8 — a Triton kernel that only compiles on +compute capability 8.9+). The GPU's compute capability is read from +``nvidia-smi`` rather than torch so the app venv needs no CUDA stack, and +every answer here is cached: the hardware cannot change mid-process. + +Everything is best-effort: when nvidia-smi is absent, errors out, or +reports something unparsable the callers get None and keep upstream +defaults instead of second-guessing the environment. +""" + +import shutil +import subprocess +from functools import lru_cache +from typing import Optional, Tuple + + +@lru_cache(maxsize=1) +def _query() -> Optional[Tuple[str, str]]: + """(name, "M.m" compute cap) for GPU 0, or None when unanswerable.""" + if shutil.which("nvidia-smi") is None: + return None + try: + proc = subprocess.run( + ["nvidia-smi", "--query-gpu=name,compute_cap", + "--format=csv,noheader,nounits", "-i", "0"], + capture_output=True, text=True, timeout=10) + except (OSError, subprocess.SubprocessError): + return None + first = proc.stdout.strip().splitlines()[:1] + if proc.returncode != 0 or not first: + return None + fields = [field.strip() for field in first[0].split(",")] + if len(fields) < 2: + return None + return fields[0], fields[1] + + +def compute_capability() -> Optional[Tuple[int, int]]: + """The first NVIDIA GPU's (major, minor) compute capability, or None. + + GPU 0 is what the managed sglang-omni pipeline stages bind to. An + unparsable capability string counts as unanswerable.""" + answer = _query() + if answer is None: + return None + try: + major, minor = (int(part) for part in answer[1].split(".")) + except ValueError: + return None + return major, minor + + +def describe() -> Optional[str]: + """A human-readable "NAME (compute capability M.m)" line, or None.""" + answer = _query() + if answer is None: + return None + return f"{answer[0]} (compute capability {answer[1]})" diff --git a/app/backends/sglomni/models.py b/app/backends/sglomni/models.py new file mode 100644 index 0000000..a5211f6 --- /dev/null +++ b/app/backends/sglomni/models.py @@ -0,0 +1,283 @@ +"""sglang-omni model weights: install state on disk, (un)install actions. + +Model weights are not part of the pip install: each server fetches its +HuggingFace repo into the standard hub cache on first start. This module +pre-fetches ("Install") and deletes ("Uninstall") those cache directories +per model — via the venv's hf CLI, exactly what a first server start +would do — plus the per-model companion packages the server needs +(``extras`` in the catalog, e.g. the qwen-tts --no-deps stack or the +Descript DAC codec). + +The cache helpers mirror huggingface_hub's own directory layout and +environment overrides (the same files ``from_pretrained`` writes), so an +install lands exactly where a server start would look. The cache is +shared with the other backends: a repo both host (Qwen3-TTS 1.7B Base / +VoiceDesign exist in the qwen backend too) is downloaded once and its +deletion affects both — the same convention every backend here accepts. +""" + +import os +import shutil +from pathlib import Path +from typing import List, Optional + +from backends import common, envs +from backends.sglomni.catalog import ModelEntry, entry_by_key, entry_by_repo +from backends.sglomni.constants import SERVER_NAME, SGLOMNI_PIP_PKG +from backends.sglomni.pythonenv import SGLOMNI_ENV, prepare_env + +# The cache directory HF keeps repos in (models--<org>--<name> folders). +# Resolution mirrors huggingface_hub.constants: HF_HUB_CACHE beats +# HUGGINGFACE_HUB_CACHE beats HF_HOME/hub beats ~/.cache/huggingface/hub. + + +def _hf_cache_dir() -> Path: + override = os.environ.get("HF_HUB_CACHE") or os.environ.get( + "HUGGINGFACE_HUB_CACHE") + if override: + return Path(override) + home = os.environ.get("HF_HOME") + if home: + return Path(home) / "hub" + return Path.home() / ".cache" / "huggingface" / "hub" + + +def repo_dir(repo_id: str) -> Path: + """The cache directory HF keeps REPO_ID's weights in.""" + return _hf_cache_dir() / ("models--" + repo_id.replace("/", "--")) + + +def model_repo_dir(entry: ModelEntry) -> Path: + """The cached-weights directory for a catalog entry.""" + return repo_dir(entry.repo) + + +def _tree_has_file(path: Path) -> bool: + """True when any file or symlink exists under PATH (recursively).""" + try: + for item in path.iterdir(): + # Snapshot files are symlinks into blobs/; count them even when + # temporarily broken (presence is what the loader checks). + if item.is_symlink() or item.is_file(): + return True + if item.is_dir() and _tree_has_file(item): + return True + except OSError: + return False + return False + + +def model_installed(entry: ModelEntry) -> bool: + """True when ENTRY's weights look complete in the local HF cache. + + A fetched repo has refs/main plus at least one file under snapshots/; + anything less counts as not installed. An interrupted download simply + resumes — via Install, or the next server start for that model. + """ + directory = model_repo_dir(entry) + if not (directory / "refs" / "main").is_file(): + return False + return _tree_has_file(directory / "snapshots") + + +def installed_entries() -> List[ModelEntry]: + """The catalog entries whose weights are already on disk.""" + return [entry for entry in _all_entries() if model_installed(entry)] + + +def installed_keys() -> List[str]: + """The installed entries' catalog keys, in catalog order.""" + return [entry.key for entry in installed_entries()] + + +def preset_voices(entry: ModelEntry) -> List[str]: + """The preset voice names ENTRY can speak with. + + Catalog-declared speakers first (the Qwen3-TTS CustomVoice table); + otherwise the checkpoint's own ``voice_embedding/*.pt`` presets are + read from the downloaded snapshot (how Voxtral ships its named + voices). Empty when the model declares none or is not downloaded. + """ + if entry.speakers: + return list(entry.speakers) + try: + snapshots = model_repo_dir(entry) / "snapshots" + for snapshot in sorted(snapshots.iterdir()): + voice_dir = snapshot / "voice_embedding" + if voice_dir.is_dir(): + names = sorted(item.stem for item in voice_dir.glob("*.pt") + if item.is_file()) + if names: + return names + except OSError: + pass + return [] + + +def _all_entries() -> List[ModelEntry]: + from backends.sglomni.catalog import ENTRIES + return list(ENTRIES) + + +def system_dep_missing(entry: ModelEntry) -> Optional[str]: + """Remediation text when ENTRY's system binary is absent (None = ok).""" + if entry.system_dep and not shutil.which(entry.system_dep): + return (f"{entry.system_dep} (system package) was not found — " + f"{entry.system_hint}. The weights still download, but the " + "server will fail to synthesize with this model until it " + "is installed.") + return None + + +def install_model(key: str, *, emit=None, cancel=None) -> int: + """Install a catalog model: companion packages, then its weights. + + Companion ``extras`` pip-install into the sglang-omni venv exactly as + upstream instructs (``--no-deps`` where upstream says so — the + Qwen3-TTS companions must not replace the pinned Transformers 5 + stack), and a missing system binary is a loud warning, not a stop: + the download is still useful and the remediation stays on screen. + The weights pre-download via the venv's hf CLI (resumable, streamed, + cancelable). Returns the exit code. + """ + entry = entry_by_key(key) + if entry is None: + print(f"[ERROR] Unknown sglang-omni model: {key!r}") + return 1 + rc = prepare_env(emit=emit, cancel=cancel) + if rc != 0: + return rc + warning = system_dep_missing(entry) + if warning: + print(f"[WARNING] {warning}") + # A GPU the model's default FP8 pipeline cannot run gets the bf16 + # fallback note up front (the install itself is still useful: the + # weights download either way). + from backends.sglomni import status as sg_status + note = sg_status.gpu_fallback_note(entry) + if note: + print(f"[WARNING] {note}") + for spec, no_deps in entry.extras: + args = ["--no-deps"] if no_deps else None + rc = common.pip_install([spec], emit=emit, cancel=cancel, + env_dir=SGLOMNI_ENV, extra_args=args) + if rc != 0: + print(f"[WARNING] pip install {spec} failed (exit {rc}); " + f"install it into {SGLOMNI_ENV} manually") + prefix = _hf_download_prefix() + if prefix is None: + print("[ERROR] No hf CLI found in the sglang-omni venv; pip " + f"install {SGLOMNI_PIP_PKG} first") + return 1 + print(f"[INFO] Downloading {entry.repo} into {_hf_cache_dir()}...") + rc = common.run_console_subprocess( + prefix + ["download", entry.repo], emit=emit, cancel=cancel) + if rc == 0: + print(f"[OK] {entry.label} downloaded.") + return rc + + +def uninstall_model(key: str, *, emit=None, cancel=None) -> int: + """Remove a model's cached weights (the inverse of install_model). + + A locally-managed server currently hosting the model is stopped first + (best-effort) so its weights are not deleted under a live process. + CANCEL is honored after that stop phase only. Returns the exit code. + """ + entry = entry_by_key(key) + if entry is None: + print(f"[ERROR] Unknown sglang-omni model: {key!r}") + return 1 + if _managed_running_repo() == entry.repo: + from backends import servers + servers.stop(SERVER_NAME) + if common.cancel_requested(cancel): + return 130 + delete_model_weights([entry]) + return 0 + + +def delete_model_weights(entries: Optional[List[ModelEntry]] = None) -> int: + """Delete the cached HF weight dirs of ENTRIES (every model by default). + + Best-effort rmtree of each ``models--<org>--<name>`` directory; returns + how many were present and removed. Only those directories are ever + touched — the rest of the HF cache may be shared with unrelated tools. + """ + if entries is None: + entries = _all_entries() + removed = 0 + for entry in entries: + directory = model_repo_dir(entry) + if not directory.is_dir(): + continue + print(f"[INFO] Removing cached {entry.repo} weights...") + shutil.rmtree(directory, ignore_errors=True) + if directory.exists(): + print(f"[WARNING] Could not fully remove {directory}") + continue + removed += 1 + if removed: + print(f"[OK] Deleted cached weights for {removed} " + f"{'model' if removed == 1 else 'models'}.") + return removed + + +def _hf_download_prefix() -> Optional[List[str]]: + """The sglang-omni venv's hf CLI argv prefix (None when absent).""" + for name in ("hf", "huggingface-cli"): + candidate = envs.env_script(name, SGLOMNI_ENV) + if candidate.is_file(): + return [str(candidate)] + return None + + +def _managed_running_repo() -> Optional[str]: + """The repo id a locally-managed, up-and-running server hosts.""" + from backends import probe, servers + from converter import config + if servers.pid_for(SERVER_NAME) is None: + return None + if not servers.alive(SERVER_NAME): + return None + return probe.sglomni_served_model(config.SGLOMNI_API_URL) + + +def resolve_model(key: Optional[str]) -> ModelEntry: + """The catalog entry a run with MODEL_KEY uses. + + An explicit KEY must exist in the catalog and be installed (a hosted + model without weights cannot boot). Without KEY the single installed + model is auto-selected; several installed models need an explicit pick + (the CLI --model flag or the Generate form's Model menu). Raises + RuntimeError with an actionable message otherwise. + """ + if key is not None: + entry = entry_by_key(key) + if entry is None: + known = ", ".join(e.key for e in _all_entries()) + raise RuntimeError( + f"Unknown sglang-omni model {key!r} (installed models are " + f"picked by catalog key; known keys: {known})") + if not model_installed(entry): + raise RuntimeError( + f"{entry.label} is not downloaded — install it via " + "Configure Backends → SGLang-Omni, or pick an installed " + "model.") + return entry + installed = installed_entries() + if not installed: + raise RuntimeError( + "No sglang-omni models are downloaded — install one via " + "Configure Backends → SGLang-Omni (Configure).") + if len(installed) > 1: + names = ", ".join(e.key for e in installed) + raise RuntimeError( + "Several sglang-omni models are installed; pick one with " + f"--model KEY (installed: {names})") + return installed[0] + + +def entry_for_served_repo(repo: Optional[str]) -> Optional[ModelEntry]: + """The catalog entry a served /v1/models repo id belongs to.""" + return entry_by_repo(repo) if repo else None diff --git a/app/backends/sglomni/pythonenv.py b/app/backends/sglomni/pythonenv.py new file mode 100644 index 0000000..11bbaa6 --- /dev/null +++ b/app/backends/sglomni/pythonenv.py @@ -0,0 +1,96 @@ +"""The sglang-omni venv: interpreter selection and provisioning. + +sglang-omni requires Python >=3.10,<3.13 while the app itself is +version-agnostic — a host whose only ``python`` is 3.13+ must still be +able to install and run this backend. ``prepare_env`` resolves a usable +interpreter for the dedicated venv (``app/envs/sglomni``), in order: + +1. the venv already exists with an acceptable interpreter (no-op); +2. the app env's interpreter or a ``python3.10/3.11/3.12`` on PATH fits + (the venv is created from it — no download); +3. uv (pip-installed into the app env) provisions a managed standalone + CPython 3.12 into ``app/envs/pythons`` and creates the venv from it — + the zero-prerequisites path for hosts like stock Arch, where only the + latest Python exists and no root/package-manager knowledge is needed. + +An existing venv built with an incompatible interpreter (an older tool +version, a since-upgraded system Python) is transparently recreated: a +venv holds nothing user-owned, and every package in it is reinstalled by +the setup that follows anyway. +""" + +import shutil +import sys +from pathlib import Path +from typing import Optional, Tuple + +from backends import common, envs +from backends.sglomni.constants import PYTHON_SPEC, PYTHON_VERSIONS + +# The dedicated venv (its interpreter may differ from the launching one). +SGLOMNI_ENV = envs.SGLOMNI_ENV_DIR + + +def env_version() -> Optional[Tuple[int, int]]: + """The (major, minor) Python version of the venv's interpreter.""" + if not envs.env_exists(SGLOMNI_ENV): + return None + return envs.python_version(envs.env_python(SGLOMNI_ENV)) + + +def env_compatible() -> bool: + """True when the venv exists with a Python the sglang-omni stack accepts.""" + version = env_version() + return version is not None and tuple(version) in {tuple(v) for v in PYTHON_VERSIONS} + + +def prepare_env(*, emit=None, cancel=None) -> int: + """Make SGLOMNI_ENV exist with a Python in PYTHON_VERSIONS. Returns 0/1. + + The three resolution paths are described in the module docstring; every + path ends with a pip-equipped venv so the regular ``python -m pip`` + helpers (requirements installs, model-companion extras) keep working. + Prints what it does so the task view shows why a download happens (or, + usually, does not). + """ + if env_compatible(): + return 0 + if envs.env_exists(SGLOMNI_ENV): + version = env_version() + found = f"{version[0]}.{version[1]}" if version else "unknown" + print(f"[WARNING] {SGLOMNI_ENV} was built with Python {found}, " + "which the sglang-omni stack does not support (needs " + "3.10-3.12); recreating it with a compatible interpreter.") + # Nothing user-owned lives in the tool-managed venv, and every + # package is reinstalled by the steps that follow this one. + shutil.rmtree(SGLOMNI_ENV, ignore_errors=True) + if SGLOMNI_ENV.exists(): + print("[ERROR] Could not remove the incompatible venv; " + f"delete {SGLOMNI_ENV} manually and re-run setup.") + return 1 + + interpreter = envs.compatible_interpreter(PYTHON_VERSIONS) + if interpreter is not None: + if interpreter == envs.env_python(envs.ENV_DIR): + print(f"[INFO] using the app environment's interpreter " + f"({sys.version_info.major}.{sys.version_info.minor}) " + "for the sglang-omni venv.") + else: + print(f"[INFO] using {interpreter} for the sglang-omni venv.") + return envs.create_env(SGLOMNI_ENV, interpreter) + + # No compatible interpreter anywhere: provision one with uv. + print("[INFO] no Python 3.10-3.12 found on this system — provisioning " + f"a managed CPython {PYTHON_SPEC} with uv (no root required).") + rc = envs.ensure_uv(emit=emit, cancel=cancel) + if rc != 0: + print("[WARNING] pip install uv failed (exit " + f"{rc}); install uv manually and re-run setup") + return rc + rc = envs.provision_env_with_uv(SGLOMNI_ENV, PYTHON_SPEC, + emit=emit, cancel=cancel) + if rc != 0: + print(f"[WARNING] uv venv failed (exit {rc}); create the venv " + f"manually, e.g.: uv venv --seed --python {PYTHON_SPEC} " + f"{SGLOMNI_ENV}") + return rc diff --git a/app/backends/sglomni/status.py b/app/backends/sglomni/status.py new file mode 100644 index 0000000..099817e --- /dev/null +++ b/app/backends/sglomni/status.py @@ -0,0 +1,207 @@ +"""detect() for the hub's backend menu: how far sglang-omni is set up.""" + +from pathlib import Path +from typing import List, Optional + +from backends import BackendStatus, ServerSpec, envs, format_launch_hint, \ + probe, servers +from backends.sglomni import gpu +from backends.sglomni.catalog import ModelEntry, entry_by_repo, \ + fallback_config_path, config_path +from backends.sglomni.constants import DEFAULT_PORT, SERVER_NAME, \ + SERVER_START_TIMEOUT, SGLOMNI_PIP_PKG +from backends.sglomni.models import installed_entries +from backends.sglomni.pythonenv import SGLOMNI_ENV, env_version +from converter import config + + +def _is_installed() -> bool: + if envs.env_script("sgl-omni", SGLOMNI_ENV).is_file(): + return True + return envs.module_available("sglang_omni", SGLOMNI_ENV) + + +def build_spec(entry: ModelEntry) -> ServerSpec: + """The managed ServerSpec hosting ENTRY on the configured port. + + Public because the run preparation (hub) and the CLI's managed-server + bootstrap need to boot exactly the model their run selected, which can + differ from detect()'s default. The config yaml is the vendored copy + from the catalog; models without one (Higgs, ZONOS2) run from + --model-path alone — except that a GPU the FP8 kernels cannot run + launches ZONOS2's vendored bf16 config instead (launch_config_path). + """ + url = config.SGLOMNI_API_URL + argv: List[str] = [ + str(envs.env_script("sgl-omni", SGLOMNI_ENV)), + "serve", "--model-path", entry.repo, + ] + yaml_path = launch_config_path(entry) + if yaml_path is not None: + argv += ["--config", str(yaml_path)] + argv += ["--port", str(_port())] + return ServerSpec(SERVER_NAME, url, argv, + identity=probe.IDENTITY_SGLOMNI, + start_timeout=SERVER_START_TIMEOUT) + + +def needs_fp8_fallback(entry: ModelEntry) -> bool: + """True when ENTRY's default FP8-quantized pipeline cannot run here. + + Only models with a declared fp8_min_compute_capability are candidates, + and only when an NVIDIA GPU actually answers: a GPU this tool cannot + read keeps upstream defaults rather than second-guessing the host.""" + if not (entry.fp8_moe and entry.fp8_min_compute_capability): + return False + capability = gpu.compute_capability() + if capability is None: + return False + return capability < entry.fp8_min_compute_capability + + +def launch_config_path(entry: ModelEntry) -> Optional[Path]: + """The vendored config yaml ENTRY's server should launch with (None = + --model-path alone, the upstream default). + + A model whose default pipeline quantizes its MoE experts to FP8 + launches its vendored bf16 config instead on GPUs the FP8 Triton + kernels cannot compile on (needs_fp8_fallback) — the server then runs + the model in bf16 at about twice the MoE VRAM. A missing fallback + file degrades to the model's normal config (the server reports the + FP8 failure itself; the boot-failure hint names it).""" + if needs_fp8_fallback(entry) and entry.bf16_config: + path = fallback_config_path(entry) + if path is not None and path.is_file(): + return path + return config_path(entry) + + +def gpu_fallback_note(entry: ModelEntry) -> Optional[str]: + """A human-readable note when ENTRY launches its bf16 fallback here. + + Printed by the install and boot flows so the (small) VRAM cost and + the reason are on the record before the server starts. None when the + model runs its default (FP8) pipeline, or when no GPU answered.""" + if not needs_fp8_fallback(entry): + return None + minimum = (".".join(str(part) + for part in entry.fp8_min_compute_capability)) + where = gpu.describe() or "unknown GPU" + path = fallback_config_path(entry) + if path is None or not path.is_file(): + return (f"{entry.label}'s default pipeline quantizes its MoE " + f"experts to FP8, which needs compute capability " + f"{minimum}+; this GPU ({where}) cannot run it, and the " + "vendored bf16 fallback config is missing — the server " + "will fail to start this model.") + return (f"{entry.label}'s default pipeline quantizes its MoE experts " + f"to FP8, which needs compute capability {minimum}+; this GPU " + f"({where}) cannot run it — launching the vendored bf16 " + "config instead (about twice the MoE VRAM).") + + +def _port() -> int: + return _port_of(config.SGLOMNI_API_URL) + + +def _port_of(url: str) -> int: + import urllib.parse + try: + return urllib.parse.urlsplit(url).port or DEFAULT_PORT + except ValueError: + return DEFAULT_PORT + + +def _managed_running_entry() -> Optional[ModelEntry]: + """The catalog model a locally-managed, running server hosts.""" + if servers.pid_for(SERVER_NAME) is None: + return None + if not servers.alive(SERVER_NAME): + return None + return entry_by_repo( + probe.sglomni_served_model(config.SGLOMNI_API_URL)) + + +def detect() -> BackendStatus: + """Detect whether sglang-omni is installed, plus the launch command. + + One managed spec exists per detection, hosting the first *installed* + catalog model (catalog order) on the single configured port — runs + needing another model boot it via their own spec (the Generate form's + Model menu / CLI --model), restarting a managed server that hosts + something else. Which model currently answers is read via the probe + (local pid alive → our URL; otherwise the remote URL) so the status + names the *running* model even when it differs from the default one. + """ + installed = _is_installed() + present = installed_entries() + configured = installed and bool(present) + details: List[str] = [] + details.append("pip: installed" if installed else + f"not installed — run setup to pip install " + f"{SGLOMNI_PIP_PKG}") + if installed: + version = env_version() + if version is not None: + details.append(f"python: {version[0]}.{version[1]}") + else: + details.append("python: unknown version") + if present: + # A model launching its bf16 fallback on this GPU says so, so the + # status line explains why its server boots with a config file. + def _model_tag(model_entry: ModelEntry) -> str: + if needs_fp8_fallback(model_entry): + return f"{model_entry.key} (bf16 fallback)" + return model_entry.key + details.append(f"models: {', '.join(_model_tag(e) for e in present)}") + else: + details.append("no models downloaded — run setup (or Configure) to " + "install one") + details.append(f"port: {_port()}") + specs: List[ServerSpec] = [build_spec(present[0])] if configured else [] + managed = servers.manages(specs) + local_models: List[str] = [] + if managed and specs and servers.alive(specs[0].name): + entry = _managed_running_entry() + if entry is not None: + local_models.append(entry.label) + remote_models, remote_urls = _detect_remote(managed) + running_models = list(dict.fromkeys(local_models + remote_models)) + return BackendStatus( + SERVER_NAME, "SGLang-Omni", + installed=installed, configured=configured, + running=managed or bool(remote_urls), + details=details, + launch_hint=format_launch_hint(specs), + servers=specs, + managed=managed, + remote=bool(remote_urls), + remote_urls=remote_urls, + remote_models=remote_models, + running_models=running_models, + partial="installed (no models)" if installed and not configured + else "") + + +def _detect_remote(managed: bool = False): + """Detect an externally-run sglang-omni server at the remote URL. + + Returns ``([model_label, ...], {spec_name: url})``. The remote URL must + answer as sglang-omni (probe identity); when it equals the local URL + and this tool started that server, it is ignored (already reported as + "[local]"). + """ + remote_models: List[str] = [] + remote_urls: dict = {} + url = (config.SGLOMNI_REMOTE_URL or "").strip() + if not url: + return remote_models, remote_urls + if managed and probe.same_endpoint(url, config.SGLOMNI_API_URL): + return remote_models, remote_urls + if probe.identify_server(url) != probe.IDENTITY_SGLOMNI: + return remote_models, remote_urls + remote_urls[SERVER_NAME] = url + entry = entry_by_repo(probe.sglomni_served_model(url)) + remote_models.append(entry.label if entry + else probe.sglomni_served_model(url) or "unknown") + return remote_models, remote_urls diff --git a/app/backends/sglomni/wizard.py b/app/backends/sglomni/wizard.py new file mode 100644 index 0000000..2083ec8 --- /dev/null +++ b/app/backends/sglomni/wizard.py @@ -0,0 +1,440 @@ +#!/usr/bin/env python3 +"""Set up the SGLang-Omni backend for the audiobook generator. + +sglang-omni is a pip package (``sglang-omni``, providing the ``sgl-omni`` +server CLI) that serves one TTS model per server process from the +OpenAI-compatible ``/v1/audio/speech`` endpoint. This module sets the +backend up end-to-end: provision a Python 3.10-3.12 venv (``app/envs/ +sglomni`` — the stack does not support 3.13+, and the interpreter is +provisioned with uv when the host has none), pip-install the package, +and install the models picked in the wizard (companion packages per the +upstream recipes + HuggingFace weight pre-download). + +It is driven by ``audiobook.py``'s hub but can also be run directly: + +Usage: + python -m backends.sglomni [--models KEY[,KEY...]] [--all] + [--skip-install] [--skip-python] + +The Configure screen (``models_screen``) manages models after the fact — +the same checkbox tree the setup wizard uses (mirroring the audio.cpp +modify flow): the installed models start checked, checking installs a +model, and unchecking one removes its cached weights after a confirm. +""" + +import argparse +import shutil +import sys +from pathlib import Path +from typing import List, Optional, Tuple + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) + +from backends import common, envs, servers, setup +from backends.sglomni import catalog as sg_catalog +from backends.sglomni import models as sg_models +from backends.sglomni.constants import SERVER_NAME, SGLOMNI_PIP_PKG, \ + UV_PIP_PKG +from backends.sglomni.pythonenv import SGLOMNI_ENV, prepare_env +from backends.sglomni.status import _is_installed +from ui import taskview, tui + +_GO_BACK = object() + + +def _nvidia_gpu_present() -> bool: + """True when an NVIDIA driver answers nvidia-smi (best effort).""" + proc = common.run_console_subprocess_quiet( + ["nvidia-smi", "-L"], timeout=10) + return proc is not None and proc.returncode == 0 + + +def _preflight() -> List[str]: + """Blocking-problem messages (empty = fine); warnings print inline.""" + problems: List[str] = [] + if sys.platform == "win32": + problems.append( + "SGLang-Omni does not support Windows: its CUDA serving stack " + "(sglang, flash-attn, NVIDIA-only wheels) has no Windows " + "builds. Use the audio.cpp, qwen or faster backend instead.") + return problems + + +def _gpu_warning() -> Optional[str]: + if _nvidia_gpu_present(): + return None + return ("No NVIDIA GPU was detected (nvidia-smi did not answer). " + "SGLang-Omni serves CUDA-only: the install will succeed but " + "the server will not start without one.") + + +def _wizard(stdscr, args: argparse.Namespace) -> Optional[dict]: + """Pick the models to install (the wizard's only question). + + A checkbox tree grouped by upstream org, pre-checked with the already + installed models (a modify flow — the same tree the Configure screen + uses): checking a model installs it, unchecking one removes its + cached weights after a confirm, and confirming an empty tree installs + the package only. The confirm's diff runs as ordered task-view steps + (venv, pip package, then removals before downloads). Esc returns None + (aborted); declining the uninstall confirm re-opens the tree. + """ + problems = _preflight() + for problem in problems: + if stdscr is not None: + tui.flash(stdscr, problem, "err") + else: + print(f"[ERROR] {problem}") + return None + + warning = _gpu_warning() + if warning is not None: + if stdscr is not None: + if not tui.confirm(stdscr, f"{warning} Install anyway?", + default=False): + return None + else: + print(f"[WARNING] {warning}") + + while True: + families = sg_catalog.install_tree_families(list(sg_catalog.ENTRIES)) + installed = set(sg_models.installed_keys()) + picked = tui.checkbox_tree( + stdscr, "Select SGLang-Omni Models to Install", families, + expand_all=True, back_value=_GO_BACK, + checked={(index, option["key"]) + for index, family in enumerate(families) + for option in family["options"] + if option["key"] in installed}, + start_on_buttons=bool(installed), allow_empty=True) + if picked is _GO_BACK: + return None + to_uninstall, to_install = _model_diff( + {key for _index, key in picked}) + if to_uninstall and not _confirm_uninstall(stdscr, to_uninstall): + continue + return { + "keys": [entry.key for entry in to_install], + "uninstall_keys": [entry.key for entry in to_uninstall], + "do_python": not args.skip_python, + "do_install": (not _is_installed()) and not args.skip_install, + } + + +def _execute_steps(settings: dict) -> List[taskview.TaskStep]: + """The ordered setup steps: Python venv, pip package, models.""" + steps: List[taskview.TaskStep] = [] + + if settings.get("do_python"): + def python(emit, cancel): + rc = prepare_env(emit=emit, cancel=cancel) + if rc != 0: + print("[WARNING] could not prepare a Python 3.10-3.12 " + "venv; see the messages above") + return rc + steps.append(taskview.TaskStep("Prepare Python 3.10-3.12", python)) + + if settings.get("do_install"): + def install(emit, cancel): + rc = common.pip_install([SGLOMNI_PIP_PKG], emit=emit, + cancel=cancel, env_dir=SGLOMNI_ENV, + extra_args=["--pre"]) + if rc != 0: + print(f"[WARNING] pip install failed (exit {rc}); install " + f"{SGLOMNI_PIP_PKG} manually") + else: + print(f"[OK] {SGLOMNI_PIP_PKG} installed") + return rc + steps.append(taskview.TaskStep(f"Install {SGLOMNI_PIP_PKG}", + install)) + + steps += _reconcile_steps( + sg_catalog.entries_by_keys(settings.get("uninstall_keys", [])), + sg_catalog.entries_by_keys(settings.get("keys", []))) + return steps + + +def _model_diff(picked_keys: set) -> Tuple[List[sg_catalog.ModelEntry], + List[sg_catalog.ModelEntry]]: + """The (to_uninstall, to_install) diff a tree selection implies. + + Both lists are in catalog order: to_uninstall holds the currently + installed models the tree left unchecked, to_install the checked ones + whose weights are not on disk yet. + """ + installed = set(sg_models.installed_keys()) + to_uninstall = [entry for entry in sg_catalog.ENTRIES + if entry.key in installed + and entry.key not in picked_keys] + to_install = [entry for entry in sg_catalog.ENTRIES + if entry.key in picked_keys + and entry.key not in installed] + return to_uninstall, to_install + + +def _confirm_uninstall(stdscr, + entries: List[sg_catalog.ModelEntry]) -> bool: + """Confirm deleting cached weights; Esc counts as a decline.""" + question = (f"Remove cached weights for {len(entries)} " + f"{'model' if len(entries) == 1 else 'models'}?") + body = ["Their cached weights are deleted — a managed server hosting", + "one of them is stopped first, and removed weights", + "re-download on the next install or server start.", ""] + body += [entry.label for entry in entries] + return tui.confirm(stdscr, question, body=body, default=False, + cancel_value=False) is True + + +def _reconcile_steps(to_uninstall: List[sg_catalog.ModelEntry], + to_install: List[sg_catalog.ModelEntry] + ) -> List[taskview.TaskStep]: + """One task-view step per model: removals first (they free disk). + + install_model and uninstall_model stream through EMIT and honor + CANCEL (the hf download and the server stop both run inside); each + closure binds its model's key. + """ + steps: List[taskview.TaskStep] = [] + for entry in to_uninstall: + steps.append(taskview.TaskStep( + f"Delete {entry.label} weights", + lambda emit, cancel, target=entry.key: sg_models.uninstall_model( + target, emit=emit, cancel=cancel))) + for entry in to_install: + steps.append(taskview.TaskStep( + f"Install {entry.label}", + lambda emit, cancel, target=entry.key: sg_models.install_model( + target, emit=emit, cancel=cancel))) + return steps + + +def _execute(settings: dict) -> int: + """Console tail: python venv, pip install, model work.""" + return taskview.run_steps_inline(_execute_steps(settings)) + + +def setup_screen(stdscr) -> int: + """Run the setup on an existing curses screen (the hub's). + + Returns 0 on completion, 1 when the user aborted (Esc in the tree, a + blocking preflight problem, or a declined GPU warning). + """ + args = build_parser().parse_args([]) + settings = _wizard(stdscr, args) + if settings is None: + return 1 + return taskview.run_steps(stdscr, "Setting up SGLang-Omni", _execute_steps(settings)) + + +def run_tui(args: Optional[argparse.Namespace] = None) -> int: + """Run the sglang-omni setup end-to-end. + + The wizard's screens need a curses session of their own (the hub runs + them on its own screen); the model work is the console tail after the + terminal is restored. + """ + if args is None: + args = build_parser().parse_args([]) + import curses + try: + settings = curses.wrapper(lambda scr: _wizard(scr, args)) + except tui.WizardCancelled: + return 1 + if settings is None: + return 1 + return _execute(settings) + + +def _collect_from_flags(args: argparse.Namespace, + parser: argparse.ArgumentParser) -> Optional[dict]: + """Build the settings dict from flags for a non-interactive run.""" + problems = _preflight() + for problem in problems: + print(f"[ERROR] {problem}") + if problems: + return None + warning = _gpu_warning() + if warning is not None: + print(f"[WARNING] {warning}") + if args.all: + keys = [entry.key for entry in sg_catalog.ENTRIES] + elif args.models: + keys = [] + for part in args.models.split(","): + key = part.strip() + if not key: + continue + if sg_catalog.entry_by_key(key) is None: + known = ", ".join(e.key for e in sg_catalog.ENTRIES) + parser.error(f"unknown model key {key!r} (known: {known})") + keys.append(key) + else: + keys = [] + print("[INFO] No --models given: installing the package only " + "(use --models KEY[,KEY...] or --all to add models, or the " + "TUI's Configure screen).") + return { + "keys": keys, + "do_python": not args.skip_python, + "do_install": (not _is_installed()) and not args.skip_install, + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Set up the SGLang-Omni backend: provision a Python " + "3.10-3.12 venv (app/envs/sglomni), pip install " + "sglang-omni, and install the selected models " + "(companion packages + HuggingFace weights).") + parser.add_argument("models_pos", nargs="*", metavar="KEY", + help="Model catalog keys to install (same as " + "--models, space separated)") + parser.add_argument("--models", type=str, default=None, metavar="KEYS", + help="Comma-separated model catalog keys to " + "install (e.g. higgs_audio_v3_tts,moss_tts)") + parser.add_argument("--all", action="store_true", + help="Install every catalog model") + parser.add_argument("--skip-install", action="store_true", + help="Do not pip install sglang-omni") + parser.add_argument("--skip-python", action="store_true", + help="Do not create/provision the venv") + return parser + + +def models_screen(stdscr) -> int: + """Per-model (un)install screen: the hub's Configure-SGLang-Omni leaf. + + The same checkbox tree the audio.cpp modify flow uses: one entry per + catalog model grouped by upstream org, with the installed models + pre-checked (Confirm accepts the tree as it stands). Checking a model + installs it — companion packages plus a pre-download of its weights + via the hf CLI — and unchecking one removes its cached weights after + a confirm (a managed server hosting the model is stopped first); the + whole diff runs as one streamed, resumable, cancelable task-view run + (removals before downloads). Install requires the pip package; with + none installed a guidance flash replaces the run. The tree re-opens + after every action, re-detecting disk state; Esc pops back to + Configure Backends. Always returns 0. + """ + while True: + families = sg_catalog.install_tree_families(list(sg_catalog.ENTRIES)) + installed = set(sg_models.installed_keys()) + picked = tui.checkbox_tree( + stdscr, "Select SGLang-Omni Models", families, + back_value=_GO_BACK, + checked={(index, option["key"]) + for index, family in enumerate(families) + for option in family["options"] + if option["key"] in installed}, + start_on_buttons=True, allow_empty=True) + if picked is _GO_BACK: + return 0 + to_uninstall, to_install = _model_diff( + {key for _index, key in picked}) + if not to_install and not to_uninstall: + continue + if to_install and not _is_installed(): + tui.flash(stdscr, "Install the SGLang-Omni backend first " + "(Configure Backends > Install Backend).", "warn") + continue + if to_uninstall and not _confirm_uninstall(stdscr, to_uninstall): + continue + steps = _reconcile_steps(to_uninstall, to_install) + rc = taskview.run_steps(stdscr, "Configure SGLang-Omni", steps, + wait_on_finish=False) + if rc == 0: + parts = [] + if to_install: + parts.append(f"{len(to_install)} installed") + if to_uninstall: + parts.append(f"{len(to_uninstall)} removed") + tui.flash(stdscr, "SGLang-Omni models updated: " + + ", ".join(parts) + ".", "ok") + else: + tui.flash(stdscr, "Could not update SGLang-Omni models.", "err") + + +def uninstall(*, emit=None, cancel=None) -> int: + """Remove the SGLang-Omni backend entirely. + + Phases: stop the managed server, pip-uninstall sglang-omni and every + catalog model's companion packages, delete every cached weight + snapshot, then remove the tool-owned venv (app/envs/sglomni — the + heavyweight CUDA stack is the install, so unlike the lighter backends + the whole environment goes) and the uv-managed interpreters under + app/envs/pythons. CANCEL is honored between phases only. Returns the + exit code (130 when cancelled before a remaining phase). + """ + if servers.pid_for(SERVER_NAME) is not None: + servers.stop(SERVER_NAME) + if common.cancel_requested(cancel): + return 130 + packages = [SGLOMNI_PIP_PKG] + for entry in sg_catalog.ENTRIES: + for spec, _no_deps in entry.extras: + name = spec.split("=")[0].split("<")[0].split(">")[0].strip() + if name and name not in packages: + packages.append(name) + if envs.env_exists(SGLOMNI_ENV): + rc = common.pip_uninstall(packages, emit=emit, env_dir=SGLOMNI_ENV) + if rc != 0: + print(f"[WARNING] pip uninstall failed (exit {rc}); the venv " + "is removed below anyway") + else: + rc = 0 + if common.cancel_requested(cancel): + return 130 + sg_models.delete_model_weights() + if common.cancel_requested(cancel): + return 130 + for directory in (SGLOMNI_ENV, envs.PYTHON_INSTALL_DIR): + if directory.is_dir(): + print(f"[INFO] Removing {directory}...") + shutil.rmtree(directory, ignore_errors=True) + if directory.exists(): + print(f"[WARNING] Could not fully remove {directory}") + else: + print(f"[OK] {directory} removed.") + return rc + + +def update(*, emit=None, cancel=None) -> int: + """Update the sglang-omni backend: pip install -U in its venv. + + A managed server that is running is stopped first (best-effort): it + imports the very package being upgraded. CANCEL is honored between + phases only. Model weights are untouched (they live in the shared + HuggingFace cache and survive package upgrades). When the venv does + not exist there is nothing to update. Returns the exit code. + """ + if servers.pid_for(SERVER_NAME) is not None: + servers.stop(SERVER_NAME) + if common.cancel_requested(cancel): + return 130 + if not envs.env_exists(SGLOMNI_ENV): + print("[INFO] SGLang-Omni is not installed; nothing to update.") + return 0 + rc = common.pip_install([SGLOMNI_PIP_PKG], emit=emit, cancel=cancel, + env_dir=SGLOMNI_ENV, upgrade=True, + extra_args=["--pre"]) + if rc != 0: + print(f"[WARNING] pip install -U failed (exit {rc}); update " + f"{SGLOMNI_PIP_PKG} manually") + else: + print(f"[OK] {SGLOMNI_PIP_PKG} is up to date (or just upgraded).") + return rc + + +def main() -> int: + parser = build_parser() + args = parser.parse_args() + if setup.interactive(): + return run_tui(args) + settings = _collect_from_flags(args, parser) + if settings is None: + return 1 + return _execute(settings) + + +if __name__ == "__main__": + sys.exit(main()) |