fn_registry/python/functions/ml/comfyui_map_a1111_params.py

"""Traduce la metadata de generación de Civitai/A1111 a parámetros de ComfyUI.

La metadata que expone Civitai (y Automatic1111) nombra el sampler y el scheduler
de forma distinta a ComfyUI: "DPM++ 2M Karras" en A1111 es
`sampler_name="dpmpp_2m"` + `scheduler="karras"` en ComfyUI. Esta función pura hace
ese mapeo y normaliza el resto de la receta a las claves que consumen los builders
del registry (`comfyui_build_txt2img_workflow`, etc.): steps, cfg, width, height,
seed, positive, negative.

Además infiere la *familia* del modelo (`sd15` / `sdxl` / `flux` / `unknown`) a
partir del nombre del modelo, el `baseModel`, los recursos y las dimensiones, para
que el pipeline de réplica pueda sustituir el checkpoint original por uno instalado
de la misma familia cuando el exacto no esté disponible. Y extrae los LoRAs tanto
de los `resources` de Civitai como de las etiquetas `<lora:nombre:peso>` embebidas
en el propio prompt (sintaxis A1111).

Función pura: sin red, sin I/O. Solo stdlib (re).
"""
import re

# Mapeo sampler A1111 -> (sampler_name ComfyUI, scheduler por defecto). El scheduler
# real puede venir como sufijo del nombre A1111 ("... Karras") y se detecta aparte.
_SAMPLER_MAP = {
    "euler": ("euler", "normal"),
    "euler a": ("euler_ancestral", "normal"),
    "euler ancestral": ("euler_ancestral", "normal"),
    "lms": ("lms", "normal"),
    "heun": ("heun", "normal"),
    "dpm2": ("dpm_2", "normal"),
    "dpm2 a": ("dpm_2_ancestral", "normal"),
    "dpm fast": ("dpm_fast", "normal"),
    "dpm adaptive": ("dpm_adaptive", "normal"),
    "dpm++ 2s a": ("dpmpp_2s_ancestral", "normal"),
    "dpm++ 2m": ("dpmpp_2m", "normal"),
    "dpm++ sde": ("dpmpp_sde", "normal"),
    "dpm++ 2m sde": ("dpmpp_2m_sde", "normal"),
    "dpm++ 2m sde heun": ("dpmpp_2m_sde", "normal"),
    "dpm++ 3m sde": ("dpmpp_3m_sde", "normal"),
    "ddim": ("ddim", "ddim_uniform"),
    "ddpm": ("ddpm", "normal"),
    "plms": ("euler", "normal"),   # PLMS no existe en ComfyUI -> fallback euler
    "unipc": ("uni_pc", "normal"),
    "lcm": ("lcm", "normal"),
    "restart": ("restart", "normal"),
}

# Sufijos de scheduler que A1111 concatena al nombre del sampler.
_SCHEDULER_SUFFIXES = [
    ("karras", "karras"),
    ("exponential", "exponential"),
    ("sgm uniform", "sgm_uniform"),
    ("simple", "simple"),
    ("ddim uniform", "ddim_uniform"),
    ("beta", "beta"),
]

_DEFAULT_SAMPLER = ("euler", "normal")

# Etiqueta A1111 de LoRA embebida en el prompt: <lora:nombre:peso>.
_LORA_TAG_RE = re.compile(r"<lora:([^:>]+)(?::([0-9.]+))?[^>]*>", re.IGNORECASE)

_SDXL_HINTS = ("xl", "pony", "sdxl", "illustrious", "noob", "animagine", "playground")
_SD15_HINTS = ("sd 1.5", "sd1.5", "sd15", "v1-5", "v1.5", "1.5")
_FLUX_HINTS = ("flux",)


def _map_sampler(raw):
    """Traduce un nombre de sampler A1111 a (sampler_name, scheduler) de ComfyUI."""
    if not raw or not isinstance(raw, str):
        return _DEFAULT_SAMPLER
    name = raw.strip().lower()
    scheduler = None
    for suffix, sched in _SCHEDULER_SUFFIXES:
        if name.endswith(" " + suffix):
            scheduler = sched
            name = name[: -len(suffix)].strip()
            break
    sampler_name, default_sched = _SAMPLER_MAP.get(name, _DEFAULT_SAMPLER)
    return sampler_name, (scheduler or default_sched)


def _num(value, cast):
    """Castea best-effort un valor que puede venir como str/num; None si no se puede."""
    if value is None or isinstance(value, bool):
        return None
    try:
        return cast(value)
    except (TypeError, ValueError):
        try:
            return cast(float(value)) if cast is int else cast(value)
        except (TypeError, ValueError):
            return None


def _dims_from_size(size):
    """Parsea 'WxH' (ej. '832x1216') a (width, height); (None, None) si no procede."""
    if not isinstance(size, str) or "x" not in size.lower():
        return None, None
    try:
        w, h = size.lower().split("x", 1)
        return int(w.strip()), int(h.strip())
    except (ValueError, AttributeError):
        return None, None


