feat: sglang backend support

1 files changed, 440 insertions, 0 deletions

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())