diff options
Diffstat (limited to 'app/backends')
| -rw-r--r-- | app/backends/__init__.py | 15 | ||||
| -rw-r--r-- | app/backends/qwen.py | 270 |
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() |