feat: manage qwen-tts model installs manually, delete model(s) when uninstalled

2 files changed, 266 insertions, 19 deletions

diff --git a/app/backends/__init__.py b/app/backends/__init__.py
index 192067c..e59c464 100644
--- a/app/backends/__init__.py
+++ b/app/backends/__init__.py
@@ -18,8 +18,10 @@ importing this package must stay cheap and dependency-free.
 Adding a backend: create ``backends/<name>.py`` exposing
 ``detect() -> BackendStatus``, ``setup_screen(stdscr) -> int`` (the setup
 wizard run on the hub's own screen) and ``uninstall() -> int``, then append
-a ``BackendInfo`` in ``_build_registry`` below. ``audiobook.py`` and the hub
-pick it up automatically. A backend's standalone CLI keeps its own
+a ``BackendInfo`` in ``_build_registry`` below. Optionally expose a
+``configure_screen(stdscr) -> int`` to give the hub's "Configure <label>"
+entry somewhere to go besides re-running the wizard. ``audiobook.py`` and
+the hub pick it up automatically. A backend's standalone CLI keeps its own
 ``run_tui()`` entry (its own curses session), which is not part of the
 registry.
 """
@@ -158,12 +160,20 @@ class BackendInfo:
     removes the backend (stops its servers, pip-uninstalls, deletes its
     files); the hub runs it inside the task view, calling it with optional
     ``emit``/``cancel`` keywords (cancel honored between phases only).
+
+    CONFIGURE_SCREEN, when given, is what the hub's "Configure <label>"
+    menu entry runs instead of SETUP_SCREEN once the backend exists — a
+    place for install-time-independent management (qwen uses it for its
+    per-model weight installs). Without one, the hub falls back to
+    SETUP_SCREEN; a backend whose wizard asks nothing (bare qwen) offers
+    no Configure entry at all.
     """
     key: str
     label: str
     detect: Callable[[], BackendStatus]
     setup_screen: Callable[[object], int]
     uninstall: Callable[..., int] = lambda *args, **kwargs: 0
+    configure_screen: Optional[Callable[[object], int]] = None
 
 
 REGISTRY: List[BackendInfo] = []
@@ -189,6 +199,7 @@ def _build_registry() -> None:
         detect=qwen.detect,
         setup_screen=qwen.setup_screen,
         uninstall=qwen.uninstall,
+        configure_screen=qwen.models_screen,
     ))
     REGISTRY.append(BackendInfo(
         key="faster",
diff --git a/app/backends/qwen.py b/app/backends/qwen.py
index c0c2cd1..3309c69 100644
--- a/app/backends/qwen.py
+++ b/app/backends/qwen.py
@@ -7,14 +7,24 @@ ONE Qwen3-TTS model per process — CustomVoice (built-in speakers), Base
 end-to-end: pip-install the package into the managed venv. There are no
 questions to ask — the port and which model to run live in
 ``app/converter/config.py`` (the model is chosen per run on the hub's
-Generate-audiobooks screen), and only one server runs at a time. It is driven
-by ``audiobook.py``'s hub but can also be run directly:
+Generate-audiobooks screen), and only one server runs at a time.
+
+Model weights are not part of the install: each demo lazily fetches its
+~4GB snapshot from HuggingFace into the standard hub cache the first time a
+server for it starts. This module tracks those three cache directories so
+weights can be pre-fetched ("Install") or deleted ("Uninstall") per model —
+via the hub's Configure-qwen-tts screen — and so the backend uninstaller can
+remove every downloaded weight alongside the package.
+
+It is driven by ``audiobook.py``'s hub but can also be run directly:
 
 Usage:
     python app/backends/qwen.py [--skip-install]
 """
 
 import argparse
+import os
+import shutil
 import sys
 from pathlib import Path
 from typing import List, Optional
@@ -33,7 +43,7 @@ from backends import (
 )
 from converter import config
 from converter.clients import QWEN3_TTS_SPEAKERS
-from ui import taskview
+from ui import taskview, tui
 
 QWEN_PIP_PKG = "qwen-tts"
 DEFAULT_PORT = 7860
@@ -60,6 +70,96 @@ DEFAULT_MODEL = "CustomVoice"
 QWEN_SPEAKERS = QWEN3_TTS_SPEAKERS
 
 
+# -- HuggingFace weight cache ------------------------------------------------
+#
+# The demo servers pull each model's snapshot into huggingface_hub's default
+# hub cache on first start (nothing in this project redirects it). These
+# helpers read and delete the same directories ``from_pretrained`` writes,
+# honoring the same environment overrides, so per-model (un)installs land
+# exactly where a server start would look.
+
+def _hf_cache_dir() -> Path:
+    """The HF hub cache dir servers fetch model weights into.
+
+    Resolution mirrors huggingface_hub.constants: HF_HUB_CACHE beats
+    HUGGINGFACE_HUB_CACHE beats HF_HOME/hub beats ~/.cache/huggingface/hub.
+    """
+    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 (models--Qwen--…)."""
+    return _hf_cache_dir() / ("models--" + repo_id.replace("/", "--"))
+
+
+def model_repo_dir(model: str) -> Path:
+    """The cached-weights directory for a MODEL_REPOS key."""
+    return repo_dir(MODEL_REPOS[model])
+
+
+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(model: str) -> bool:
+    """True when MODEL'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(model)
+    if not (directory / "refs" / "main").is_file():
+        return False
+    return _tree_has_file(directory / "snapshots")
+
+
+def installed_models() -> List[str]:
+    """The MODEL_REPOS keys whose weights are already on disk."""
+    return [name for name in MODEL_REPOS if model_installed(name)]
+
+
+def delete_model_weights(models: Optional[List[str]] = None) -> int:
+    """Delete the cached HF weight dirs of MODELS (every model by default).
+
+    Best-effort rmtree of each ``models--Qwen--…`` 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. Prints its
+    progress, which streams into the task view under curses too.
+    """
+    names = sorted(MODEL_REPOS) if models is None else list(models)
+    removed = 0
+    for name in names:
+        directory = model_repo_dir(name)
+        if not directory.is_dir():
+            continue
+        print(f"[INFO] Removing cached {MODEL_REPOS[name]} weights...")
+        shutil.rmtree(directory, ignore_errors=True)
+        removed += 1
+    if removed:
+        print(f"[OK] Deleted cached weights for {removed} "
+              f"{'model' if removed == 1 else 'models'}.")
+    return removed
+
+
 def _is_installed() -> bool:
     if envs.env_script("qwen-tts-demo").is_file():
         return True
@@ -174,6 +274,30 @@ def build_parser() -> argparse.ArgumentParser:
     return parser
 
 
+def _build_spec(model: str) -> ServerSpec:
+    """The single managed ServerSpec hosting MODEL on the configured port."""
+    url = config.QWEN_API_URL
+    return ServerSpec(
+        "qwen", url,
+        [str(envs.env_script("qwen-tts-demo")), MODEL_REPOS[model],
+         "--ip", "127.0.0.1",
+         "--port", str(_config_port(url, DEFAULT_PORT))],
+        identity=desired_identity(model))
+
+
+def _managed_running_model() -> Optional[str]:
+    """The model a locally-managed, up-and-running demo answers as.
+
+    None when no pid file exists (this tool never started that server), the
+    process is gone, or the probe cannot identify which model it hosts.
+    """
+    if servers.pid_for("qwen") is None:
+        return None
+    if not servers.alive("qwen"):
+        return None
+    return model_for_identity(probe.identify_server(config.QWEN_API_URL))
+
+
 def detect() -> BackendStatus:
     """Detect whether qwen-tts is installed, plus the launch command.
 
@@ -192,13 +316,7 @@ def detect() -> BackendStatus:
     details.append(f"port: {_config_port(url, DEFAULT_PORT)}")
     details.append(f"model: {model}")
     details.append(f"speaker: {config.SPEAKER}")
-    demo = str(envs.env_script("qwen-tts-demo"))
-    specs = [
-        ServerSpec("qwen", url,
-                   [demo, MODEL_REPOS[model], "--ip", "127.0.0.1",
-                    "--port", str(_config_port(url, DEFAULT_PORT))],
-                   identity=desired_identity(model)),
-    ]
+    specs = [_build_spec(model)]
     managed = servers.manages(specs)
     # A locally-managed server names its running model via the probe of the
     # managed URL; a remotely-run demo names it via the remote-URL probe.
@@ -244,20 +362,72 @@ def _detect_remote(managed: bool = False):
     return remote_models, remote_urls
 
 
+def _hf_download_prefix() -> Optional[List[str]]:
+    """The venv's hf CLI argv prefix (None when neither script is present)."""
+    for name in ("hf", "huggingface-cli"):
+        candidate = envs.env_script(name)
+        if candidate.is_file():
+            return [str(candidate)]
+    return None
+
+
+def install_model(model: str, *, emit=None, cancel=None) -> int:
+    """Pre-download MODEL's weights into the HF cache via the venv's hf CLI.
+
+    Exactly what the first server start does implicitly, made explicit:
+    streamed progress through EMIT (percent bars feed the task view),
+    CANCEL kills the download process group mid-flight, and re-running
+    resumes where a previous attempt left off. Returns the exit code.
+    """
+    prefix = _hf_download_prefix()
+    if prefix is None:
+        print("[ERROR] No hf CLI found in the managed venv; pip install "
+              f"{QWEN_PIP_PKG} first")
+        return 1
+    repo = MODEL_REPOS[model]
+    print(f"[INFO] Downloading {repo} into {_hf_cache_dir()}...")
+    rc = common.run_console_subprocess(prefix + ["download", repo],
+                                       emit=emit, cancel=cancel)
+    if rc == 0:
+        print(f"[OK] {repo} downloaded.")
+    return rc
+
+
+def uninstall_model(model: str, *, emit=None, cancel=None) -> int:
+    """Remove MODEL's cached weights (the inverse of install_model).
+
+    A locally-managed server that currently answers as MODEL is stopped
+    first (best-effort) so its weights are not deleted under a live
+    process; no server, a down one, or one serving another model leaves
+    everything else untouched. CANCEL is honored after that stop phase
+    only. Returns the exit code.
+    """
+    if _managed_running_model() == model:
+        servers.stop("qwen")
+    if common.cancel_requested(cancel):
+        return 130
+    delete_model_weights([model])
+    return 0
+
+
 def uninstall(*, emit=None, cancel=None) -> int:
-    """Remove the qwen-tts backend entirely: stop its server, pip uninstall.
+    """Remove the qwen-tts backend entirely: stop its server, pip uninstall,
+    then delete every downloaded model.
 
     qwen-tts is a pip package (``qwen_tts`` + the ``qwen-tts-demo`` script)
     installed into the managed venv, so uninstalling it removes the backend.
-    Any server this tool started is stopped first (best-effort). Model
-    weights already fetched into the HuggingFace cache stay on disk.
+    Any server this tool started is stopped first (best-effort). The three
+    models' weight snapshots — multi-GB directories lazily fetched into the
+    HuggingFace cache (~/.cache/huggingface/hub) — are deleted too, matching
+    the hub's confirmation dialog; only those three directories are removed,
+    never the shared cache itself.
 
     With EMIT given (the in-TUI task view) pip runs piped, streaming into
     EMIT, so its output never touches the terminal behind curses. CANCEL is
-    a ``threading.Event`` honored between phases only (after the server has
-    been stopped, before pip starts) — a started phase always completes,
-    so pip is never killed mid-run. Returns the exit code (130 when
-    cancelled before pip ran).
+    a ``threading.Event`` honored between phases only (stop servers / pip /
+    delete models) — a started phase always completes, so pip is never
+    killed mid-run. Returns the exit code (130 when cancelled before a
+    remaining phase).
     """
     if servers.pid_for("qwen") is not None:
         # Only stop when a pid file exists: without one this tool never
@@ -272,9 +442,75 @@ def uninstall(*, emit=None, cancel=None) -> int:
               f"{QWEN_PIP_PKG} from the managed venv manually")
     else:
         print(f"[OK] {QWEN_PIP_PKG} removed.")
+    # Weights go even when the pip step failed: the package can be
+    # re-installed any time, multi-GB snapshots are what actually cost disk.
+    if common.cancel_requested(cancel):
+        return 130
+    delete_model_weights()
     return rc
 
 
+def models_screen(stdscr) -> int:
+    """Per-model (un)install screen: the hub's Configure-qwen-tts leaf.
+
+    Each of the three models gets exactly one action reflecting disk state:
+    Install pre-downloads its weights via the hf CLI (a streamed, resumable,
+    cancelable task-view run) and Uninstall deletes them again (stopping a
+    managed server that answers as that model first). Install requires the
+    pip package; without one a guidance flash replaces the download, since
+    pre-fetched weights without a backend to serve them buy nothing. The
+    status table and options re-render after every action, so Esc pops back
+    to Configure backends. Always returns 0.
+    """
+    while True:
+        present = installed_models()
+        rows = [(name, "installed", "ok") if name in present
+                else ("not installed", "warn") for name in MODEL_REPOS]
+        options = []
+        for name in MODEL_REPOS:
+            if name in present:
+                options.append((f"Uninstall {name}", ("uninstall", name)))
+            elif _is_installed():
+                options.append((f"Install {name}", ("install", name)))
+            else:
+                options.append((f"{name} (backend not installed)",
+                                ("noop", name)))
+
+        def make_work(action: str, target: str):
+            def work(emit, cancel) -> int:
+                if action == "install":
+                    return install_model(target, emit=emit, cancel=cancel)
+                return uninstall_model(target, emit=emit, cancel=cancel)
+            return work
+
+        choice = tui.menu(
+            stdscr, "Configure qwen-tts", options,
+            back_value=tui.Wizard.BACK,
+            help_lines=["Models are normally pulled when their server first",
+                        "starts; Install pre-downloads one right now."],
+            table_title="Model state", table_rows=rows)
+        if choice is tui.Wizard.BACK:
+            return 0
+        action, name = choice
+        if action == "noop":
+            continue
+        if action == "install" and not _is_installed():
+            tui.flash(stdscr, "Install the qwen-tts backend first "
+                      "(Configure backends > Install Backend).", "warn")
+            continue
+        title = (f"Download {MODEL_REPOS[name]}" if action == "install"
+                 else f"Delete {name} weights")
+        step = taskview.TaskStep(title, make_work(action, name))
+        rc = taskview.run_steps(stdscr, title, [step], wait_on_finish=False)
+        if rc == 0:
+            tui.flash(stdscr,
+                      f"{name} downloaded." if action == "install"
+                      else f"{name} weights removed.", "ok")
+        else:
+            verb = "download" if action == "install" else "remove"
+            tui.flash(stdscr, f"Could not {verb} {name}.", "err")
+
+
 def main() -> int:
     parser = build_parser()
     args = parser.parse_args()