EricRollei-AAA_Metadata_System/eric_metadata/workflow/generation.py

"""Helpers for deriving generation parameters from ComfyUI workflows.

These utilities focus on identifying sampler configuration, model selection,
latent dimensions, and other generation-related values. They are intentionally
lenient so preparation work can stabilise before we introduce strict schema
expectations.
"""

from __future__ import annotations

import logging
from typing import Any, Dict, Iterable, Optional

from .parsing import AssetDiscovery, WorkflowGraph, WorkflowParsingService

logger = logging.getLogger(__name__)


def extract_generation_parameters(
    workflow: Dict[str, Any],
    *,
    parsing_service: Optional[WorkflowParsingService] = None,
) -> Dict[str, Any]:
    """Return a best-effort dictionary of generation controls.

    The intent is not to be perfect, but to provide stable hooks for upcoming
    refactors and tests. When values cannot be derived they are simply omitted.
    """

    if not isinstance(workflow, dict):
        return {}

    parsing = parsing_service or WorkflowParsingService()
    graph = WorkflowGraph.from_workflow(workflow)
    nodes = graph.nodes

    generation: Dict[str, Any] = {}

    sampler_payload = _extract_sampler_settings(nodes)
    if sampler_payload:
        generation.update(sampler_payload)

    dimensions = _extract_dimensions(nodes)
    if dimensions:
        generation.update(dimensions)

    model_name = _extract_model_name(nodes)
    if model_name:
        generation["model"] = model_name

    vae_name = _extract_first_string(nodes, ("vae", "vae_name"))
    if vae_name:
        generation["vae"] = vae_name

    assets = parsing.discover_assets(workflow)
    lora_entries = _build_lora_entries(assets)
    if lora_entries:
        generation["loras"] = lora_entries

    return generation


def _extract_sampler_settings(nodes: Dict[str, Any]) -> Dict[str, Any]:
    best_candidate: Optional[Dict[str, Any]] = None
    best_score = -1

    for node_id, node in nodes.items():
        class_type = _get_class_type(node)
        lowered = class_type.lower()
        if "sampler" not in lowered:
            continue

        inputs = _extract_inputs(node)
        score = _score_sampler_inputs(inputs)
        if score > best_score:
            best_candidate = inputs
            best_score = score

    if not best_candidate:
        return {}

    payload: Dict[str, Any] = {}

    sampler_name = _first_string(best_candidate, ("sampler", "sampler_name"))
    if sampler_name:
        payload["sampler"] = sampler_name

    scheduler = _first_string(best_candidate, ("scheduler", "scheduler_name"))
    if scheduler:
        payload["scheduler"] = scheduler

    steps = _first_numeric(best_candidate, ("steps", "max_steps", "sampling_steps"), coerce=int)
    if steps is not None:
        payload["steps"] = steps

    cfg = _first_numeric(best_candidate, ("cfg_scale", "cfg", "guidance_scale", "scale"))
    if cfg is not None:
        payload["cfg_scale"] = cfg

    seed = _first_numeric(best_candidate, ("seed", "noise_seed", "latent_seed"), coerce=int)
    if seed is not None:
        payload["seed"] = seed

    denoise = _first_numeric(best_candidate, ("denoise", "denoise_strength", "sigma_start"))
    if denoise is not None:
        payload["denoise"] = denoise

    return payload


def _score_sampler_inputs(inputs: Dict[str, Any]) -> int:
    if not inputs:
        return -1
    score = 0
    for key in ("steps", "cfg", "cfg_scale", "sampler", "sampler_name", "seed"):
        if key in inputs:
            score += 1
    return score


def _extract_dimensions(nodes: Dict[str, Any]) -> Dict[str, Any]:
    for node in nodes.values():
        inputs = _extract_inputs(node)
        width = _first_numeric(inputs, ("width", "w"), coerce=int)
        height = _first_numeric(inputs, ("height", "h"), coerce=int)
        if width is not None and height is not None:
            return {"width": width, "height": height}
    return {}


def _extract_model_name(nodes: Dict[str, Any]) -> Optional[str]:
    preferred_order = (
        "CheckpointLoaderSimple",
        "CheckpointLoader",
        "ModelLoader",
        "DiffusersLoader",
        "unCLIPCheckpointLoader",
    )

    best_name: Optional[str] = None
    best_priority = len(preferred_order)

    for node in nodes.values():
        class_type = _get_class_type(node)
        inputs = _extract_inputs(node)
        candidate = _first_string(
            inputs,
            (
                "ckpt_name",
                "model_name",
                "model",
                "checkpoint",
                "checkpoint_name",
            ),
        )
        if not candidate:
            continue

        try:
            priority = preferred_order.index(class_type)
        except ValueError:
            priority = len(preferred_order)

        if best_name is None or priority < best_priority:
            best_name = candidate
            best_priority = priority

    return best_name


def _extract_first_string(nodes: Dict[str, Any], keys: Iterable[str]) -> Optional[str]:
    for node in nodes.values():
        inputs = _extract_inputs(node)
        candidate = _first_string(inputs, keys)
        if candidate:
            return candidate
    return None


def _build_lora_entries(assets: AssetDiscovery) -> Optional[list]:
    if not assets.loras:
        return None

    entries = []
    for record in assets.loras:
        strength = record.extra.get("strength_model")
        if strength is None:
            strength = record.extra.get("strength")
        strength_value = _coerce_float(strength) or 1.0

        clip_strength = _coerce_float(record.extra.get("clip_strength"))

        payload: Dict[str, Any] = {
            "name": record.name,
            "strength": strength_value,
            "source": record.source,
        }
        if clip_strength is not None:
            payload["clip_strength"] = clip_strength

        entries.append(payload)

    return entries or None


def _get_class_type(node: Dict[str, Any]) -> str:
    value = node.get("class_type") or node.get("type") or ""
    return str(value)


def _extract_inputs(node: Dict[str, Any]) -> Dict[str, Any]:
    raw = node.get("inputs")
    if isinstance(raw, dict):
        return raw

    if isinstance(raw, list):
        mapping: Dict[str, Any] = {}
        for entry in raw:
            if isinstance(entry, dict):
                name = entry.get("name") or entry.get("key")
                if name:
                    mapping[str(name)] = entry.get("value")
        return mapping

    return {}


def _first_string(inputs: Dict[str, Any], keys: Iterable[str]) -> Optional[str]:
    value = _first(inputs, keys)
    if isinstance(value, bytes):
        try:
            return value.decode("utf-8", errors="ignore")
        except Exception:
            return None
    if isinstance(value, str):
        stripped = value.strip()
        return stripped or None
    return None


def _first_numeric(
    inputs: Dict[str, Any],
    keys: Iterable[str],
    *,
    coerce=float,
) -> Optional[float]:
    value = _first(inputs, keys)
    return _coerce_numeric(value, coerce=coerce)


def _first(inputs: Dict[str, Any], keys: Iterable[str]) -> Any:
    for key in keys:
        if key in inputs:
            return inputs[key]
    return None


def _coerce_numeric(value: Any, *, coerce=float) -> Optional[float]:
    if value is None:
        return None
    if isinstance(value, (int, float)):
        try:
            return coerce(value)
        except (TypeError, ValueError):
            return None
    if isinstance(value, str):
        stripped = value.strip()
        if not stripped:
            return None
        try:
            return coerce(float(stripped) if coerce is float else int(float(stripped)))
        except (TypeError, ValueError):
            return None
    return None


def _coerce_float(value: Any) -> Optional[float]:
    return _coerce_numeric(value, coerce=float)