194 lines
7.0 KiB
Python
194 lines
7.0 KiB
Python
"""
|
||
Unified key categorisation and architecture detection.
|
||
|
||
Consolidates key-detection implementations found across the codebase:
|
||
|
||
1. ``categorize_key()`` – :file:`utils.py`
|
||
2. ``categorize_checkpoint_key()`` – :file:`utils.py`
|
||
3. ``_categorize_lora_key()`` – :file:`engine/baking_processor_matching.py`
|
||
4. ``_is_te_key()`` – :file:`engine/serialization_factory.py`
|
||
5. ``detect_checkpoint_architecture`` – :file:`engine/musubi_checkpoint_studio.py` (fallback)
|
||
|
||
Usage
|
||
-----
|
||
>>> from engine.key_utils import categorize_key, is_te_key, is_vae_key, is_clip_key, detect_architecture
|
||
|
||
>>> categorize_key("model.diffusion_model.input_blocks.0.0.weight")
|
||
'unet'
|
||
|
||
>>> categorize_key("te1.text_model.encoder.layers.0.weight")
|
||
'te'
|
||
|
||
>>> categorize_key("first_stage_model.decoder.conv_in.weight")
|
||
'vae'
|
||
|
||
>>> categorize_key("visual.encoder.layers.0.weight")
|
||
'clip'
|
||
|
||
>>> is_te_key("te2.text_model.encoder.layers.0.weight")
|
||
True
|
||
"""
|
||
|
||
from typing import Iterable, Optional, Set
|
||
|
||
import torch
|
||
|
||
|
||
# ===================================================================
|
||
# Pattern collections
|
||
# ===================================================================
|
||
|
||
# ── Text Encoder patterns ──────────────────────────────────────────
|
||
# Checked FIRST to avoid 'encoder' in text_model.encoder.layers
|
||
# matching the VAE 'encoder' pattern.
|
||
_TE_PATTERNS: Set[str] = {
|
||
# From categorize_checkpoint_key (utils.py)
|
||
"te1.", "te.", "text_model", "conditioner",
|
||
"clip_l", "t5", "transformer.text_model",
|
||
"token_embedding", "positional_embedding",
|
||
# From _categorize_lora_key (baking_processor_matching.py) – extra lora_te variants
|
||
"lora_te", "lora_te1_", "lora_te2_",
|
||
"lora_te1_text_model.", "lora_te2_text_model.",
|
||
# From _is_te_key (serialization_factory.py)
|
||
"transformer.model.layers",
|
||
"text_encoder.",
|
||
"clip_g.",
|
||
# From _is_te_key (musubi_checkpoint_studio.py)
|
||
"te2.",
|
||
}
|
||
|
||
# ── VAE patterns ───────────────────────────────────────────────────
|
||
_VAE_PATTERNS: Set[str] = {
|
||
"first_stage_model", ".vae.", "vae.", "vae_", "vae/",
|
||
"decoder", "encoder", "post_quant_conv",
|
||
".conv", ".conv1", ".conv2", # broad – matches after TE check
|
||
# From _is_vae_key (musubi_checkpoint_studio.py)
|
||
"quant_conv.",
|
||
}
|
||
|
||
# ── CLIP visual patterns ───────────────────────────────────────────
|
||
_CLIP_PATTERNS: Set[str] = {
|
||
"visual.", "clip.", "clip_visual", "vision_model",
|
||
"image_encoder", "image_projection",
|
||
# From _is_clip_key (musubi_checkpoint_studio.py)
|
||
"clip.",
|
||
}
|
||
|
||
# ── UNet / diffusion model patterns ────────────────────────────────
|
||
_UNET_PATTERNS: Set[str] = {
|
||
"model.", "model_", "model/", "unet", "diffusion",
|
||
"time_embed", "input_blocks", "middle_block", "output_blocks",
|
||
"out.", "out_", "conv_in", "conv_out", "norm", "attentions",
|
||
"resnets", "down_blocks", "up_blocks",
|
||
"double_blocks", "single_blocks", "img_attn", "txt_attn",
|
||
"img_mlp", "txt_mlp", "feed_forward", "linear1", "linear2",
|
||
"proj", "layers", "attention", "to_q", "to_k", "to_v", "to_out",
|
||
"mlp.fc", "adaln", "modulation", "transformer", "diffusion_model",
|
||
}
|
||
|
||
# ── Architecture-detection patterns ────────────────────────────────
|
||
_ARCHITECTURE_PATTERNS = [
|
||
("Anima", "net.blocks"),
|
||
("Flux", "double_blocks"),
|
||
("Flux", "single_blocks"),
|
||
("Lumina2", "cap_embedder"),
|
||
("Z-Image", "z_image"),
|
||
("Z-Image", "layers"), # heuristic – Z-Image often has 'layers' in attention
|
||
("SDXL", "model.diffusion_model"),
|
||
("SD1.5", "diffusion_model"),
|
||
]
|
||
|
||
|
||
# ===================================================================
|
||
# Public API
|
||
# ===================================================================
|
||
|
||
def categorize_key(key: str) -> str:
|
||
"""
|
||
Categorise a key (LoRA or checkpoint) into a component name.
|
||
|
||
Returns one of ``'unet'``, ``'te'``, ``'clip'``, ``'vae'``, ``'other'``.
|
||
|
||
**Order of checks (important)**
|
||
The function checks TE patterns **before** VAE patterns so that
|
||
``text_model.encoder.layers`` is classified as ``'te'`` rather than
|
||
``'vae'`` (which also checks for ``'encoder'``).
|
||
"""
|
||
key_lower = key.lower()
|
||
|
||
# 1. Text Encoder (checked first – see docstring)
|
||
if any(p in key_lower for p in _TE_PATTERNS):
|
||
return 'te'
|
||
|
||
# 2. VAE
|
||
if any(p in key_lower for p in _VAE_PATTERNS):
|
||
return 'vae'
|
||
|
||
# 3. CLIP visual
|
||
if any(p in key_lower for p in _CLIP_PATTERNS):
|
||
return 'clip'
|
||
|
||
# 4. UNet / diffusion model
|
||
if any(p in key_lower for p in _UNET_PATTERNS):
|
||
return 'unet'
|
||
|
||
# 5. Unrecognised
|
||
return 'other'
|
||
|
||
|
||
def is_te_key(key: str) -> bool:
|
||
"""Return ``True`` if *key* belongs to a text encoder (CLIP/T5)."""
|
||
return categorize_key(key) == 'te'
|
||
|
||
|
||
def is_vae_key(key: str) -> bool:
|
||
"""Return ``True`` if *key* belongs to the VAE."""
|
||
return categorize_key(key) == 'vae'
|
||
|
||
|
||
def is_clip_key(key: str) -> bool:
|
||
"""Return ``True`` if *key* belongs to the CLIP visual encoder."""
|
||
return categorize_key(key) == 'clip'
|
||
|
||
|
||
def is_unet_key(key: str) -> bool:
|
||
"""Return ``True`` if *key* belongs to the UNet / diffusion model."""
|
||
return categorize_key(key) == 'unet'
|
||
|
||
|
||
def detect_architecture(keys: Iterable[str]) -> str:
|
||
"""
|
||
Detect checkpoint architecture from a collection of keys.
|
||
|
||
Returns one of ``'Anima'``, ``'Flux'``, ``'Lumina2'``, ``'Z-Image'``,
|
||
``'SDXL'``, ``'SD1.5'``, ``'Unknown'``.
|
||
|
||
Patterns are checked in order — the first match wins.
|
||
"""
|
||
key_set: Set[str] = set(keys)
|
||
for arch, pattern in _ARCHITECTURE_PATTERNS:
|
||
if any(pattern in k for k in key_set):
|
||
return arch
|
||
return "Unknown"
|
||
|
||
|
||
# ===================================================================
|
||
# Safetensors dtype mappings
|
||
# ===================================================================
|
||
# These were previously duplicated in _SafetensorsStreamWriter.DTYPE_MAP and
|
||
# _LazyCheckpointMapping._SAFETENSORS_DTYPE_MAP inside musubi_checkpoint_studio.py.
|
||
# Extracted here to eliminate fragile cross-module coupling (baker_node.py:62-64).
|
||
|
||
SAFETENSORS_DTYPE_MAP = {
|
||
'F64': torch.float64, 'F32': torch.float32, 'F16': torch.float16,
|
||
'BF16': torch.bfloat16, 'F8_E4M3': torch.float8_e4m3fn,
|
||
'F8_E5M2': torch.float8_e5m2, 'F8': torch.float8_e4m3fn,
|
||
'I64': torch.int64, 'I32': torch.int32, 'I16': torch.int16,
|
||
'I8': torch.int8, 'U8': torch.uint8, 'BOOL': torch.bool,
|
||
}
|
||
|
||
DTYPE_MAP = {k: (v, v.itemsize) for k, v in SAFETENSORS_DTYPE_MAP.items()}
|
||
DTYPE_MAP_REVERSE = {v: k for k, v in SAFETENSORS_DTYPE_MAP.items()}
|
||
# Override: force float8_e4m3fn → valid 'F8_E4M3' (not legacy 'F8' which collides and overwrites)
|
||
DTYPE_MAP_REVERSE[torch.float8_e4m3fn] = 'F8_E4M3'
|