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

#!/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())