1 files changed, 283 insertions, 0 deletions
diff --git a/app/backends/sglomni/models.py b/app/backends/sglomni/models.py
new file mode 100644
index 0000000..a5211f6
--- /dev/null
+++ b/app/backends/sglomni/models.py
@@ -0,0 +1,283 @@
+"""sglang-omni model weights: install state on disk, (un)install actions.
+
+Model weights are not part of the pip install: each server fetches its
+HuggingFace repo into the standard hub cache on first start. This module
+pre-fetches ("Install") and deletes ("Uninstall") those cache directories
+per model — via the venv's hf CLI, exactly what a first server start
+would do — plus the per-model companion packages the server needs
+(``extras`` in the catalog, e.g. the qwen-tts --no-deps stack or the
+Descript DAC codec).
+
+The cache helpers mirror huggingface_hub's own directory layout and
+environment overrides (the same files ``from_pretrained`` writes), so an
+install lands exactly where a server start would look. The cache is
+shared with the other backends: a repo both host (Qwen3-TTS 1.7B Base /
+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
+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.
+
+
+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"
+
+
+def repo_dir(repo_id: str) -> Path:
+    """The cache directory HF keeps REPO_ID's weights in."""
+    return _hf_cache_dir() / ("models--" + repo_id.replace("/", "--"))
+
+
+def model_repo_dir(entry: ModelEntry) -> Path:
+    """The cached-weights directory for a catalog entry."""
+    return repo_dir(entry.repo)
+
+
+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
+
+
+def model_installed(entry: ModelEntry) -> bool:
+    """True when ENTRY's weights look complete in the local HF cache.
+
+    A fetched repo has refs/main plus at least one file under snapshots/;
+    anything less counts as not installed. An interrupted download simply
+    resumes — via Install, or the next server start for that model.
+    """
+    directory = model_repo_dir(entry)
+    if not (directory / "refs" / "main").is_file():
+        return False
+    return _tree_has_file(directory / "snapshots")
+
+
+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)]
+
+
+def installed_keys() -> List[str]:
+    """The installed entries' catalog keys, in catalog order."""
+    return [entry.key for entry in installed_entries()]
+
+
+def preset_voices(entry: ModelEntry) -> List[str]:
+    """The preset voice names ENTRY can speak with.
+
+    Catalog-declared speakers first (the Qwen3-TTS CustomVoice table);
+    otherwise the checkpoint's own ``voice_embedding/*.pt`` presets are
+    read from the downloaded snapshot (how Voxtral ships its named
+    voices). Empty when the model declares none or is not downloaded.
+    """
+    if entry.speakers:
+        return list(entry.speakers)
+    try:
+        snapshots = model_repo_dir(entry) / "snapshots"
+        for snapshot in sorted(snapshots.iterdir()):
+            voice_dir = snapshot / "voice_embedding"
+            if voice_dir.is_dir():
+                names = sorted(item.stem for item in voice_dir.glob("*.pt")
+                               if item.is_file())
+                if names:
+                    return names
+    except OSError:
+        pass
+    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):
+        return (f"{entry.system_dep} (system package) was not found — "
+                f"{entry.system_hint}. The weights still download, but the "
+                "server will fail to synthesize with this model until it "
+                "is installed.")
+    return None
+
+
+def install_model(key: str, *, emit=None, cancel=None) -> int:
+    """Install a catalog model: companion packages, then its weights.
+
+    Companion ``extras`` pip-install into the sglang-omni venv exactly as
+    upstream instructs (``--no-deps`` where upstream says so — the
+    Qwen3-TTS companions must not replace the pinned Transformers 5
+    stack), and a missing system binary is a loud warning, not a stop:
+    the download is still useful and the remediation stays on screen.
+    The weights pre-download via the venv's hf CLI (resumable, streamed,
+    cancelable). Returns the exit code.
+    """
+    entry = entry_by_key(key)
+    if entry is None:
+        print(f"[ERROR] Unknown sglang-omni model: {key!r}")
+        return 1
+    rc = prepare_env(emit=emit, cancel=cancel)
+    if rc != 0:
+        return rc
+    warning = system_dep_missing(entry)
+    if warning:
+        print(f"[WARNING] {warning}")
+    # A GPU the model's default FP8 pipeline cannot run gets the bf16
+    # fallback note up front (the install itself is still useful: the
+    # weights download either way).
+    from backends.sglomni import status as sg_status
+    note = sg_status.gpu_fallback_note(entry)
+    if note:
+        print(f"[WARNING] {note}")
+    for spec, no_deps in entry.extras:
+        args = ["--no-deps"] if no_deps else None
+        rc = common.pip_install([spec], emit=emit, cancel=cancel,
+                                env_dir=SGLOMNI_ENV, extra_args=args)
+        if rc != 0:
+            print(f"[WARNING] pip install {spec} failed (exit {rc}); "
+                  f"install it into {SGLOMNI_ENV} manually")
+    prefix = _hf_download_prefix()
+    if prefix is None:
+        print("[ERROR] No hf CLI found in the sglang-omni venv; pip "
+              f"install {SGLOMNI_PIP_PKG} first")
+        return 1
+    print(f"[INFO] Downloading {entry.repo} into {_hf_cache_dir()}...")
+    rc = common.run_console_subprocess(
+        prefix + ["download", entry.repo], emit=emit, cancel=cancel)
+    if rc == 0:
+        print(f"[OK] {entry.label} downloaded.")
+    return rc
+
+
+def uninstall_model(key: str, *, emit=None, cancel=None) -> int:
+    """Remove a model's cached weights (the inverse of install_model).
+
+    A locally-managed server currently hosting the model is stopped first
+    (best-effort) so its weights are not deleted under a live process.
+    CANCEL is honored after that stop phase only. Returns the exit code.
+    """
+    entry = entry_by_key(key)
+    if entry is None:
+        print(f"[ERROR] Unknown sglang-omni model: {key!r}")
+        return 1
+    if _managed_running_repo() == entry.repo:
+        from backends import servers
+        servers.stop(SERVER_NAME)
+    if common.cancel_requested(cancel):
+        return 130
+    delete_model_weights([entry])
+    return 0
+
+
+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.
+    """
+    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
+
+
+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
+
+
+def _managed_running_repo() -> Optional[str]:
+    """The repo id a locally-managed, up-and-running server hosts."""
+    from backends import probe, servers
+    from converter import config
+    if servers.pid_for(SERVER_NAME) is None:
+        return None
+    if not servers.alive(SERVER_NAME):
+        return None
+    return probe.sglomni_served_model(config.SGLOMNI_API_URL)
+
+
+def resolve_model(key: Optional[str]) -> ModelEntry:
+    """The catalog entry a run with MODEL_KEY uses.
+
+    An explicit KEY must exist in the catalog and be installed (a hosted
+    model without weights cannot boot). Without KEY the single installed
+    model is auto-selected; several installed models need an explicit pick
+    (the CLI --model flag or the Generate form's Model menu). Raises
+    RuntimeError with an actionable message otherwise.
+    """
+    if key is not None:
+        entry = entry_by_key(key)
+        if entry is None:
+            known = ", ".join(e.key for e in _all_entries())
+            raise RuntimeError(
+                f"Unknown sglang-omni model {key!r} (installed models are "
+                f"picked by catalog key; known keys: {known})")
+        if not model_installed(entry):
+            raise RuntimeError(
+                f"{entry.label} is not downloaded — install it via "
+                "Configure Backends → SGLang-Omni, or pick an installed "
+                "model.")
+        return entry
+    installed = installed_entries()
+    if not installed:
+        raise RuntimeError(
+            "No sglang-omni models are downloaded — install one via "
+            "Configure Backends → SGLang-Omni (Configure).")
+    if len(installed) > 1:
+        names = ", ".join(e.key for e in installed)
+        raise RuntimeError(
+            "Several sglang-omni models are installed; pick one with "
+            f"--model KEY (installed: {names})")
+    return installed[0]
+
+
+def entry_for_served_repo(repo: Optional[str]) -> Optional[ModelEntry]:
+    """The catalog entry a served /v1/models repo id belongs to."""
+    return entry_by_repo(repo) if repo else None