diff options
Diffstat (limited to 'app/backends/sglomni/wizard.py')
| -rw-r--r-- | app/backends/sglomni/wizard.py | 440 |
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()) |