def _infer_family(meta, resources, width, height):
    """Infiere 'sd15' | 'sdxl' | 'flux' | 'unknown' de la receta."""
    blob_parts = [
        str(meta.get("Model") or ""),
        str(meta.get("model") or ""),
        str(meta.get("baseModel") or ""),
    ]
    for res in resources or []:
        if isinstance(res, dict):
            blob_parts.append(str(res.get("modelName") or ""))
            blob_parts.append(str(res.get("baseModel") or ""))
    blob = " ".join(blob_parts).lower()

    if any(h in blob for h in _FLUX_HINTS):
        return "flux"
    if any(h in blob for h in _SDXL_HINTS):
        return "sdxl"
    if any(h in blob for h in _SD15_HINTS):
        return "sd15"
    # Sin pistas en los nombres: deducir por la dimensión mayor.
    longest = max(width or 0, height or 0)
    if longest >= 900:
        return "sdxl"
    if longest > 0:
        return "sd15"
    return "unknown"


def _checkpoint_hint(meta, resources):
    """Nombre del checkpoint original (Model de la meta o primer resource checkpoint)."""
    model = meta.get("Model") or meta.get("model")
    if isinstance(model, str) and model.strip():
        return model.strip()
    for res in resources or []:
        if isinstance(res, dict) and str(res.get("modelType", "")).lower() in (
            "checkpoint", "model"
        ):
            nm = res.get("modelName") or res.get("name")
            if nm:
                return str(nm)
    return ""


def _loras(meta, resources):
    """Extrae LoRAs de los resources Civitai y de las etiquetas <lora:..> del prompt."""
    out = []
    seen = set()
    for res in resources or []:
        if isinstance(res, dict) and str(res.get("modelType", "")).lower() == "lora":
            nm = res.get("modelName") or res.get("name")
            if nm and str(nm) not in seen:
                seen.add(str(nm))
                w = res.get("weight")
                weight = w if isinstance(w, (int, float)) and not isinstance(w, bool) else 1.0
                out.append({"name": str(nm), "weight": float(weight), "source": "resource"})
    for m in _LORA_TAG_RE.finditer(str(meta.get("prompt") or "")):
        nm = m.group(1).strip()
        if nm and nm not in seen:
            seen.add(nm)
            weight = float(m.group(2)) if m.group(2) else 1.0
            out.append({"name": nm, "weight": weight, "source": "prompt_tag"})
    return out


def _clean_prompt(prompt):
    """Quita las etiquetas <lora:..> del prompt (ComfyUI las maneja como nodos)."""
    return _LORA_TAG_RE.sub("", str(prompt or "")).strip().strip(",").strip()


def comfyui_map_a1111_params(meta, resources=None):
    """Traduce metadata de generación Civitai/A1111 a parámetros de ComfyUI.

    Args:
        meta: dict de generación estilo A1111/Civitai. Claves reconocidas: `prompt`,
            `negativePrompt`, `Model`/`model`, `baseModel`, `sampler`, `steps`,
            `cfgScale`, `seed`, `Size` ('WxH'), `clipSkip`.
        resources: lista de recursos de Civitai ({modelType, modelName, weight,
            baseModel, ...}) para detectar checkpoint, LoRAs y familia. Opcional.

    Returns:
        dict {sampler_name, scheduler, steps, cfg, width, height, seed, positive,
        negative, family, checkpoint_hint, loras, clip_skip}. Los valores numéricos
        son None cuando la meta no los aporta (el caller pone defaults por familia).
        `family` ∈ {sd15, sdxl, flux, unknown}. `loras` = [{name, weight, source}].
        `positive` viene sin las etiquetas <lora:..> (que pasan a `loras`).
    """
    meta = meta or {}
    resources = resources or []

    sampler_name, scheduler = _map_sampler(meta.get("sampler"))
    width, height = _dims_from_size(meta.get("Size"))
    family = _infer_family(meta, resources, width, height)

    return {
        "sampler_name": sampler_name,
        "scheduler": scheduler,
        "steps": _num(meta.get("steps"), int),
        "cfg": _num(meta.get("cfgScale"), float),
        "width": width,
        "height": height,
        "seed": _num(meta.get("seed"), int),
        "positive": _clean_prompt(meta.get("prompt")),
        "negative": str(meta.get("negativePrompt") or "").strip(),
        "family": family,
        "checkpoint_hint": _checkpoint_hint(meta, resources),
        "loras": _loras(meta, resources),
        "clip_skip": _num(meta.get("clipSkip"), int),
    }


if __name__ == "__main__":
    import json

    demo_meta = {
        "prompt": "cinematic portrait of a knight <lora:detail_tweaker:0.6>, sharp focus",
        "negativePrompt": "blurry, lowres",
        "Model": "juggernautXL_v11",
        "sampler": "DPM++ 2M Karras",
        "steps": 30, "cfgScale": 5.5, "seed": 12345, "Size": "832x1216",
    }
    print(json.dumps(comfyui_map_a1111_params(demo_meta), ensure_ascii=False, indent=2))