"""The model_specs catalog, server.json building and selection views.""" import json import re from pathlib import Path from typing import Dict, List, Optional, Set, Tuple from .constants import TASK_TTS DESIGN_PACKAGE_RE = re.compile(r"voice[\s_\-]?design", re.IGNORECASE) _BACKEND_DESCRIPTIONS = ( ("cuda", "NVIDIA GPUs (fastest)"), ("vulkan", "cross-vendor GPU"), ("hip", "AMD GPUs"), ("cpu", "no GPU required"), ) def _backend_options(detected: Optional[str] = None ) -> Tuple[List[Tuple[str, str]], int]: """Build the aligned backend menu options and the default index. The backend names are padded to a common width so the ``-`` dashes before the descriptions line up. When DETECTED matches one of the options, that option gets ``[auto-detected]`` appended and is the default (cursor/start) selection; otherwise the first option is the default as before. Returns (options, default_index). """ width = max(len(name) for name, _ in _BACKEND_DESCRIPTIONS) options: List[Tuple[str, str]] = [] default_index = 0 for index, (name, desc) in enumerate(_BACKEND_DESCRIPTIONS): label = f"{name.ljust(width)} - {desc}" if detected == name: label += " [auto-detected]" default_index = index options.append((label, name)) return options, default_index _BACKEND_TOKEN_RE = re.compile(r"-(cuda|vulkan|hip|cpu|metal)(?:-|$)") def detect_backend(audiocpp_dir: Path) -> Optional[str]: """Best-effort detection of the backend audiocpp_server was built for. Scans ``audiocpp_dir/build/*`` for build directories that contain a built ``bin/audiocpp_server`` (``.exe`` allowed on Windows) and reads the backend token out of the directory name (``-cuda-``, ``-vulkan-``, ``-hip-`` or ``-cpu-``; ``-metal-`` is mapped to ``cpu``). Returns the backend only when exactly one distinct backend was built, so a checkout with builds for several backends does not silently pick one. Returns None when there is no ``build/`` directory, no built server, or more than one distinct backend. """ build_root = audiocpp_dir / "build" if not build_root.is_dir(): return None backends: Set[str] = set() try: build_dirs = sorted(build_root.iterdir(), key=lambda p: p.name.lower()) except OSError: return None for build_dir in build_dirs: if not build_dir.is_dir(): continue server = build_dir / "bin" / "audiocpp_server" if not server.exists(): server_exe = build_dir / "bin" / "audiocpp_server.exe" if not server_exe.exists(): continue match = _BACKEND_TOKEN_RE.search(build_dir.name.lower()) if not match: continue token = match.group(1) backends.add("cpu" if token == "metal" else token) if len(backends) == 1: return next(iter(backends)) return None def _default_package(packages: List[dict]) -> Optional[dict]: """Pick the default package from a list of packages. Prefers the package flagged ``default: true``, then the first GGUF package, then the first package overall. Returns None for an empty list. """ if not packages: return None for package in packages: if package.get("default"): return package for package in packages: if package.get("format") == "gguf": return package return packages[0] def load_model_catalog(audiocpp_dir: Path) -> List[dict]: """Read model_specs/*.json and return the TTS-capable families. Each returned entry has: family, display_name, description, languages, clone_capable, packages (the full list from the spec), install_id (recommended package id), and default_path (``models/``). All families are treated equally and listed in alphabetical order by display name. """ specs_dir = audiocpp_dir / "model_specs" if not specs_dir.is_dir(): raise NotADirectoryError( f"{audiocpp_dir} has no model_specs/ directory; re-run setup " "to refresh the audio.cpp checkout") entries: List[dict] = [] for spec_path in sorted(specs_dir.glob("*.json")): try: spec = json.loads(spec_path.read_text(encoding="utf-8")) except (OSError, ValueError): continue tasks = spec.get("tasks") or [] if "tts" not in tasks and spec.get("category") != "tts": continue family = spec.get("family") or spec_path.stem packages = spec.get("packages") or [] package = _default_package(packages) if package is None: # No installable package: skip (cannot be hosted from a path). continue target_directory = package.get("target_directory") or family languages = spec.get("languages") or [] display_name = spec.get("display_name") or family description = spec.get("description") or "" entries.append({ "family": family, "display_name": display_name, "description": description, "languages": languages, "tasks": list(tasks), "clone_capable": "clone" in tasks, "packages": packages, "install_id": package.get("id") or family, "default_path": f"models/{target_directory}", }) # All families are treated equally: alphabetical by display name. entries.sort(key=lambda entry: entry["display_name"].lower()) return entries def is_design_package(package: dict) -> bool: """Return True when a package's name marks it a voice-design model. audio.cpp voice-design packages (whose id, display name, or target directory mentions "voice design") are the only packages that must be hosted with task "vdes"; their role is not in the schema, only in those strings, so it is detected from them. """ text = " ".join(str(package.get(key, "")) for key in ("id", "display_name", "target_directory")) return bool(DESIGN_PACKAGE_RE.search(text)) def package_dir_options(entry: dict) -> List[dict]: """Return one option per distinct target_directory of a family's packages. Each option is a dict with: target_directory, install_id (the recommended package id inside that directory), design (voice-design package flag), and recommended (whether it holds the family's default package). Precisions that share a directory (q8_0/bf16/...) collapse to a single option. """ packages = entry.get("packages") or [] default_pkg = _default_package(packages) default_dir = (default_pkg or {}).get("target_directory") or entry["family"] by_dir: Dict[str, List[dict]] = {} order: List[str] = [] for package in packages: directory = package.get("target_directory") or entry["family"] if directory not in by_dir: by_dir[directory] = [] order.append(directory) by_dir[directory].append(package) options: List[dict] = [] for directory in order: package = _default_package(by_dir[directory]) options.append({ "target_directory": directory, "install_id": (package or {}).get("id") or directory, "design": is_design_package(package or {}), "recommended": directory == default_dir, }) # Put the recommended package first for a friendlier checklist. options.sort(key=lambda opt: not opt["recommended"]) return options def build_model_entry(family: str, model_id: str, model_path: str, task: str = TASK_TTS) -> dict: """Assemble one server.json model entry. ``task`` defaults to "tts"; voice design packages are hosted with "vdes" so the server runs its design session for speech requests (audiobook.py then requires --instructions with that entry). """ return { "id": model_id, "family": family, "path": model_path, "task": task, "mode": "offline", } def build_server_config(host: str, port: int, backend: str, lazy_load: bool, model_entries: List[dict], voice_dir: Optional[str] = None) -> dict: """Assemble the server.json document. ``voice_dir`` is a server-level cloning voice library; when set, every hosted clone-capable family can use its voices with ``--voice``. """ config_doc = { "host": host, "port": port, "backend": backend, "lazy_load": lazy_load, "models": model_entries, } if voice_dir: config_doc["voice_dir"] = voice_dir return config_doc def load_server_config(server_json: Path) -> Optional[dict]: """Read server.json into a dict, or None when it cannot be used. Returns None for a missing file, unreadable content, or a non-dict document. Used by the wizard's modify flow to pre-fill its screens from an existing config instead of prompting to overwrite it. """ if not server_json.exists(): return None try: data = json.loads(server_json.read_text(encoding="utf-8")) except (OSError, ValueError): return None if not isinstance(data, dict): return None return data def server_config_selections(server_config: dict, catalog: List[dict] ) -> Tuple[Dict[str, List[str]], Dict[Tuple[str, str], str]]: """Map an existing server.json's models back to catalog selections. Returns ``(selected_dirs, tasks)``: ``selected_dirs`` maps a catalog family to the target directories it hosts (``models/`` paths with the ``models/`` prefix stripped, in server.json order), and ``tasks`` maps ``(family, target_directory)`` to the entry's task (``"tts"`` or ``"vdes"``) so the wizard can preserve how design packages were hosted. Entries whose family is not in the CATALOG are ignored — the wizard cannot offer them again. """ families = {entry["family"] for entry in catalog} selected_dirs: Dict[str, List[str]] = {} tasks: Dict[Tuple[str, str], str] = {} for entry in server_config.get("models") or []: if not isinstance(entry, dict): continue family = entry.get("family") if not isinstance(family, str) or family not in families: continue path = entry.get("path") if not isinstance(path, str): continue target = path[len("models/"):] if path.startswith("models/") else path if family not in selected_dirs: selected_dirs[family] = [] if target not in selected_dirs[family]: selected_dirs[family].append(target) tasks[(family, target)] = str(entry.get("task") or TASK_TTS) return selected_dirs, tasks