Terpentinas-EasyLoRAMerger/engine/checkpoint_normalizer.py

"""
checkpoint_normalizer.py — Architecture-Aware Checkpoint Key Normalizer

Detects model architecture from checkpoint tensor keys and normalizes keys
to a canonical form so that checkpoints from different sources (e.g., with
or without the "model.diffusion_model." prefix) can be properly matched
and merged.

Extensible design: add new architectures to CHECKPOINT_ARCHITECTURES dict.
"""

from typing import Dict, List, Tuple, Optional

# ============================================================================
# ARCHITECTURE REGISTRY — Add new models here
# ============================================================================
# Each entry defines:
#   detect:     List of substrings; if ANY key in the checkpoint contains ALL
#               of these (AND-logic per entry), the architecture matches.
#               Multiple entries = OR logic (first matching entry wins).
#   strip_prefixes: List of prefixes to remove from keys to reach canonical form.
#   canonical_form: Description of canonical form for documentation.

CHECKPOINT_ARCHITECTURES = {
    "anima": {
        "detect": [
            # Anima models use net.blocks.* pattern (FLOW-based diffusion)
            ["net.blocks"],
            # Also detect by net.final_layer as secondary marker
            ["net.final_layer"],
            # Detect bare blocks.* keys without net. prefix (common in fine-tuned
            # checkpoints). Uses AND logic: both "blocks." and "adaln_modulation"
            # must exist in the keyset. "blocks." alone would also match SDXL's
            # input_blocks.* / output_blocks.* and Flux's double_blocks.*, so
            # "adaln_modulation" narrows it to Anima-specific key patterns.
            ["blocks.", "adaln_modulation"],
        ],
        # Strip "net." prefix so that net.blocks.* and bare blocks.* keys
        # normalize to the same canonical form for matching.
        "strip_prefixes": ["net."],
        "canonical_form": "bare blocks. prefix — keys like blocks.0.adaln_modulation_cross_attn.1.weight",
        "description": "Anima / Anime Diffusion (FLOW-based model with blocks architecture)",
    },
    "flux": {
        "detect": ["double_blocks", "single_blocks"],
        "strip_prefixes": ["model.diffusion_model."],
        "canonical_form": "no_prefix — keys like double_blocks.0.img_attn.proj.weight",
        "description": "FLUX / FLUX.1-dev / FLUX.1-schnell (Black Forest Labs)",
    },
    "lumina2": {
        "detect": [
            # Primary: cap_embedder is unique to NextDiT / Lumina2
            ["cap_embedder"],
        ],
        "strip_prefixes": ["model.diffusion_model."],
        "canonical_form": "no prefix — keys like layers.0.attention.to_q.weight, cap_embedder.0.weight",
        "description": "Lumina2 / Lumina Image 2.0 / NextDiT (Alpha-VLLM, base for Z-Image Turbo)",
    },
    "z_image": {
        "detect": [
            # Primary: Z-Image has 'layers' patterns
            ["layers", "attention"],
            # Fallback: explicit z_image marker
            ["z_image"],
        ],
        "strip_prefixes": ["model.diffusion_model."],
        "canonical_form": "no_prefix — keys like diffusion_model.layers.0.attention.to_q.weight",
        "description": "Z-Image (Lightrick / Z-inversion models)",
    },
    "sdxl": {
        "detect": [
            # SDXL has input_blocks / middle_block / output_blocks WITH model.diffusion_model. prefix
            ["model.diffusion_model", "input_blocks"],
            ["model.diffusion_model", "middle_block"],
            ["model.diffusion_model", "output_blocks"],
            # Catch-all: any key with model.diffusion_model. prefix is SDXL
            # (FLUX and Z-Image are checked first, so they won't reach here)
            ["model.diffusion_model"],
        ],
        "strip_prefixes": [],
        "canonical_form": "with model.diffusion_model. prefix — keys like model.diffusion_model.input_blocks.0.0.weight",
        "description": "Stable Diffusion XL",
    },
    "sd15": {
        "detect": [
            # SD1.5 has input_blocks WITHOUT model.diffusion_model prefix (but has diffusion_model)
            # Only match if NOT sdxl (checked by detect order)
            ["diffusion_model", "input_blocks"],
            ["diffusion_model", "middle_block"],
            ["diffusion_model", "output_blocks"],
        ],
        "strip_prefixes": [],
        "canonical_form": "with diffusion_model. prefix — keys like diffusion_model.input_blocks.0.0.weight",
        "description": "Stable Diffusion 1.5",
    },
}


