fix: sglomni retry contract, booting detection, CLI keys, companion refresh, ffmpeg depHEADmain

7 files changed, 104 insertions, 124 deletions

diff --git a/app/backends/sglomni/__init__.py b/app/backends/sglomni/__init__.py
index be94a0e..ed7d4ed 100644
--- a/app/backends/sglomni/__init__.py
+++ b/app/backends/sglomni/__init__.py
@@ -22,7 +22,6 @@ from .constants import (
     SERVER_NAME,
     SERVER_START_TIMEOUT,
     SGLOMNI_PIP_PKG,
-    UV_PIP_PKG,
 )
 from .gpu import (
     compute_capability,
@@ -79,7 +78,7 @@ from .wizard import (
 __all__ = [
     # constants
     "CONFIGS_DIR", "DEFAULT_PORT", "PYTHON_SPEC", "PYTHON_VERSIONS",
-    "SERVER_NAME", "SERVER_START_TIMEOUT", "SGLOMNI_PIP_PKG", "UV_PIP_PKG",
+    "SERVER_NAME", "SERVER_START_TIMEOUT", "SGLOMNI_PIP_PKG",
     # gpu
     "compute_capability", "describe",
     # catalog
diff --git a/app/backends/sglomni/catalog.py b/app/backends/sglomni/catalog.py
index f20509f..8bfb1fa 100644
--- a/app/backends/sglomni/catalog.py
+++ b/app/backends/sglomni/catalog.py
@@ -52,6 +52,11 @@ class ModelEntry:
     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)
+    # NOTE(unverified upstream): only the two Base entries are known to
+    # accept a request-scoped seed (Voxtral rejects one outright); qwen's
+    # demo client does send seeds to the CustomVoice/VoiceDesign models,
+    # so those pipelines may accept one too — verify before flipping the
+    # flag (CONSTANT_SEED currently no-ops for every other entry).
     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
@@ -103,6 +108,12 @@ _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 decode their codec assets through
+# the ffmpeg binary (and the client concatenates multi-part chunks with
+# it); the weights download fine without it, the server just fails to
+# synthesize — so the install flow warns, like it does for sox.
+_FFMPEG_HINT = ("install the ffmpeg system package (e.g. sudo pacman -S "
+                "ffmpeg, sudo apt install ffmpeg, brew install ffmpeg)")
 # The Fish Audio and ZONOS2 pipelines use the Descript DAC codec, which
 # upstream installs WITH dependencies — but descript-audiotools carries a
 # vestigial 2021-era pin, protobuf<3.20 (its code never imports protobuf),
@@ -255,6 +266,7 @@ ENTRIES: Tuple[ModelEntry, ...] = (
         capability=CAPABILITY_CLONE,
         requires_reference=False,
         extras=_DAC_EXTRAS,
+        system_dep="ffmpeg", system_hint=_FFMPEG_HINT,
         notes="zero-shot narration or cloning from a reference clip",
     ),
     ModelEntry(
@@ -265,6 +277,7 @@ ENTRIES: Tuple[ModelEntry, ...] = (
         capability=CAPABILITY_CLONE,
         requires_reference=True,
         extras=_DAC_EXTRAS,
+        system_dep="ffmpeg", system_hint=_FFMPEG_HINT,
         notes="voice cloning, 44.1 kHz DAC vocoder",
         fp8_moe=True,
         fp8_min_compute_capability=(8, 9),
diff --git a/app/backends/sglomni/constants.py b/app/backends/sglomni/constants.py
index 94a07bd..945ce8a 100644
--- a/app/backends/sglomni/constants.py
+++ b/app/backends/sglomni/constants.py
@@ -6,9 +6,6 @@ from pathlib import Path
 # (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
diff --git a/app/backends/sglomni/models.py b/app/backends/sglomni/models.py
index 607e36a..990673a 100644
--- a/app/backends/sglomni/models.py
+++ b/app/backends/sglomni/models.py
@@ -16,31 +16,25 @@ 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, extra_import_name
+from backends.sglomni.catalog import ENTRIES, Extra, ModelEntry, \
+    entry_by_key, entry_by_repo, extra_import_name
 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.
+# The HF-cache primitives live in backends.common (shared with the qwen
+# backend, which fetches into the same cache); the module-level wrappers
+# below keep the names internal callers (and the tests) reference.
 
 
 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"
+    """The cache directory HF keeps repos in (models--<org>--<name>
+    folders); common.hf_cache_dir resolves the environment overrides."""
+    return common.hf_cache_dir()
 
 
 def repo_dir(repo_id: str) -> Path:
@@ -55,17 +49,7 @@ def model_repo_dir(entry: ModelEntry) -> Path:
 
 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
+    return common.hf_tree_has_file(path)
 
 
 def model_installed(entry: ModelEntry) -> bool:
@@ -83,7 +67,7 @@ def model_installed(entry: ModelEntry) -> bool:
 
 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)]
+    return [entry for entry in ENTRIES if model_installed(entry)]
 
 
 def installed_keys() -> List[str]:
@@ -115,11 +99,6 @@ def preset_voices(entry: ModelEntry) -> List[str]:
     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):
@@ -147,15 +126,20 @@ def missing_companions(entry: ModelEntry) -> List[Extra]:
                                          SGLOMNI_ENV)]
 
 
