#!/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 [KEY ...] [--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 gpu as sg_gpu from backends.sglomni import models as sg_models from backends.sglomni.constants import SERVER_NAME, SGLOMNI_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). The same probe the launch decisions use (sglomni.gpu): nvidia-smi names GPU 0, or there is no usable answer.""" return sg_gpu.describe() is not None 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, 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 try: curses.curs_set(1) # restore the text cursor hidden by the TUI except curses.error: pass 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] else: # Positional keys and --models both feed the same list (the # positional form is the --models shorthand's space-separated # twin); duplicates collapse, unknown keys stop the run with the # known set. keys: List[str] = [] for part in list(args.models_pos or []) + \ (args.models or "").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})") if key not in keys: keys.append(key) if not keys: print("[INFO] No models given: installing the package only " "(pass model keys — positional or --models KEY[,KEY…] — " "or --all to add models, or use 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, delete every catalog model's 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. No pip-uninstall phase: the venv removal IS the cleanup, and pip-ing the package plus every companion out of an environment that is about to be deleted is minutes of pure wait time. 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 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 0 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. The upgrade is followed by a companion refresh — every installed model's extras re-run (a pin already satisfied is a pip no-op, so this is cheap when nothing drifted) — so a newer sglang-omni's companion requirements are met the way a fresh install would meet them; a failing extra warns and leaves the update successful (the import probe re-heals it at the next model install or server start). 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. CANCEL is honored between phases only. 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") return rc print(f"[OK] {SGLOMNI_PIP_PKG} is up to date (or just upgraded).") for entry in sg_models.installed_entries(): crc = sg_models.install_companions(entry, emit=emit, cancel=cancel, force=True) if crc != 0: print(f"[WARNING] {entry.label}'s companion packages could " "not all be refreshed; the next install or server start " "retries what the import probe finds missing.") return 0 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())