def _matches_architecture(keys: List[str], arch_config: dict) -> bool:
    """
    Check if any key pattern in arch_config['detect'] matches the given keys.
    Each entry in 'detect' is EITHER a string (all keys are checked for that
    substring) OR a list of strings (all substrings must be present, AND-logic).
    Multiple entries = OR logic.
    """
    detect_patterns = arch_config.get("detect", [])
    for pattern in detect_patterns:
        if isinstance(pattern, str):
            # Single substring — any key containing it
            if any(pattern in k for k in keys):
                return True
        elif isinstance(pattern, list):
            # List of substrings — ALL must be present in at least one key each
            # (but not necessarily the same key)
            if all(any(sub in k for k in keys) for sub in pattern):
                return True
    return False


def detect_checkpoint_architecture(keys: List[str]) -> str:
    """
    Detect the model architecture from a list of tensor keys.

    Checks architectures in order of specificity (most specific first):
        flux → lumina2 → z_image → baked_sd15 → sdxl → sd15 → unknown

    The ``baked_sd15`` pre-check catches SD1.5 checkpoints whose UNet keys
    have been prefixed with ``model.diffusion_model.`` (e.g. by the baking
    pipeline).  Without this, such checkpoints incorrectly match the SDXL
    catch-all detector (``["model.diffusion_model"]``) before the SD15
    detector can fire.

    Args:
        keys: List of tensor key strings from the checkpoint header.

    Returns:
        Architecture name string: "flux", "lumina2", "z_image", "sdxl",
        "sd15", or "unknown".
    """
    # Ordered by specificity: Anima (net.blocks is unique) first,
    # then FLUX (double_blocks are unique),
    # then Lumina2 (cap_embedder is unique to NextDiT),
    # then mixed patterns, then SD variants.
    check_order = ["anima", "flux", "lumina2", "z_image", "sdxl", "sd15"]

    # ── Pre-check: baked SD1.5 ─────────────────────────────────────────
    # Baked SD1.5 checkpoints have model.diffusion_model.input_blocks.xxx
    # (baker adds "model." prefix to UNet keys) which would match the SDXL
    # catch-all below.  Detect them by looking for SD1.5-specific CLIP keys.
    # "cond_stage_model.transformer.text_model" is exclusive to SD1.5 —
    # SDXL uses "model.conditioner.embedders.0.transformer.text_model" or
    # "model.text_model".
    if _matches_architecture(keys, {"detect": [
        ["cond_stage_model.transformer.text_model", "model.diffusion_model.input_blocks"],
    ]}):
        return "sd15"

    for arch_name in check_order:
        arch_config = CHECKPOINT_ARCHITECTURES.get(arch_name, {})
        if _matches_architecture(keys, arch_config):
            return arch_name

    return "unknown"


def normalize_checkpoint_key(key: str, architecture: str) -> str:
    """
    Normalize a single checkpoint key to canonical form for the given architecture.

    For FLUX / Z-Image: strips 'model.diffusion_model.' prefix.
    For SDXL / SD1.5: keeps existing prefix as-is.

    Args:
        key: Original tensor key.
        architecture: Architecture name from detect_checkpoint_architecture().

    Returns:
        Normalized key in canonical form.
    """
    arch_config = CHECKPOINT_ARCHITECTURES.get(architecture, {})
    strip_prefixes = arch_config.get("strip_prefixes", [])

    normalized = key
    for prefix in strip_prefixes:
        if normalized.startswith(prefix):
            normalized = normalized[len(prefix):]
            break  # only strip first matching prefix

    return normalized


def normalize_checkpoint_header(
    header: Dict[str, dict],
    architecture: Optional[str] = None,
) -> Tuple[Dict[str, dict], Dict[str, str]]:
    """
    Normalize all keys in a checkpoint header to canonical form.

    Args:
        header: Dict mapping key -> {shape, dtype, ...} as read from safetensors.
        architecture: Detected architecture. If None, auto-detect.

    Returns:
        (normalized_header, key_mapping) where:
            normalized_header: Same values, but with canonical keys.
            key_mapping: Dict mapping canonical_key -> original_key.
    """
    if not header:
        return {}, {}

    if architecture is None:
        architecture = detect_checkpoint_architecture(list(header.keys()))

    normalized = {}
    mapping = {}

    for orig_key, value in header.items():
        norm_key = normalize_checkpoint_key(orig_key, architecture)
        # Handle collisions: keep first occurrence with warning
        if norm_key not in normalized:
            normalized[norm_key] = value
            mapping[norm_key] = orig_key
        else:
            import logging
            logging.warning(
                f"normalize_checkpoint_header: key collision — "
                f"'{orig_key}' and '{mapping[norm_key]}' both normalize to '{norm_key}'. "
                f"Keeping first occurrence."
            )

    return normalized, mapping