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

"""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, 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.


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 missing_companions(entry: ModelEntry) -> List[Extra]:
    """ENTRY's companion packages absent from the sglang-omni venv.

    One import probe per extra (its top-level module, see
    ``extra_import_name``) against the venv's interpreter, so a model whose
    weights landed in the shared HF cache by another route — or whose
    companions failed to pip-install during setup, which install_model only
    warns about — is detected before the server dies on the import. A venv
    that does not exist at all yields no verdict: the start flow fails on
    the missing ``sgl-omni`` executable anyway."""
    if not entry.extras or not envs.env_exists(SGLOMNI_ENV):
        return []
    return [extra for extra in entry.extras
            if not envs.module_available(extra_import_name(extra[0]),
                                         SGLOMNI_ENV)]


def install_companions(entry: ModelEntry, *, emit=None, cancel=None) -> int:
    """pip-install ENTRY's missing 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):
        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"[ERROR] pip install {spec} failed (exit {rc}); "
                  f"install it into {SGLOMNI_ENV} manually")
            return rc
    return 0


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 missing_companions(entry):
        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