fn_registry/python/functions/ml/comfyui_search_civitai_images.py

"""Busca imágenes en Civitai via su API pública GET /api/v1/images.

Hermana de `comfyui_search_civitai_models` pero para imágenes: devuelve cada
resultado con la `url` de su PNG (que suele traer el workflow ComfyUI embebido en
los chunks de texto). Permite filtrar NSFW y ordenar. El token se resuelve desde
`pass civitai/api-token` si no se pasa explícitamente.

Sirve para cosechar de Civitai imágenes con su receta para clonar su calidad y
alimentar la librería de skills (grupo `comfyui-skill`): cada `url` se descarga
con `comfyui_fetch_civitai_image` y se destila con `comfyui_extract_recipe_from_png`
(la receta sale del workflow embebido en el PNG, ver Gotchas).

El endpoint de imágenes de Civitai NO admite búsqueda textual (un `query` crudo
devuelve HTTP 500). Por eso, cuando se pasa `query` sin `model_version_id`, esta
función primero resuelve el query a una versión de modelo con
`comfyui_search_civitai_models` y consulta las imágenes de esa versión (filtro
fiable). Si el query no casa ningún modelo, cae al feed global ordenado
(best-effort, ver Gotchas).

Impura: red (HTTP GET a civitai.com) + subprocess (`pass`). Solo stdlib salvo la
función hermana del registry que compone.
"""
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
if _THIS_DIR not in sys.path:
    sys.path.insert(0, _THIS_DIR)

from comfyui_search_civitai_models import comfyui_search_civitai_models  # noqa: E402

_API = "https://civitai.com/api/v1/images"
_TIMEOUT = 30.0
_RETRIES = 3  # la API de imágenes devuelve 503 transitorios bajo carga


def _resolve_token(token: str | None) -> str | None:
    """Devuelve el token explícito o lo lee de `pass civitai/api-token` (best-effort)."""
    if token:
        return token
    try:
        out = subprocess.run(
            ["pass", "civitai/api-token"],
            capture_output=True, text=True, timeout=10,
        )
        if out.returncode == 0:
            lines = out.stdout.strip().splitlines()
            return lines[0].strip() if lines and lines[0].strip() else None
    except Exception:  # noqa: BLE001 — pass no instalado / entrada inexistente
        pass
    return None


def comfyui_search_civitai_images(
    *,
    query: str | None = None,
    model_version_id: int | None = None,
    nsfw: str = "None",
    sort: str = "Most Reactions",
    limit: int = 20,
    token: str | None = None,
) -> dict:
    """Busca imágenes en Civitai y devuelve resultados normalizados.

    Args:
        query: texto de búsqueda libre. La API de imágenes de Civitai NO tiene un
            parámetro de búsqueda textual fuerte (a diferencia de la de modelos):
            se envía best-effort y puede ser ignorado por el servidor. Para
            filtrar de verdad por un modelo concreto usa `model_version_id`.
            keyword-only.
        model_version_id: filtra las imágenes generadas con una versión de modelo
            concreta (el `version_id` que devuelve `comfyui_search_civitai_models`).
            Es el filtro fiable. None no filtra. keyword-only.
        nsfw: nivel NSFW Civitai: "None" (solo SFW, default), "Soft", "Mature",
            "X", o "true"/"false". keyword-only.
        sort: orden Civitai: "Most Reactions" (default), "Most Comments", "Newest".
            keyword-only.
        limit: número máximo de resultados (1-200). keyword-only.
        token: API token de Civitai (header Authorization Bearer). Si None se
            resuelve de `pass civitai/api-token`. No hardcodear. keyword-only.

    Returns:
        dict {ok, items, count, resolved_model_version_id, error}. items = lista de
        {id, url, width, height, nsfw, nsfw_level, base_model, model_version_ids,
        meta}. `meta` es el dict de generación de Civitai SI lo expone (hoy suele
        venir null, ver Gotchas); la receta real se destila del PNG embebido.
        `resolved_model_version_id` indica a qué versión se resolvió un `query`.
        ok=False con error si la petición falla; 0 resultados devuelve ok=True con
        items=[] (no es error).
    """
    tok = _resolve_token(token)

    # El endpoint /images no admite query textual: resolver query → versión de modelo.
    resolved_vid = model_version_id
    if resolved_vid is None and query:
        mr = comfyui_search_civitai_models(query, types="Checkpoint", limit=1, token=tok)
        if mr.get("ok") and mr.get("items"):
            resolved_vid = mr["items"][0].get("version_id")

    params = [
        ("limit", str(max(1, min(int(limit), 200)))),
        ("sort", sort),
        ("nsfw", str(nsfw)),
    ]
    if resolved_vid is not None:
        params.append(("modelVersionId", str(resolved_vid)))

    url = f"{_API}?{urllib.parse.urlencode(params)}"
    headers = {"User-Agent": "fn-registry/comfyui_search_civitai_images"}
    if tok:
        headers["Authorization"] = f"Bearer {tok}"

    data = None
    last_err = ""
    for attempt in range(_RETRIES):
        try:
            req = urllib.request.Request(url, headers=headers)
            with urllib.request.urlopen(req, timeout=_TIMEOUT) as resp:
                data = json.loads(resp.read())
            break
        except urllib.error.HTTPError as exc:
            body = exc.read().decode(errors="replace")[:300]
            last_err = f"HTTP {exc.code} en /api/v1/images: {body}"
            if exc.code in (502, 503, 429) and attempt < _RETRIES - 1:
                time.sleep(1.5 * (attempt + 1))
                continue
            return {"ok": False, "items": [], "count": 0,
                    "resolved_model_version_id": resolved_vid, "error": last_err}
        except urllib.error.URLError as exc:
            return {"ok": False, "items": [], "count": 0,
                    "resolved_model_version_id": resolved_vid,
                    "error": f"no se pudo conectar a civitai.com: {exc.reason}"}
        except json.JSONDecodeError as exc:
            return {"ok": False, "items": [], "count": 0,
                    "resolved_model_version_id": resolved_vid,
                    "error": f"respuesta no es JSON válido: {exc}"}
    if data is None:
        return {"ok": False, "items": [], "count": 0,
                "resolved_model_version_id": resolved_vid,
                "error": last_err or "sin respuesta tras reintentos"}

    items = []
    for img in data.get("items", []) or []:
        lvl = img.get("nsfwLevel")
        is_nsfw = bool(img.get("nsfw", False))
        if isinstance(lvl, str) and lvl not in ("", "None"):
            is_nsfw = True
        elif isinstance(lvl, (int, float)) and not isinstance(lvl, bool) and lvl > 1:
            is_nsfw = True
        items.append({
            "id": img.get("id"),
            "url": img.get("url"),
            "width": img.get("width"),
            "height": img.get("height"),
            "nsfw": is_nsfw,
            "nsfw_level": lvl,
            "base_model": img.get("baseModel"),
            "model_version_ids": img.get("modelVersionIds") or [],
            "meta": img.get("meta") or {},
        })
    return {"ok": True, "items": items, "count": len(items),
            "resolved_model_version_id": resolved_vid, "error": ""}


if __name__ == "__main__":
    out = comfyui_search_civitai_images(query="cinematic portrait", nsfw="None", limit=5)
    print(out["ok"], out["count"])
    for it in out["items"]:
        meta = it["meta"]
        prompt = (meta.get("prompt") or "")[:60]
        print(f"  img {it['id']} nsfw={it['nsfw']} model={meta.get('Model')!r} prompt={prompt!r}")