"""Model install state: what is on disk, what is missing, how to fetch it.""" import json import os import shutil import sys import tempfile from pathlib import Path from typing import Callable, Dict, List, Optional, Set, Tuple from backends import common from . import catalog as _catalog def _install_models(audiocpp_dir: Path, install_guidance: List[Tuple[str, str]], download: bool, emit=None, cancel=None) -> int: """Print and optionally run the model install commands. One ``python install `` command per hosted model (de-duped by install id). When DOWNLOAD is True each command is run in the audio.cpp checkout via ``subprocess`` so the models are downloaded automatically; a failing install is reported as a warning and does not abort the remaining downloads. When DOWNLOAD is False (or the model manager is missing) the commands are only printed, copy-pasteable as before. With EMIT given (the in-TUI task view) each download streams its output to EMIT and — when the checkout's ``model_manager_v2.py`` supports it — runs with ``--progress --cancel-file`` so the view can show a real byte progress bar and cancel gracefully. CANCEL aborts a running download. Returns 0 when every command succeeded (or nothing needed running), 130 when cancelled, 1 when any download failed. """ manager = audiocpp_dir / "tools" / "model_manager_v2.py" seen: Set[str] = set() install_ids: List[str] = [] for _, install_id in install_guidance: if install_id not in seen: seen.add(install_id) install_ids.append(install_id) supports_progress = emit is not None and _manager_supports_progress(manager) if download and not manager.is_file(): print(f"[WARNING] {manager} not found; printing the install commands " "instead of running them") download = False failed = False for install_id in install_ids: command = f"python {manager} install {install_id}" if not download: print(command) continue print(f"[INFO] Downloading {install_id}...") argv = [sys.executable, str(manager), "install", install_id] cancel_file: Optional[Path] = None on_cancel = None if supports_progress: fd, cancel_path = tempfile.mkstemp( prefix="audiocpp_cancel_", suffix=".cancel") os.close(fd) cancel_file = Path(cancel_path) cancel_file.unlink() # absent = not cancelled argv += ["--progress", "--cancel-file", str(cancel_file)] on_cancel = cancel_file.touch try: rc = common.run_console_subprocess( argv, cwd=str(audiocpp_dir), emit=emit, cancel=cancel, on_cancel=on_cancel) except OSError as exc: print(f"[WARNING] Could not run {command}: {exc}") rc = 1 finally: if cancel_file is not None: try: cancel_file.unlink() except OSError: pass if rc == 130 or (cancel is not None and cancel.is_set()): return 130 if rc != 0: failed = True print(f"[WARNING] install {install_id} exited with code " f"{rc}; the model may need to be downloaded " "by hand") return 1 if failed else 0 def _manager_supports_progress(manager: Path) -> bool: """True when MANAGER (model_manager_v2.py) supports --progress output. The ``--progress``/``--cancel-file`` flags are relatively recent; an older audio.cpp checkout may not have them, so probe the script source once instead of failing the download with an unknown flag. """ try: text = manager.read_text(encoding="utf-8", errors="ignore") except OSError: return False return "AUDIOCPP_PROGRESS" in text and "--cancel-file" in text def _decide_download(audiocpp_dir: Path, model_entries: List[dict], confirm: Callable[[str, bool], bool]) -> bool: """Ask whether to download the selected models now. CONFIRM asks the yes/no question (ask_bool for the line prompts, a TUI confirm for the wizard). When the audio.cpp model manager is missing the prompt is skipped and False is returned, so the install commands are only printed rather than offered to run. The prompt is also skipped (False) when every selected model is already on disk (see ``_all_models_present``), so an already-configured checkout is not asked to re-download models it already has. """ manager = audiocpp_dir / "tools" / "model_manager_v2.py" if not manager.is_file(): return False if _all_models_present(audiocpp_dir, model_entries): return False return confirm( "Automatically download the selected models with model_manager_v2.py " "now?", True) def _build_tree_families(catalog: List[dict]) -> List[dict]: """Shape the catalog into the checkbox_tree widget's family list.""" families: List[dict] = [] for entry in catalog: capabilities = ["tts"] if "clone" in entry["tasks"]: capabilities.append("cloning") if "design" in entry["tasks"]: capabilities.append("design") name = entry["display_name"] options = [] for opt in _catalog.package_dir_options(entry): options.append({ "key": opt["target_directory"], "label": opt["install_id"], "recommended": opt["recommended"], }) families.append({ "label": name, "detail": ", ".join(capabilities), "options": options, }) return families def _model_path_present(path: Path) -> bool: """True when a server.json model path holds actual model files. A present path is either a file (a single-model package) or a non-empty directory (the usual GGUF package target directory; an empty one means a download that never ran or was cleaned up halfway). """ try: if path.is_file(): return True if path.is_dir(): return any(path.iterdir()) except OSError: return False return False def _all_models_present(audiocpp_dir: Path, model_entries: List[dict]) -> bool: """True when every selected model entry's path already holds files on disk. Paths resolve against the checkout (where model_manager_v2.py installs them), honoring absolute paths. Used by the wizard to skip the "Automatically download the selected models" prompt when nothing is actually missing. An empty selection is treated as not-present. """ if not model_entries: return False for entry in model_entries: rel = entry.get("path") if not isinstance(rel, str) or not rel: return False path = Path(rel) if Path(rel).is_absolute() else audiocpp_dir / rel if not _model_path_present(path): return False return True def missing_model_entries(server_json: Path) -> List[dict]: """Return the server.json model entries whose files are not on disk. Paths resolve exactly like audiocpp_server resolves them (relative paths against the server.json's directory). Each returned entry carries the entry ``id`` and ``rel`` (the configured path string); used by ``detect`` to warn that a conversion would fail until the models are installed. """ try: data = json.loads(server_json.read_text(encoding="utf-8")) except (OSError, ValueError): return [] if not isinstance(data, dict): return [] base = server_json.parent missing: List[dict] = [] for entry in data.get("models") or []: if not isinstance(entry, dict): continue rel = entry.get("path") if not isinstance(rel, str) or not rel: continue path = Path(rel) if Path(rel).is_absolute() else base / rel if _model_path_present(path): continue missing.append({"id": str(entry.get("id") or rel), "rel": rel}) return missing def _install_id_by_path(audiocpp_dir: Path) -> Dict[str, str]: """Map ``models/`` -> catalog install id. The catalog package that installs a model is derived from the ``default_path`` of each TTS family; an entry whose path matches no catalog package has no install id. """ by_path: Dict[str, str] = {} try: for entry in _catalog.load_model_catalog(audiocpp_dir): by_path[entry["default_path"]] = entry["install_id"] except (NotADirectoryError, OSError): pass return by_path def installed_model_entries(server_json: Path) -> List[dict]: """Return the server.json model entries whose files ARE on disk. The complement of ``missing_model_entries``: each returned entry carries the entry ``id`` and ``rel`` (the configured path string), resolved exactly like ``missing_model_entries`` (relative against the server.json's directory). Used by the wizard's "Delete unused models?" step to find already-downloaded models that were unselected. """ try: data = json.loads(server_json.read_text(encoding="utf-8")) except (OSError, ValueError): return [] if not isinstance(data, dict): return [] base = server_json.parent installed: List[dict] = [] for entry in data.get("models") or []: if not isinstance(entry, dict): continue rel = entry.get("path") if not isinstance(rel, str) or not rel: continue path = Path(rel) if Path(rel).is_absolute() else base / rel if _model_path_present(path): installed.append({"id": str(entry.get("id") or rel), "rel": rel}) return installed def missing_model_install_guidance(audiocpp_dir: Path, missing: List[dict]) -> List[Tuple[str, str]]: """Map MISSING model entries to (display name, install id) pairs. The install id is derived from each entry's configured path via the catalog (see ``_install_id_by_path``); entries whose path matches no catalog package are skipped (there is no ``model_manager_v2.py install`` command for them). Feeds ``_install_models`` for the "Download Missing Models" action. """ by_path = _install_id_by_path(audiocpp_dir) guidance: List[Tuple[str, str]] = [] for item in missing: install_id = by_path.get(item["rel"]) if install_id: guidance.append((item["id"], install_id)) return guidance def model_install_hints(audiocpp_dir: Path, missing: List[dict]) -> List[str]: """Remediation lines for MISSING model entries (see missing_model_entries). Maps each entry's configured path back to the catalog package that installs it (``models/`` -> install id) so the line carries the exact ``model_manager_v2.py install`` command; entries whose directory matches no catalog package just name the path. """ by_path = _install_id_by_path(audiocpp_dir) hints: List[str] = [] for item in missing: install_id = by_path.get(item["rel"]) hint = f"model not downloaded: {item['id']} ({item['rel']})" if install_id: hint += (f" — install with: python tools/model_manager_v2.py " f"install {install_id}") hints.append(hint) return hints def install_models(audiocpp_dir: Path, guidance: List[Tuple[str, str]], emit=None, cancel=None) -> int: """Download the (display name, install id) models via the helper script. Runs ``model_manager_v2.py install`` for each de-duped install id in the checkout, streaming to the console (or to EMIT, the in-TUI task view); a failing install is reported as a warning and does not abort the rest. Returns 0 when every download succeeded, 130 when cancelled, 1 when any failed. Used by the hub's "Download Missing Models" action (see ``missing_model_install_guidance`` for the mapping). """ return _install_models(audiocpp_dir, guidance, download=True, emit=emit, cancel=cancel) def hand_install_guidance(audiocpp_dir: Path, missing: List[dict]) -> str: """Explain how to install MISSING model entries by hand. Returns a multi-line message listing each missing model's id and the path its files must be placed in (``rel``, resolved against the checkout). Used when the missing models cannot be mapped to a ``model_manager_v2.py install`` command, so the user still knows what to download and where to put it. """ lines = [ "None of the missing models map to a model_manager_v2.py install " "command.", "Download them by hand and place the files at these paths:", ] for item in missing: lines.append(f" {item['id']} -> {item['rel']}") lines.append(f"(paths are relative to {audiocpp_dir})") return "\n".join(lines) def unused_installed_entries(server_json: Path, new_paths: Set[str]) -> List[dict]: """Return installed server.json entries whose path is not in NEW_PATHS. The already-downloaded models (see ``installed_model_entries``) that the new selection does not host any more — the candidates for the wizard's "Delete unused models?" prompt. Entries whose files are not on disk are never listed (there is nothing to delete). """ return [entry for entry in installed_model_entries(server_json) if entry["rel"] not in new_paths] def delete_model_files(server_json: Path, entries: List[dict]) -> int: """Remove the on-disk model files for ENTRIES ({id, rel}) from disk. Each entry's ``rel`` is resolved exactly like the server resolves it (relative against ``server_json``'s directory; absolute paths honored), then removed as a directory tree or a single file. Missing entries are ignored. Returns the number of paths removed. Used by the wizard's "Delete unused models?" step — the regenerated server.json already only lists the kept models, so no entry cleanup is needed here. """ base = server_json.parent removed = 0 for item in entries: rel = item.get("rel") if not isinstance(rel, str) or not rel: continue path = Path(rel) if Path(rel).is_absolute() else base / rel try: if not path.exists(): continue if path.is_dir(): shutil.rmtree(path, ignore_errors=True) else: path.unlink() except OSError as exc: print(f"[WARNING] Could not remove {path}: {exc}") continue print(f"[OK] Removed unused model {path}") removed += 1 return removed