fn_registry/python/functions/ml/comfyui_install_3d_model.py

"""Instala un checkpoint Hunyuan3D-2 en la carpeta checkpoints/ de ComfyUI.

ComfyUI 0.26.0 reconstruye mallas 3D con los nodos nativos de Hunyuan3D-2, que
cargan un checkpoint self-contained (DiT de forma + VAE 3D + encoder de imagen en
un solo .safetensors) via ImageOnlyCheckpointLoader. Esta funcion resuelve el repo
de HuggingFace de la variante pedida, REUTILIZA la cache de HF si ya esta bajado
(sin re-descargar), y copia el .safetensors a la carpeta checkpoints/ (la ruta real
que declara extra_model_paths.yaml) con el nombre que espera el loader nativo.

Cascada: (1) si el destino ya existe -> reutiliza; (2) si esta en la cache de HF
-> copia desde la cache; (3) si no -> descarga con huggingface_hub (token de
`pass` si la variante fuera gated).

Impura: lectura de YAML, escritura en disco, posible red (HTTP) y subprocess (pass).
"""
import os
import shutil
import subprocess

# variant -> (repo_id HF, ruta del archivo dentro del repo, nombre destino en checkpoints/)
_VARIANTS = {
    "mini": (
        "tencent/Hunyuan3D-2mini",
        "hunyuan3d-dit-v2-mini/model.fp16.safetensors",
        "3D_hunyuan3d-dit-v2-mini.safetensors",
    ),
    "standard": (
        "tencent/Hunyuan3D-2",
        "hunyuan3d-dit-v2-0/model.fp16.safetensors",
        "3D_hunyuan3d-dit-v2-0.safetensors",
    ),
    "mv": (
        "tencent/Hunyuan3D-2mv",
        "hunyuan3d-dit-v2-mv/model.fp16.safetensors",
        "3D_hunyuan3d-dit-v2-mv.safetensors",
    ),
}

_MIN_BYTES = 1_000_000  # un .safetensors real pesa GBs; descarta restos/HTML.


def _checkpoints_dir(comfyui_dir: str) -> str:
    """Resuelve el directorio real de checkpoints de ComfyUI.

    Lee extra_model_paths.yaml (prefiere la seccion con is_default) para devolver
    `<base_path>/<checkpoints_subdir>`. Si el YAML no existe o no se puede parsear,
    cae a la ruta nativa `<comfyui_dir>/models/checkpoints`.
    """
    base = os.path.expanduser(comfyui_dir)
    native = os.path.join(base, "models", "checkpoints")
    yml = os.path.join(base, "extra_model_paths.yaml")
    if not os.path.isfile(yml):
        return native
    try:
        import yaml
        with open(yml, encoding="utf-8") as fh:
            data = yaml.safe_load(fh) or {}
    except Exception:  # noqa: BLE001 — YAML/PyYAML no disponible: usar nativa.
        return native
    if not isinstance(data, dict):
        return native
    fallback = None
    for section in data.values():
        if not isinstance(section, dict):
            continue
        sub = section.get("checkpoints")
        if not sub:
            continue
        bp = os.path.expanduser(str(section.get("base_path", "")))
        first_line = str(sub).splitlines()[0].strip()
        resolved = os.path.join(bp, first_line)
        if section.get("is_default"):
            return resolved
        if fallback is None:
            fallback = resolved
    return fallback or native


def _find_in_hf_cache(repo_id: str, repo_filename: str) -> str | None:
    """Busca el archivo en la cache local de HuggingFace, sin red.

    Layout: ~/.cache/huggingface/hub/models--<org>--<name>/snapshots/<hash>/...
    Resuelve el symlink al blob real y verifica un tamano minimo. Devuelve la ruta
    real o None.
    """
    org_name = repo_id.replace("/", "--")
    hub = os.path.expanduser("~/.cache/huggingface/hub")
    cache_root = os.path.join(hub, f"models--{org_name}", "snapshots")
    if not os.path.isdir(cache_root):
        return None
    target = os.path.basename(repo_filename)
    for snap in os.listdir(cache_root):
        snap_dir = os.path.join(cache_root, snap)
        if not os.path.isdir(snap_dir):
            continue
        for root, _dirs, files in os.walk(snap_dir):
            if target in files:
                real = os.path.realpath(os.path.join(root, target))
                if os.path.isfile(real) and os.path.getsize(real) >= _MIN_BYTES:
                    return real
    return None