-def install_companions(entry: ModelEntry, *, emit=None, cancel=None) -> int:
-    """pip-install ENTRY's missing companion packages into the venv.
+def install_companions(entry: ModelEntry, *, emit=None, cancel=None,
+                       force: bool = False) -> int:
+    """pip-install ENTRY's companion packages into the venv.
 
     The same recipe ``install_model`` runs (the catalog's ``--no-deps``
     flags preserved — the Qwen3-TTS companions must not replace the pinned
     Transformers 5 stack), limited to what the import probe found absent,
     so a start-time heal touches as little of the pinned environment as
-    possible. Returns the first failing exit code, 0 when all present."""
-    for spec, no_deps in missing_companions(entry):
+    possible. With FORCE every extra's pip spec re-runs instead — a
+    satisfied pin is a pip no-op, so the update flow uses that to heal
+    version drift the import probe cannot see (the protobuf re-pin
+    especially). Returns the first failing exit code, 0 when all present."""
+    wanted = list(entry.extras) if force else missing_companions(entry)
+    for spec, no_deps in wanted:
         args = ["--no-deps"] if no_deps else None
         rc = common.pip_install([spec], emit=emit, cancel=cancel,
                                 env_dir=SGLOMNI_ENV, extra_args=args)
@@ -237,36 +221,19 @@ def uninstall_model(key: str, *, emit=None, cancel=None) -> int:
 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.
+    Delegates to common.hf_delete_model_weights: best-effort rmtree of
+    each ``models--<org>--<name>`` directory; only those directories are
+    ever touched — the rest of the HF cache may be shared with unrelated
+    tools. Returns how many were present and removed.
     """
     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
+        entries = list(ENTRIES)
+    return common.hf_delete_model_weights([entry.repo for entry in entries])
 
 
 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
+    return common.hf_download_prefix(SGLOMNI_ENV)
 
 
 def _managed_running_repo() -> Optional[str]:
@@ -292,7 +259,7 @@ def resolve_model(key: Optional[str]) -> ModelEntry:
     if key is not None:
         entry = entry_by_key(key)
         if entry is None:
-            known = ", ".join(e.key for e in _all_entries())
+            known = ", ".join(e.key for e in ENTRIES)
             raise RuntimeError(
                 f"Unknown sglang-omni model {key!r} (installed models are "
                 f"picked by catalog key; known keys: {known})")
diff --git a/app/backends/sglomni/pythonenv.py b/app/backends/sglomni/pythonenv.py
index 11bbaa6..cb19710 100644
--- a/app/backends/sglomni/pythonenv.py
+++ b/app/backends/sglomni/pythonenv.py
@@ -21,10 +21,9 @@ 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 import envs
 from backends.sglomni.constants import PYTHON_SPEC, PYTHON_VERSIONS
 
 # The dedicated venv (its interpreter may differ from the launching one).
diff --git a/app/backends/sglomni/status.py b/app/backends/sglomni/status.py
index 2b77307..6a47cff 100644
--- a/app/backends/sglomni/status.py
+++ b/app/backends/sglomni/status.py
@@ -3,8 +3,8 @@
 from pathlib import Path
 from typing import List, Optional
 
-from backends import BackendStatus, ServerSpec, envs, format_launch_hint, \
-    probe, servers
+from backends import BackendStatus, ServerSpec, common, 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
@@ -101,15 +101,9 @@ def gpu_fallback_note(entry: ModelEntry) -> Optional[str]:
 
 
 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
+    """The managed server's port (the configured URL's explicit port,
+    else the backend default)."""
+    return common.port_of(config.SGLOMNI_API_URL, DEFAULT_PORT)
 
 
 def _managed_running_entry() -> Optional[ModelEntry]:
diff --git a/app/backends/sglomni/wizard.py b/app/backends/sglomni/wizard.py
index 96679d6..f39e754 100644
--- a/app/backends/sglomni/wizard.py
+++ b/app/backends/sglomni/wizard.py
@@ -13,7 +13,7 @@ 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]
+    python -m backends.sglomni [KEY ...] [--models KEY[,KEY...]] [--all]
         [--skip-install] [--skip-python]
 
 The Configure screen (``models_screen``) manages models after the fact —
@@ -32,9 +32,9 @@ 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, \
-    UV_PIP_PKG
+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
@@ -43,10 +43,11 @@ _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
+    """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]:
@@ -240,6 +241,10 @@ def run_tui(args: Optional[argparse.Namespace] = None) -> int:
         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)
@@ -258,21 +263,27 @@ def _collect_from_flags(args: argparse.Namespace,
         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(","):
+    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})")
-            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).")
+            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,
@@ -357,33 +368,20 @@ def models_screen(stdscr) -> int:
 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).
+    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
-    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
@@ -395,17 +393,23 @@ def uninstall(*, emit=None, cancel=None) -> int:
                 print(f"[WARNING] Could not fully remove {directory}")
             else:
                 print(f"[OK] {directory} removed.")
-    return rc
+    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. 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.
+    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)
@@ -420,9 +424,16 @@ def update(*, emit=None, cancel=None) -> int:
     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
+        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: