fn_registry/python/functions/ml/comfyui_download_model.py

"""Descarga un checkpoint / LoRA / VAE a la carpeta correcta de ComfyUI.

Descarga por HTTP a `<comfyui_dir>/models/<dest_subdir>/<filename>` siguiendo
redirects. Soporta Civitai (`https://civitai.com/api/download/models/<versionId>`,
token opcional via `?token=` y header `Authorization: Bearer`) y HuggingFace (URL
directa de resolve). Antes de aceptar el archivo VALIDA que la respuesta no sea
una pagina HTML de error (Cloudflare, login wall, 404 estilizado) y que, si el
nombre termina en `.safetensors`, tenga una cabecera de safetensors valida. Asi
no deja "modelos" que en realidad son HTML de 2 KB.

Funcion impura: hace red (HTTP GET) y escribe en disco. Solo stdlib.
"""
import json
import os
import struct
import urllib.error
import urllib.parse
import urllib.request

_HTML_SNIFF = (b"<!doctype", b"<html", b"<head", b"<?xml")


def _derive_filename(url: str, content_disposition: str) -> str:
    """Deriva el nombre de archivo del Content-Disposition o, si no, de la URL."""
    if content_disposition:
        # filename="x" | filename=x | filename*=UTF-8''x
        for part in content_disposition.split(";"):
            part = part.strip()
            for key in ("filename*=", "filename="):
                if part.lower().startswith(key):
                    raw = part[len(key):].strip().strip('"')
                    if "''" in raw:  # RFC 5987: UTF-8''<pct-encoded>
                        raw = raw.split("''", 1)[1]
                    name = urllib.parse.unquote(os.path.basename(raw))
                    if name:
                        return name
    name = os.path.basename(urllib.parse.urlparse(url).path)
    return name or "model.bin"


def _is_valid_safetensors(path: str) -> bool:
    """True si el archivo tiene cabecera de safetensors coherente.

    Formato: 8 bytes little-endian con la longitud N del header JSON, seguidos de
    N bytes que empiezan por '{'. Rechaza HTML/errores disfrazados de .safetensors.
    """
    try:
        size = os.path.getsize(path)
        if size < 9:
            return False
        with open(path, "rb") as fh:
            n = struct.unpack("<Q", fh.read(8))[0]
            if n <= 0 or n > size - 8 or n > 100_000_000:
                return False
            return fh.read(1) == b"{"
    except Exception:  # noqa: BLE001 — archivo ilegible = invalido
        return False


def comfyui_download_model(
    url: str,
    dest_subdir: str = "checkpoints",
    *,
    comfyui_dir: str = "~/ComfyUI",
    filename: str | None = None,
    token: str | None = None,
    overwrite: bool = False,
    timeout_s: float = 1800.0,
) -> dict:
    """Descarga un modelo a `<comfyui_dir>/models/<dest_subdir>/<filename>`.

    Args:
        url: URL directa de descarga (Civitai api/download, HuggingFace resolve,
            o cualquier HTTP que sirva el binario).
        dest_subdir: subcarpeta dentro de `models/` (checkpoints, loras, vae,
            controlnet, ...). Default "checkpoints".
        comfyui_dir: raiz de la instalacion de ComfyUI (se expande ~).
        filename: nombre destino del archivo. Si None, se deriva del
            Content-Disposition de la respuesta o del path de la URL.
        token: token de API (Civitai). Se añade como `?token=` y como header
            `Authorization: Bearer <token>`. None lo omite.
        overwrite: si False y el destino ya existe, no descarga y devuelve error.
        timeout_s: timeout de la peticion HTTP en segundos.

    Returns:
        dict {ok: bool, path: str, size_bytes: int, error: str}. ok False si la
        respuesta era HTML de error, si un .safetensors no valida su cabecera, o
        si fallo la red/escritura. En esos casos no deja basura en disco.
    """
    base = os.path.expanduser(comfyui_dir)
    dest_dir = os.path.join(base, "models", dest_subdir)

    req_url = url
    headers = {"User-Agent": "fn-registry/comfyui_download_model"}
    if token:
        sep = "&" if "?" in req_url else "?"
        req_url = f"{req_url}{sep}token={urllib.parse.quote(token)}"
        headers["Authorization"] = f"Bearer {token}"

    req = urllib.request.Request(req_url, headers=headers)
    tmp_path = None
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            content_type = resp.headers.get("Content-Type", "")
            disp = resp.headers.get("Content-Disposition", "")
            name = filename or _derive_filename(resp.geturl(), disp)

            os.makedirs(dest_dir, exist_ok=True)
            final_path = os.path.join(dest_dir, name)
            if os.path.exists(final_path) and not overwrite:
                return {
                    "ok": False,
                    "path": final_path,
                    "size_bytes": os.path.getsize(final_path),
                    "error": f"ya existe (overwrite=False): {final_path}",
                }

            # Rechazo temprano por content-type HTML.
            if "text/html" in content_type.lower():
                return {
                    "ok": False,
                    "path": "",
                    "size_bytes": 0,
                    "error": (
                        f"la respuesta es HTML (Content-Type: {content_type}), "
                        "no un binario de modelo. Revisa la URL/token."
                    ),
                }

            tmp_path = final_path + ".part"
            first = resp.read(512)
            # Sniff de los primeros bytes: HTML aunque el content-type mienta.
            low = first.lower().lstrip()
            if any(low.startswith(sig) for sig in _HTML_SNIFF):
                return {
                    "ok": False,
                    "path": "",
                    "size_bytes": 0,
                    "error": "la respuesta empieza con HTML (pagina de error/login), no un modelo.",
                }

            size = 0
            with open(tmp_path, "wb") as fh:
                fh.write(first)
                size += len(first)
                while True:
                    chunk = resp.read(1024 * 256)
                    if not chunk:
                        break
                    fh.write(chunk)
                    size += len(chunk)
    except urllib.error.HTTPError as exc:
        body = exc.read().decode(errors="replace")[:300]
        _cleanup(tmp_path)
        return {"ok": False, "path": "", "size_bytes": 0,
                "error": f"HTTP {exc.code} en {url}: {body}"}
    except Exception as exc:  # noqa: BLE001 — red/DNS/escritura
        _cleanup(tmp_path)
        return {"ok": False, "path": "", "size_bytes": 0,
                "error": f"fallo descargando {url}: {exc}"}

    # Validacion de tamaño minimo (una pagina de error suele ser < 2 KB).
    if size < 1024:
        _cleanup(tmp_path)
        return {"ok": False, "path": "", "size_bytes": size,
                "error": f"descarga sospechosamente pequeña ({size} bytes); probable error, no un modelo."}

    # Validacion de cabecera safetensors si aplica.
    if name.endswith(".safetensors") and not _is_valid_safetensors(tmp_path):
        _cleanup(tmp_path)
        return {"ok": False, "path": "", "size_bytes": size,
                "error": f"{name} no tiene una cabecera safetensors valida; descarga corrupta o HTML disfrazado."}

    os.replace(tmp_path, final_path)
    return {"ok": True, "path": final_path, "size_bytes": size, "error": ""}


def _cleanup(path: str | None) -> None:
    if path and os.path.exists(path):
        try:
            os.remove(path)
        except OSError:
            pass


if __name__ == "__main__":
    import sys

    out = comfyui_download_model(
        sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:8188/",
        dest_subdir="checkpoints",
        filename="smoke_fake.safetensors",
    )
    print(json.dumps(out, ensure_ascii=False, indent=2))