path: root/app/backends/audiocpp/catalog.py

"""The model_specs catalog, server.json building and selection views."""

import json
import re
import sys
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

from converter.clients import AUDIOCPP_CLONE_ONLY_FAMILIES, audiocpp_family_spec_tasks

from .constants import TASK_CLON, TASK_TTS, TASK_VDES

DESIGN_PACKAGE_RE = re.compile(r"voice[\s_\-]?design", re.IGNORECASE)

# The spec task vocabulary for "can this family turn text into audio":
# plain synthesis ("tts"), reference-voice synthesis ("clone"; specs spell
# it out, unlike the hosted "clon" task) and described-voice synthesis
# ("vdes"). Families whose tasks name none of these only transform audio
# (speech-to-speech, voice conversion, ...) and are kept out of the
# install catalog.
NARRATION_TASKS = frozenset({TASK_TTS, "clone", TASK_VDES})


def request_options_families(audiocpp_dir: Path) -> Dict[str, dict]:
    """Map the families whose spec defines per-request options.

    Reads every ``model_specs/<family>.json`` in AUDIOCPP_DIR once and
    returns ``{family_key: {"display_name": ...}}`` for the specs that
    list request options (a non-empty ``options.request`` array) — the
    families a "Request options" field makes sense for. The key is the
    spec's ``family`` field (falling back to the file stem), matching
    what ``GET /v1/models`` reports, so callers can look an entry up by
    its family id. A missing or unreadable specs directory yields {}
    (every family then counts as unknown rather than unsupported).
    """
    specs_dir = audiocpp_dir / "model_specs"
    if not specs_dir.is_dir():
        return {}
    families: Dict[str, 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
        options = spec.get("options")
        request = options.get("request") if isinstance(options, dict) else None
        if not isinstance(request, list) or not request:
            continue
        family = str(spec.get("family") or spec_path.stem)
        families[family] = {
            "display_name": str(spec.get("display_name") or family),
        }
    return families


def supports_request_options(families: Dict[str, dict],
                             family: str) -> Optional[bool]:
    """Whether FAMILY accepts per-request options — None when unknown.

    True only when FAMILIES (from request_options_families) lists the
    family; False when it was read but does not define request options;
    None when support cannot be determined from the local specs (no
    checkout, or a family the specs do not describe).
    """
    if not families:
        return None
    return family in families


_BACKEND_DESCRIPTIONS = (
    ("cuda", "NVIDIA GPUs"),
    ("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).

    On macOS only ``cpu`` is offered: Metal (via ``build_metal.sh``) is
    the only buildable inference backend there, and it is recorded as
    ``cpu`` (see the ``-metal-`` -> ``cpu`` mapping in ``detect_backend``
    and ``built_server_binary``), so cuda/vulkan/hip — which cannot build
    on macOS — never reach the build step.
    """
    if sys.platform == "darwin":
        label = "cpu - Apple Metal"
        if detected == "cpu":
            label += " [auto-detected]"
        return [(label, "cpu")], 0
    width = max(len(name) for name, _ in _BACKEND_DESCRIPTIONS)
    options = []
    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/<target_directory>``).
    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 tasks:
            # A task list that names no text-synthesis capability means the
            # family cannot narrate text at all (e.g. PersonaPlex:
            # categorized "tts" but speech-to-speech only) — hosting one
            # fails every request, so it is never offered for install.
            if not (NARRATION_TASKS & set(tasks)):
                continue
        elif spec.get("category") != "tts":
            # No task list: fall back to the category as before.
            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 is_clone_only_family(family: str, tasks: Optional[Set[str]] = None
                         ) -> bool:
    """True when FAMILY's audio.cpp implementation rejects plain TTS.

    Such families can only synthesize by cloning a reference voice, so
    their server entries must be hosted with task "clon" — hosting them
    with "tts" fails every speech request at session-creation time.
    Families are classified from the explicit known-clone-only set
    (AUDIOCPP_CLONE_ONLY_FAMILIES, which also covers specs that wrongly
    claim "tts" — Chatterbox) or from a spec task list that names only
    "clone" (TASKS, when the caller has it; without one the specs are
    read best-effort).
    """
    if family in AUDIOCPP_CLONE_ONLY_FAMILIES:
        return True
    if tasks is None:
        tasks = audiocpp_family_spec_tasks(family)
    return bool(tasks) and set(tasks) == {"clone"}


def hosting_task(entry: dict) -> str:
    """The server.json task a family's non-design packages are hosted with.

    Clone-only families (see ``is_clone_only_family``) get "clon" so their
    cloning sessions can be created at all; every other family keeps
    "tts", which serves plain TTS and — where the family supports it —
    cloning through the request's voice field alike.
    """
    if is_clone_only_family(str(entry.get("family") or ""),
                            tasks=set(entry.get("tasks") or []) or None):
        return TASK_CLON
    return TASK_TTS


def rehost_clone_only_entries(server_json: Path, data: dict) -> List[str]:
    """Re-host "tts"-tasked clone-only entries in DATA as "clon", in place.

    Server.json files written before clone-only hosting existed carry
    task "tts" for families whose audio.cpp implementation rejects plain
    TTS sessions (e.g. Chatterbox), so every speech request fails with
    HTTP 500. Each such entry's task is rewritten to "clon"; when anything
    changed, the document is written back to SERVER_JSON (same layout the
    wizard writes). Returns the repaired entries' ids, in order — empty
    when nothing needed changing (or the document is unusable).
    """
    models = data.get("models")
    if not isinstance(models, list):
        return []
    repaired: List[str] = []
    for entry in models:
        if not isinstance(entry, dict):
            continue
        if str(entry.get("task") or "") != TASK_TTS:
            continue
        family = str(entry.get("family") or "")
        if not family or not is_clone_only_family(family):
            continue
        entry["task"] = TASK_CLON
        repaired.append(str(entry.get("id") or family))
    if repaired:
        try:
            with server_json.open("w", encoding="utf-8") as handle:
                json.dump(data, handle, indent=2, ensure_ascii=False)
                handle.write("\n")
        except OSError:
            # The in-memory document is fixed either way; a failed write
            # only means the fix does not survive the process.
            pass
    return repaired


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/<target>`` 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