gart/admin/api/models.py

"""Catálogo de modelos: listar, baixar, escolher, apagar.

Extraído de models_api.py — a tabela de comandos segue lá.
"""

from __future__ import annotations

import shutil
import subprocess
import threading

from fcpxml.diarize import (
    diarization_capability,
)
from fcpxml.model_manager import (
    download_model,
    get_models_dir,
    is_model_downloaded,
    list_installed_models,
    load_catalog,
    load_hf_token,
    load_num_speakers,
    load_selected_model,
    load_transcript_language,
    model_cache_dir,
    save_models_dir,
    save_selected_model,
    save_transcript_language,
)

from . import shared
from .shared import (
    RECOMMENDED,
)

# Downloads em andamento, para o comando `cancel` conseguir interrompê-los.
# Mora aqui, e não no shared, porque só `download` e `cancel` o tocam — e o
# lock é próprio: ele protege este dicionário, não a saída em stdout.
_CANCEL: dict[str, threading.Event] = {}
_CANCEL_LOCK = threading.Lock()


def cmd_catalog() -> None:
    catalog = load_catalog()
    installed = list_installed_models()
    diar_ok, diar_msg = diarization_capability(load_hf_token())
    shared.emit(
        {
            "models": catalog,
            "installed": installed,
            "selected": load_selected_model(),
            "language": load_transcript_language(),
            "models_dir": str(get_models_dir()),
            "installed_count": len(installed),
            "recommended": list(RECOMMENDED),
            "diarization": diar_ok,
            "diarization_message": diar_msg,
            "hf_token_set": bool(load_hf_token()),
            "num_speakers": load_num_speakers(),
        }
    )

def cmd_download(args: dict) -> int:
    model = str(args.get("model", ""))
    if model not in _model_names():
        shared.emit({"type": "error", "message": f"Modelo desconhecido: {model}"})
        return 1
    ev = threading.Event()
    with _CANCEL_LOCK:
        _CANCEL[model] = ev
    try:
        download_model(model, progress_cb=lambda f: shared.emit({"type": "progress", "fraction": f}), cancel_event=ev)
        installed = is_model_downloaded(model)
        shared.emit({"type": "done", "installed": installed})
        if installed:
            save_selected_model(model)
        return 0 if installed else 1
    except Exception as exc:
        shared.emit({"type": "error", "message": str(exc)})
        return 1
    finally:
        with _CANCEL_LOCK:
            _CANCEL.pop(model, None)

def cmd_cancel(args: dict) -> None:
    model = str(args.get("model", ""))
    ev = _CANCEL.get(model)
    if ev is not None:
        ev.set()
    shared.emit({"ok": True})

def cmd_select(args: dict) -> None:
    model = str(args.get("model", ""))
    if not is_model_downloaded(model):
        shared.emit({"ok": False, "error": "Modelo não está instalado."})
        return
    save_selected_model(model)
    shared.emit({"ok": True, "selected": load_selected_model()})

def cmd_set_language(args: dict) -> int:
    """Persist the transcription language (the default for every transcription)."""
    lang = str(args.get("language", "auto"))
    try:
        saved = save_transcript_language(lang)
    except ValueError as exc:
        shared.emit({"ok": False, "error": str(exc)})
        return 1
    shared.emit({"ok": True, "language": saved})
    return 0

def cmd_delete(args: dict) -> None:
    model = str(args.get("model", ""))
    try:
        shutil.rmtree(model_cache_dir(model), ignore_errors=True)
    except Exception:
        pass
    shared.emit({"ok": True})

def cmd_open_finder(args: dict) -> None:
    target = str(args.get("path") or model_cache_dir(str(args.get("model", ""))))
    try:
        subprocess.Popen(["open", target])
    except OSError:
        pass
    shared.emit({"ok": True})

def cmd_set_models_dir(args: dict) -> int:
    try:
        d = save_models_dir(str(args.get("dir", "")))
        shared.emit({"ok": True, "models_dir": d})
        return 0
    except ValueError as exc:
        shared.emit({"ok": False, "error": str(exc)})
        return 1

def _model_names() -> list[str]:
    return [m["internal_name"] for m in load_catalog()]