def _pass_hf_token() -> str | None:
    """Lee el token de HuggingFace de `pass API_TOKEN_huggingFace`, o None."""
    try:
        out = subprocess.run(
            ["pass", "show", "API_TOKEN_huggingFace"],
            capture_output=True, text=True, timeout=10,
        )
        if out.returncode == 0:
            tok = out.stdout.splitlines()[0].strip() if out.stdout.strip() else ""
            return tok or None
    except (OSError, subprocess.SubprocessError):
        pass
    return None


def comfyui_install_3d_model(
    variant: str = "mini",
    *,
    hf_token: str | None = None,
    comfyui_dir: str = "~/ComfyUI",
) -> dict:
    """Instala el checkpoint Hunyuan3D-2 de la variante pedida en checkpoints/.

    Args:
        variant: "mini" (≈5 GB VRAM, default), "standard" (dit-v2-0, ≈6 GB) o
            "mv" (multiview). Determina el repo de HF y el nombre destino.
        hf_token: token de HuggingFace si la variante fuera gated. Si None y hace
            falta descargar, se intenta leer de `pass show API_TOKEN_huggingFace`.
            keyword-only.
        comfyui_dir: raiz de la instalacion de ComfyUI (se expande ~). La carpeta
            real de checkpoints se resuelve via extra_model_paths.yaml. keyword-only.

    Returns:
        dict {ok, path, bytes, reused_cache, error}. path = ruta del checkpoint en
        checkpoints/; reused_cache=True si ya estaba instalado o se copio de la
        cache de HF (sin descarga de red). Si falla, ok=False y error explica.
    """
    if variant not in _VARIANTS:
        return {"ok": False, "path": "", "bytes": 0, "reused_cache": False,
                "error": f"variant {variant!r} no valida; usa {sorted(_VARIANTS)}"}
    repo_id, repo_filename, dest_name = _VARIANTS[variant]
    ckpt_dir = _checkpoints_dir(comfyui_dir)
    dest = os.path.join(ckpt_dir, dest_name)

    # 1. Ya instalado.
    if os.path.isfile(dest) and os.path.getsize(dest) >= _MIN_BYTES:
        return {"ok": True, "path": dest, "bytes": os.path.getsize(dest),
                "reused_cache": True, "error": ""}

    # 2. En la cache de HF -> copiar (sin red).
    cached = _find_in_hf_cache(repo_id, repo_filename)
    if cached:
        try:
            os.makedirs(ckpt_dir, exist_ok=True)
            shutil.copy2(cached, dest)
        except OSError as exc:
            return {"ok": False, "path": "", "bytes": 0, "reused_cache": False,
                    "error": f"no se pudo copiar de la cache HF a {dest}: {exc}"}
        return {"ok": True, "path": dest, "bytes": os.path.getsize(dest),
                "reused_cache": True, "error": ""}

    # 3. Descargar con huggingface_hub (lazy; usa su propia cache).
    token = hf_token or _pass_hf_token()
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        return {"ok": False, "path": "", "bytes": 0, "reused_cache": False,
                "error": ("no esta en la cache de HF y huggingface_hub no esta "
                          "instalado en este venv. Instala huggingface_hub o baja "
                          f"el archivo {repo_filename!r} de {repo_id!r} a mano (o con "
                          "comfyui_download_model usando la URL de resolve de HF).")}
    try:
        local = hf_hub_download(repo_id=repo_id, filename=repo_filename, token=token)
        os.makedirs(ckpt_dir, exist_ok=True)
        shutil.copy2(local, dest)
    except Exception as exc:  # noqa: BLE001 — red/auth/gated/escritura.
        return {"ok": False, "path": "", "bytes": 0, "reused_cache": False,
                "error": f"fallo descargando {repo_filename} de {repo_id}: {exc}"}
    return {"ok": True, "path": dest, "bytes": os.path.getsize(dest),
            "reused_cache": False, "error": ""}


if __name__ == "__main__":
    import json

    print(json.dumps(comfyui_install_3d_model("mini"), ensure_ascii=False, indent=2))