1420 lines
58 KiB
Python
1420 lines
58 KiB
Python
# klein_normalizer.py - UNIVERSAL NORMALIZER WITH ALL FORMATS (CLEANED)
|
||
import torch
|
||
import re
|
||
from collections import defaultdict
|
||
from typing import Dict, List
|
||
try:
|
||
from ..utils import safe_get_rank, silent_pad_or_truncate
|
||
except ImportError:
|
||
from utils import safe_get_rank, silent_pad_or_truncate
|
||
|
||
# ==================== PERFORMANCE & REPAIR CONFIGURATION ====================
|
||
|
||
ENABLE_AUTO_REPAIR = True
|
||
ENABLE_BATCH_RANK_FIX = True
|
||
DEBUG_KEY_DIFF = False # Enable for debugging conversions
|
||
|
||
_current_mapping_ref = None # optional mapping dict for identity tracking
|
||
|
||
# Cache for conversion results keyed by (format_type, key_fingerprint)
|
||
# Stores the converted state dict to avoid redundant conversions.
|
||
_NORMALIZE_CACHE: dict = {}
|
||
|
||
# Increment this version counter to bust all cached normalization results.
|
||
# Use when normalization logic changes (e.g., new key conversion rules).
|
||
# Version 1 → 2: SD1.5 TE key normalization fix (text_model. prefix handling)
|
||
NORMALIZE_CACHE_VERSION = 2
|
||
|
||
|
||
# ==================== PRE-COMPILED REGEX PATTERNS ====================
|
||
|
||
LAYER_PATTERN = re.compile(r'layers_(\d+)_')
|
||
BLOCK_PATTERN = re.compile(r'(double|single)_blocks_(\d+)_')
|
||
ATTN_PATTERN = re.compile(r'(img|txt)_attn_(proj|qkv)')
|
||
TO_OUT_PATTERN = re.compile(r'to_out_(\d+)')
|
||
MLP_PATTERN = re.compile(r'(img|txt)_mlp_(\d+)')
|
||
FEED_FORWARD_PATTERN = re.compile(r'feed_forward_w(\d)')
|
||
RESBLOCK_PATTERN = re.compile(r'resblocks\.(\d+)_')
|
||
ANIMA_BLOCK_PATTERN = re.compile(r'blocks_(\d+)_')
|
||
|
||
# ==================== UTILITY FUNCTIONS ====================
|
||
|
||
def fast_cleanup_key(key):
|
||
"""Standardizes underscores into dots for Flux/Z-Image compatibility."""
|
||
# 1. Remove common prefixes that shouldn't be there
|
||
key = key.replace("lora_unet_", "diffusion_model.")
|
||
|
||
# 2. Fix the 'attention_to' vs 'attention.to' (Z-Image & Flux)
|
||
key = key.replace("attention_to_", "attention.to_")
|
||
|
||
# 3. Fix the 'img_attn_qkv' and 'txt_attn_proj' (Flux 4B/9B)
|
||
# This handles the mismatch where one side has dots and the other underscores
|
||
replacements = {
|
||
"img_attn_": "img_attn.",
|
||
"txt_attn_": "txt_attn.",
|
||
"img_mlp_": "img_mlp.",
|
||
"txt_mlp_": "txt_mlp.",
|
||
"to_out_": "to_out.",
|
||
"feed_forward_w": "feed_forward.w"
|
||
}
|
||
|
||
for old, new in replacements.items():
|
||
key = key.replace(old, new)
|
||
|
||
# 4. Final safety: remove double dots if they were accidentally created
|
||
while ".." in key:
|
||
key = key.replace("..", ".")
|
||
|
||
return key
|
||
|
||
def debug_key_change(old, new):
|
||
"""Print key transformations when debugging is enabled."""
|
||
if DEBUG_KEY_DIFF and old != new:
|
||
print(f"🔧 {old} → {new}")
|
||
|
||
def safe_insert(converted, key, value, original_key=None):
|
||
"""
|
||
Prevent silent overwriting if two keys collapse into one.
|
||
Keeps first occurrence, warns on collision.
|
||
"""
|
||
if key in converted:
|
||
print("WARNING: Key collision detected!")
|
||
if original_key:
|
||
print(f" Original: {original_key}")
|
||
print(f" Collides with existing: {key}")
|
||
print(" Keeping first occurrence.")
|
||
return
|
||
converted[key] = value
|
||
# Store mapping if global mapping ref exists and original_key provided
|
||
if original_key is not None and _current_mapping_ref is not None:
|
||
_current_mapping_ref[key] = original_key
|
||
|
||
def get_alpha_value(state_dict, key, default=16.0):
|
||
"""Determine alpha value, checking for prefixed metadata from previous merges."""
|
||
|
||
# 1. Check for alpha key inside the state_dict (tensors)
|
||
base_key = (key.replace(".lora_A.weight", "")
|
||
.replace(".lora_B.weight", "")
|
||
.replace(".lora_down.weight", "")
|
||
.replace(".lora_up.weight", "")
|
||
.replace(".lora.down.weight", "")
|
||
.replace(".lora.up.weight", ""))
|
||
|
||
alpha_patterns = [
|
||
f"{base_key}.alpha",
|
||
f"{base_key}_alpha",
|
||
key.replace("lora_A.weight", "alpha"),
|
||
key.replace("lora_B.weight", "alpha"),
|
||
]
|
||
|
||
for ak in alpha_patterns:
|
||
if ak in state_dict:
|
||
val = state_dict[ak]
|
||
return val.item() if isinstance(val, torch.Tensor) else float(val)
|
||
|
||
# 2. METADATA CHECK (3-way check for EasyLoRAMerger styles)
|
||
meta = getattr(state_dict, 'metadata', {})
|
||
if meta:
|
||
# Check standard, then Preserve A prefix, then Preserve B prefix
|
||
for mk in ['ss_network_alpha', 'lora_a_ss_network_alpha', 'lora_b_ss_network_alpha']:
|
||
if mk in meta:
|
||
try:
|
||
return float(meta[mk])
|
||
except:
|
||
continue
|
||
|
||
# 3. FALLBACK: RANK INFERENCE (For 'None' mode merges)
|
||
if key in state_dict:
|
||
tensor = state_dict[key]
|
||
if hasattr(tensor, 'shape') and len(tensor.shape) >= 2:
|
||
# Pick dim 0 for A/down, dim 1 for B/up
|
||
if "lora_A" in key or "lora_down" in key or "lora.down" in key:
|
||
rank = tensor.shape[0]
|
||
else:
|
||
rank = tensor.shape[1]
|
||
return float(rank)
|
||
|
||
return default
|
||
|
||
# ==================== DETECTION & REPAIR ====================
|
||
|
||
def detect_broken_lora(state_dict):
|
||
"""Detect common LoRA corruption issues."""
|
||
issues = []
|
||
|
||
a_keys = [k for k in state_dict if "lora_A" in k]
|
||
b_keys = [k for k in state_dict if "lora_B" in k]
|
||
|
||
# If there are no LoRA keys, this is not a LoRA - skip detection
|
||
if len(a_keys) == 0 and len(b_keys) == 0:
|
||
return issues
|
||
|
||
if len(a_keys) != len(b_keys):
|
||
issues.append("Mismatch between lora_A and lora_B counts")
|
||
|
||
for key, tensor in state_dict.items():
|
||
if hasattr(tensor, "shape"):
|
||
# Skip alpha tensors - they can be 0D (scalars)
|
||
if '.alpha' in key and len(tensor.shape) == 0:
|
||
continue # This is fine!
|
||
# Only check shape for LoRA-related tensors
|
||
if "lora" in key.lower() and len(tensor.shape) < 2 and '.alpha' not in key:
|
||
issues.append(f"Bad tensor shape: {key}")
|
||
|
||
if issues:
|
||
print("WARNING: Broken LoRA detected:")
|
||
for i in issues[:5]:
|
||
print(" -", i)
|
||
|
||
return issues
|
||
|
||
def repair_missing_alpha(state_dict):
|
||
"""Add missing alpha keys."""
|
||
repaired = dict(state_dict)
|
||
added = 0
|
||
|
||
for key in list(state_dict.keys()):
|
||
if "lora_A.weight" in key:
|
||
alpha_key = key.replace(".lora_A.weight", ".alpha")
|
||
if alpha_key not in repaired:
|
||
from .scale_utils import find_alpha_value
|
||
meta = getattr(state_dict, 'metadata', {})
|
||
alpha_val = find_alpha_value(state_dict, key, default=16.0, metadata=meta)
|
||
if alpha_val is None:
|
||
# Rank inference fallback (original behavior)
|
||
alpha_val = 16.0
|
||
repaired[alpha_key] = torch.tensor(alpha_val)
|
||
added += 1
|
||
|
||
if added:
|
||
print(f"🛠Added {added} missing alpha keys")
|
||
|
||
return repaired
|
||
|
||
def batch_fix_lora_ranks(state_dict, target_rank=None):
|
||
"""Fix mismatched LoRA ranks across a batch."""
|
||
if not ENABLE_BATCH_RANK_FIX:
|
||
return state_dict
|
||
|
||
grouped = defaultdict(list)
|
||
modified_count = 0
|
||
|
||
for key, tensor in state_dict.items():
|
||
if "lora_" in key and hasattr(tensor, "shape"):
|
||
grouped[key.split(".lora_")[0]].append((key, tensor))
|
||
|
||
fixed = dict(state_dict)
|
||
|
||
for group, entries in grouped.items():
|
||
ranks = []
|
||
for k, t in entries:
|
||
if len(t.shape) >= 2:
|
||
ranks.append(min(t.shape[:2]))
|
||
|
||
if not ranks:
|
||
continue
|
||
|
||
desired_rank = target_rank or min(ranks)
|
||
|
||
# Check if any tensor needs adjustment
|
||
for key, tensor in entries:
|
||
current_rank = safe_get_rank(tensor, key)
|
||
if current_rank != desired_rank:
|
||
fixed[key] = silent_pad_or_truncate(tensor, desired_rank, key)
|
||
modified_count += 1
|
||
|
||
if modified_count > 0:
|
||
print(f"WARNING: Batch rank fix adjusted {modified_count} tensors")
|
||
else:
|
||
print("* Batch rank fix complete (no changes needed)")
|
||
|
||
return fixed
|
||
|
||
def auto_repair_lora(state_dict):
|
||
"""Automatically detect and repair common LoRA issues."""
|
||
if not ENABLE_AUTO_REPAIR:
|
||
return state_dict
|
||
|
||
# Preserve metadata
|
||
metadata = getattr(state_dict, "metadata", None)
|
||
|
||
issues = detect_broken_lora(state_dict)
|
||
repaired = dict(state_dict) # Convert to dict to avoid metadata issues
|
||
|
||
if issues:
|
||
repaired = repair_missing_alpha(repaired)
|
||
repaired = batch_fix_lora_ranks(repaired)
|
||
|
||
# Restore metadata if it was a safetensors object
|
||
if metadata and hasattr(state_dict, "metadata"):
|
||
# If we need to preserve metadata, we'd need to return to safetensors format
|
||
# This is complex - maybe just log a warning
|
||
print("Note: Metadata preservation requires safetensors object")
|
||
|
||
return repaired
|
||
|
||
def accelerate_lora_load(state_dict):
|
||
"""Run repair pipeline on LoRA load."""
|
||
print("* Accelerating LoRA load...")
|
||
state_dict = auto_repair_lora(state_dict)
|
||
return state_dict
|
||
|
||
# ==================== FORMAT DETECTION ====================
|
||
|
||
def detect_lora_format(state_dict, mapping_ref=None):
|
||
"""Detect which LoRA format we're dealing with."""
|
||
if not state_dict:
|
||
return "unknown"
|
||
|
||
keys = list(state_dict.keys())
|
||
if not keys:
|
||
return "unknown"
|
||
|
||
# Infer separator and naming styles (for universal key normalization)
|
||
if mapping_ref is not None:
|
||
separator_style = infer_separator_style(keys)
|
||
naming_style = infer_naming_style(keys)
|
||
mapping_ref["__separator_style__"] = separator_style
|
||
mapping_ref["__naming_style__"] = naming_style
|
||
|
||
# Check for SDXL Kohya format (with or without text encoder)
|
||
# Check for both original (underscore) and normalized (dot) notation
|
||
has_sdxl_patterns = any("input_blocks_" in k or "middle_block_" in k or "output_blocks_" in k or
|
||
"input_blocks." in k or "middle_block." in k or "output_blocks." in k for k in keys)
|
||
|
||
if has_sdxl_patterns:
|
||
has_lora_a_b = any("lora_A" in k or "lora_B" in k for k in keys)
|
||
|
||
if has_lora_a_b:
|
||
# SDXL-specific patterns
|
||
# Check for SDXL-specific architecture patterns
|
||
# IMPORTANT: Check ALL keys for SDXL architecture, not just lora_A/lora_B keys
|
||
# This fixes issue where wai-A3-V140.safetensors was misdetected as SD1.5
|
||
has_emb_layers = any("emb_layers" in k for k in keys)
|
||
has_proj_in_out = any("proj_in" in k or "proj_out" in k for k in keys)
|
||
has_in_out_layers = any("in_layers" in k or "out_layers" in k for k in keys)
|
||
has_time_embed = any("time_embed" in k for k in keys)
|
||
has_transformer_blocks = any("transformer_blocks" in k for k in keys)
|
||
|
||
has_sdxl_architecture = (
|
||
has_emb_layers or
|
||
has_proj_in_out or
|
||
has_in_out_layers or
|
||
has_time_embed or
|
||
has_transformer_blocks
|
||
)
|
||
|
||
|
||
if has_sdxl_architecture:
|
||
return "sdxl_kohya"
|
||
else:
|
||
return "sd15_kohya"
|
||
|
||
# Count different key patterns
|
||
lora_down_count = sum(1 for k in keys if "lora_down" in k and "weight" in k)
|
||
lora_up_count = sum(1 for k in keys if "lora_up" in k and "weight" in k)
|
||
lora_a_count = sum(1 for k in keys if "lora_A" in k and "weight" in k)
|
||
lora_b_count = sum(1 for k in keys if "lora_B" in k and "weight" in k)
|
||
|
||
# Check for SDXL with Text Encoder (must come before SDXL standard check)
|
||
# Check for both original (lora_te) and normalized (te1., te.) TE patterns
|
||
has_te_original = any("lora_te" in k for k in keys)
|
||
has_te_normalized = any("te1." in k or "te." in k for k in keys)
|
||
has_te = has_te_original or has_te_normalized
|
||
|
||
if has_te:
|
||
# Check if it has SDXL patterns (original or normalized)
|
||
has_sdxl_original = any("lora_unet_" in k for k in keys) and \
|
||
any("input_blocks_" in k or "middle_block_" in k or "output_blocks_" in k for k in keys)
|
||
has_sdxl_normalized = any("input_blocks." in k or "middle_block." in k or "output_blocks." in k for k in keys)
|
||
has_sdxl = has_sdxl_original or has_sdxl_normalized
|
||
|
||
if has_sdxl:
|
||
# This is SDXL with TE, not Pony
|
||
# Check for both original (lora_down/lora_up) and normalized (lora_A/lora_B) formats
|
||
has_down_up = (lora_down_count > 0 and lora_up_count > 0)
|
||
has_a_b = (lora_a_count > 0 and lora_b_count > 0)
|
||
|
||
if has_down_up:
|
||
return "sdxl_with_te_lora_down_up"
|
||
elif has_a_b:
|
||
return "sdxl_with_te_lora_a_b"
|
||
else:
|
||
# Fallback based on key count
|
||
return "sdxl_with_te_lora_down_up" if lora_down_count > lora_a_count else "sdxl_with_te_lora_a_b"
|
||
|
||
# Check for SDXL standard format (lora_unet_ prefix with lora_down/lora_up)
|
||
if any("lora_unet_" in k for k in keys):
|
||
if any("input_blocks_" in k or "middle_block_" in k or "output_blocks_" in k for k in keys):
|
||
if any("lora_down" in k or "lora_up" in k for k in keys):
|
||
# Check if it's a small/partial LoRA (Illustrious style)
|
||
if len(keys) < 50:
|
||
return "illustrious_sdxl"
|
||
return "sdxl_standard"
|
||
|
||
# Check for Musubi Text Encoder format (needs conversion)
|
||
if any("lora_te_res_blocks" in k for k in keys):
|
||
if any("diffusion_model" in k for k in keys) or any("layers" in k for k in keys):
|
||
return "musubi_zimage_te_needs_conversion"
|
||
else:
|
||
return "musubi_te_needs_conversion"
|
||
|
||
# Check for Musubi format (needs conversion)
|
||
if any("lora_unet_double_blocks" in k for k in keys):
|
||
return "musubi_flux_needs_conversion"
|
||
|
||
if any("lora_unet_layers" in k for k in keys):
|
||
return "musubi_zimage_needs_conversion"
|
||
|
||
# Check if partially converted
|
||
if any("diffusion_model" in k for k in keys) and any("lora_unet" in k for k in keys):
|
||
return "musubi_partially_converted"
|
||
|
||
# Check for Anima format (lora_unet_blocks_ without double/single prefix)
|
||
if any("lora_unet_blocks_" in k for k in keys):
|
||
if not any("double_blocks" in k or "single_blocks" in k for k in keys):
|
||
return "anima_needs_conversion"
|
||
|
||
# =====================================================================
|
||
# 🔥 SD1.5 Diffusers format detection
|
||
# Detects LoRAs using HuggingFace Diffusers naming convention:
|
||
# down_blocks_X / up_blocks_X / mid_block_X
|
||
# instead of ComfyUI-native:
|
||
# input_blocks_X / output_blocks_X / middle_block_X
|
||
#
|
||
# Must come BEFORE the generic musubi catch-all (line 375) because
|
||
# these keys also contain "lora_unet_" and would be misrouted.
|
||
#
|
||
# Examples:
|
||
# lora_unet_down_blocks_0_attentions_1_transformer_blocks_0_attn2_to_out.0
|
||
# lora_unet_mid_block_attentions_0_transformer_blocks_0_attn1_to_q
|
||
# lora_unet_up_blocks_2_attentions_0_transformer_blocks_0_attn2_to_v
|
||
#
|
||
# Excludes Flux (double_blocks/single_blocks — different structure)
|
||
# and Anima (blocks_N_ without down/up prefix — already caught above).
|
||
# =====================================================================
|
||
if any("lora_unet_" in k for k in keys):
|
||
has_diffusers_sd15 = any(
|
||
"_down_blocks_" in k or "_up_blocks_" in k or "_mid_block_" in k
|
||
for k in keys
|
||
)
|
||
has_flux_patterns = any("double_blocks" in k or "single_blocks" in k for k in keys)
|
||
|
||
if has_diffusers_sd15 and not has_flux_patterns:
|
||
return "sd15_diffusers_needs_conversion"
|
||
|
||
if any("lora_unet_" in k for k in keys):
|
||
return "musubi_other_needs_conversion"
|
||
|
||
# Check for Pony Diffusion (only if not SDXL with TE)
|
||
if any("lora_te" in k for k in keys):
|
||
if lora_down_count > 0 and lora_up_count > 0:
|
||
return "pony_diffusion_lora_down_up"
|
||
elif lora_a_count > 0 and lora_b_count > 0:
|
||
return "pony_diffusion_lora_a_b"
|
||
|
||
# Check for Flux Klein
|
||
if any("diffusion_model" in k for k in keys):
|
||
if any("double_blocks" in k for k in keys):
|
||
if lora_down_count > 0 and lora_up_count > 0:
|
||
return "flux_klein_lora_down_up"
|
||
elif lora_a_count > 0 and lora_b_count > 0:
|
||
return "flux_klein_lora_a_b"
|
||
|
||
# Check for Z-Image format
|
||
if any("diffusion_model.layers" in k for k in keys):
|
||
if lora_a_count > 0 and lora_b_count > 0:
|
||
return "z_image_lora_a_b"
|
||
elif lora_down_count > 0 and lora_up_count > 0:
|
||
return "z_image_lora_down_up"
|
||
|
||
# Check for SDXL Diffusers format
|
||
if any("lora.down" in k or "lora.up" in k for k in keys):
|
||
return "sdxl_diffusers"
|
||
|
||
# Check for Z-Image variants
|
||
if any("diffusion_model.layers." in k for k in keys):
|
||
if any("adaLN_modulation" in k for k in keys):
|
||
return "z_image_ai_toolkit"
|
||
else:
|
||
return "z_image_other"
|
||
|
||
if any(".lora_A.default.weight" in k or ".lora_B.default.weight" in k for k in keys):
|
||
return "z_image_default"
|
||
|
||
# Check for Flux Klein variants
|
||
if any("diffusion_model.double_blocks." in k for k in keys):
|
||
# Check for 9B by tensor shape
|
||
for k in keys[:5]:
|
||
if "lora_A" in k and state_dict[k].shape[1] == 4096:
|
||
return "flux_klein_9b"
|
||
return "flux_klein_4b"
|
||
|
||
if any("diffusion_model.single_blocks." in k for k in keys):
|
||
return "flux_klein_4b"
|
||
|
||
if any("diffusion_model." in k and ("double_blocks" in k or "single_blocks" in k) for k in keys):
|
||
return "flux_klein_4b"
|
||
|
||
# SD1.5 Kohya
|
||
if any("lora_A" in k or "lora_B" in k for k in keys):
|
||
return "sd15_kohya"
|
||
|
||
# LyCORIS
|
||
if any("lora." in k and ("alpha" in k or "dyn" in k) for k in keys):
|
||
return "lycoris"
|
||
|
||
return "unknown"
|
||
|
||
# ==================== INFERENCE FUNCTIONS ====================
|
||
|
||
def infer_separator_style(keys):
|
||
"""
|
||
Infer separator style from a list of keys.
|
||
Returns "dot", "underscore", or "hybrid".
|
||
"""
|
||
dot_pattern = re.compile(r'\.\d+\.') # .0.
|
||
hybrid_pattern = re.compile(r'_\d+\.') # _0.
|
||
underscore_pattern = re.compile(r'_\d+_') # _0_
|
||
|
||
dot_count = 0
|
||
hybrid_count = 0
|
||
underscore_count = 0
|
||
|
||
for key in keys:
|
||
if hybrid_pattern.search(key):
|
||
hybrid_count += 1
|
||
elif dot_pattern.search(key):
|
||
dot_count += 1
|
||
elif underscore_pattern.search(key):
|
||
underscore_count += 1
|
||
|
||
if hybrid_count > 0:
|
||
return "hybrid"
|
||
if dot_count > underscore_count:
|
||
return "dot"
|
||
else:
|
||
return "underscore"
|
||
|
||
|
||
def infer_naming_style(keys):
|
||
"""
|
||
Infer naming style from a list of keys.
|
||
Returns "lora_a_b" or "lora_down_up".
|
||
"""
|
||
lora_a_count = sum(1 for k in keys if "lora_A" in k)
|
||
lora_b_count = sum(1 for k in keys if "lora_B" in k)
|
||
lora_down_count = sum(1 for k in keys if "lora_down" in k)
|
||
lora_up_count = sum(1 for k in keys if "lora_up" in k)
|
||
|
||
if lora_a_count + lora_b_count > lora_down_count + lora_up_count:
|
||
return "lora_a_b"
|
||
else:
|
||
return "lora_down_up"
|
||
|
||
def repair_hybrid_separators(key):
|
||
"""Convert hybrid underscore‑before‑index patterns to pure dot notation."""
|
||
# Pattern: _(\d+)\. (underscore before index, dot after)
|
||
hybrid = re.compile(r'_(\d+)\.')
|
||
# Replace with .\1.
|
||
new_key = hybrid.sub(r'.\1.', key)
|
||
return new_key
|
||
|
||
# ==================== MUSUBI CONVERTERS ====================
|
||
|
||
def convert_musubi_te_to_standard(state_dict):
|
||
"""
|
||
Convert Musubi Text Encoder format to standard ComfyUI format.
|
||
|
||
Musubi TE: lora_te_res_blocks_0_mlp_c_fc.lora_down.weight
|
||
Standard: conditioner.embedders.0.transformer.resblocks.0.mlp.c_fc.lora_A.weight
|
||
"""
|
||
print("🔄 Converting Musubi Text Encoder format...")
|
||
converted = {}
|
||
|
||
for key, value in state_dict.items():
|
||
new_key = key
|
||
|
||
# Convert resblock pattern
|
||
new_key = RESBLOCK_PATTERN.sub(r'resblocks.\1.', new_key)
|
||
|
||
# Convert lora_down/up to lora_A/B
|
||
new_key = (new_key
|
||
.replace("lora_down.weight", "lora_A.weight")
|
||
.replace("lora_up.weight", "lora_B.weight"))
|
||
|
||
new_key = fast_cleanup_key(new_key)
|
||
|
||
debug_key_change(key, new_key)
|
||
safe_insert(converted, new_key, value, key)
|
||
|
||
print(f"✅ Converted {len(converted)} TE keys")
|
||
return converted
|
||
|
||
def convert_musubi_to_standard(state_dict):
|
||
"""Convert Musubi Flux/Klein format to standard."""
|
||
print("🔄 Converting Musubi format...")
|
||
converted = {}
|
||
|
||
# Check if this is actually Pony format
|
||
is_pony = any("lora_te" in k for k in state_dict.keys())
|
||
|
||
for key, value in state_dict.items():
|
||
new_key = key
|
||
|
||
# PONY BYPASS - don't convert these!
|
||
if is_pony and "lora_te" in key:
|
||
converted[key] = value
|
||
continue
|
||
|
||
# Prefix normalization
|
||
if not new_key.startswith("diffusion_model."):
|
||
new_key = new_key.replace("lora_unet_", "diffusion_model.")
|
||
|
||
# Repair hybrid separators (e.g., _0. -> .0.)
|
||
new_key = repair_hybrid_separators(new_key)
|
||
|
||
# Pattern substitutions
|
||
new_key = LAYER_PATTERN.sub(r'layers.\1.', new_key)
|
||
new_key = BLOCK_PATTERN.sub(r'\1_blocks.\2.', new_key)
|
||
new_key = ATTN_PATTERN.sub(r'\1_attn.\2', new_key)
|
||
new_key = TO_OUT_PATTERN.sub(r'to_out.\1', new_key)
|
||
new_key = MLP_PATTERN.sub(r'\1_mlp.\2', new_key)
|
||
new_key = FEED_FORWARD_PATTERN.sub(r'feed_forward.w\1', new_key)
|
||
|
||
# Attention routing
|
||
new_key = (new_key
|
||
.replace("attention_to_k", "attention.to_k")
|
||
.replace("attention_to_q", "attention.to_q")
|
||
.replace("attention_to_v", "attention.to_v")
|
||
.replace("attention_to_out_0", "attention.to_out.0"))
|
||
|
||
# Index normalization
|
||
new_key = re.sub(r'\.(\d+)_', r'.\1.', new_key)
|
||
new_key = re.sub(r'_(\d+)\.', r'.\1.', new_key)
|
||
|
||
# LoRA naming
|
||
new_key = (new_key
|
||
.replace("lora_down.weight", "lora_A.weight")
|
||
.replace("lora_up.weight", "lora_B.weight"))
|
||
|
||
new_key = fast_cleanup_key(new_key)
|
||
|
||
debug_key_change(key, new_key)
|
||
safe_insert(converted, new_key, value, key)
|
||
|
||
print(f"✅ Converted {len(converted)} keys")
|
||
return converted
|
||
|
||
|
||
def _convert_sd15_diffusers_key(key: str) -> str:
|
||
"""
|
||
Convert a single SD1.5 Diffusers-format key to ComfyUI-native format.
|
||
|
||
This is the core conversion logic extracted as a pure string transformation
|
||
so it can be reused by convert_sd15_diffusers_to_comfyui() and tests.
|
||
|
||
See convert_sd15_diffusers_to_comfyui() for full mapping documentation.
|
||
"""
|
||
new_key = key
|
||
|
||
# Step 1: Convert lora_unet_ prefix → diffusion_model.
|
||
if new_key.startswith("lora_unet_"):
|
||
new_key = "diffusion_model." + new_key[len("lora_unet_"):]
|
||
|
||
# Step 1b: Handle lora_te_ prefix (Text Encoder keys for SD1.5 Diffusers)
|
||
# Raw input: lora_te_text_model_encoder_layers_0_self_attn_k_proj
|
||
# Output: text_model.encoder.layers.0.self_attn.k_proj
|
||
# This normalizes the TE path so the reverse map (Strategy 0) can match
|
||
# against the checkpoint's cond_stage_model.transformer.text_model.* keys.
|
||
if new_key.startswith("lora_te_"):
|
||
new_key = new_key[len("lora_te_"):] # Remove lora_te_ prefix
|
||
# Convert text_model_encoder_layers_N_ → text_model.encoder.layers.N.
|
||
new_key = re.sub(
|
||
r'text_model_encoder_layers_(\d+)_',
|
||
r'text_model.encoder.layers.\1.',
|
||
new_key
|
||
)
|
||
# Convert self_attn_k_proj → self_attn.k_proj
|
||
new_key = re.sub(
|
||
r'self_attn_(q|k|v|out)_proj',
|
||
r'self_attn.\1_proj',
|
||
new_key
|
||
)
|
||
# Convert mlp_fc1 → mlp.fc1, mlp_fc2 → mlp.fc2
|
||
new_key = re.sub(r'mlp_fc(\d)', r'mlp.fc\1', new_key)
|
||
|
||
# Step 2: Convert down_blocks_X → input_blocks.(X+1)
|
||
# (input_blocks.0 = conv_in, so down_blocks.0 → input_blocks.1, etc.)
|
||
def _shift_down_block(m):
|
||
idx = int(m.group(1))
|
||
return f"input_blocks_{idx + 1}_"
|
||
new_key = re.sub(r'down_blocks_(\d+)_', _shift_down_block, new_key)
|
||
|
||
# Step 3: Convert up_blocks_X → output_blocks_X (1:1)
|
||
new_key = re.sub(r'up_blocks_(\d+)_', r'output_blocks_\1_', new_key)
|
||
|
||
# Step 4: Convert mid_block_ → middle_block_ (1:1)
|
||
new_key = new_key.replace("mid_block_", "middle_block_")
|
||
|
||
# Step 5: Convert attentions_N_ → sub-block index 1_ (ComfyUI SD1.5)
|
||
# 🔥 CRITICAL FIX: ComfyUI SD1.5 places attention at sub-block index 1
|
||
# (e.g., input_blocks.1.1.transformer_blocks.0), NOT at the Diffusers
|
||
# attention index N. Previous code used '\1_' (the attention index),
|
||
# which produced sub-block .0 for all attention keys, making them
|
||
# un-matchable via the reverse key map (Strategy 0). All attention
|
||
# deltas then fell through to Energy Redirect, causing shape mismatches
|
||
# and OOM crashes.
|
||
# Reference: SD1.5 sub-block map — .0=resnet, .1=attention
|
||
new_key = re.sub(r'attentions_(\d+)_', r'1_', new_key)
|
||
|
||
# Step 5b: Convert resnets_N_ → sub-block index N_ (ComfyUI SD1.5)
|
||
# ComfyUI SD1.5 places resnet at sub-block index 0 for N=0.
|
||
# Keeping the original index N works for N=0 (most common case).
|
||
# For N>0 (second resnet at sub-block .2), this is approximately correct
|
||
# since LoRAs rarely target both resnet and attention in the same block.
|
||
new_key = re.sub(r'resnets_(\d+)_', r'\1_', new_key)
|
||
|
||
# Step 6: Convert remaining underscore numerics to dot numerics
|
||
# e.g., input_blocks_1_1_ → input_blocks.1.1.
|
||
new_key = re.sub(r'_(\d+)_', r'.\1.', new_key)
|
||
|
||
# Step 6b: Fix missing dots from adjacent underscore-numeric patterns
|
||
# After step 5 strips "attentions_N_" → "N_", we get patterns like
|
||
# "1_1_transformer" where the middle "1_" isn't caught by step 6
|
||
# because its leading underscore was consumed by the previous match.
|
||
# This converts "1_transformer" → "1.transformer".
|
||
# Also handles "1_resnets" → "1.resnets" for ResNet submodules.
|
||
new_key = re.sub(r'(\d+)_([a-zA-Z])', r'\1.\2', new_key)
|
||
|
||
# Step 7: Convert attention projection names (underscore → dot)
|
||
# attn1_to_q → attn1.to.q, attn2_to_out_0 → attn2.to.out.0
|
||
new_key = re.sub(r'attn(\d)_to_q\b', r'attn\1.to_q', new_key)
|
||
new_key = re.sub(r'attn(\d)_to_k\b', r'attn\1.to_k', new_key)
|
||
new_key = re.sub(r'attn(\d)_to_v\b', r'attn\1.to_v', new_key)
|
||
new_key = re.sub(r'attn(\d)_to_out_(\d+)', r'attn\1.to_out.\2', new_key)
|
||
new_key = re.sub(r'attn(\d)_to_out\b', r'attn\1.to_out', new_key)
|
||
|
||
# Step 7b: Catch any remaining to_out_N (underscore) → to_out.N (dot)
|
||
# The SD1.5 Diffusers format produces keys like "attn2_to_out.0" where
|
||
# the dot is already present after "to_out", so Step 7's
|
||
# attn(\d)_to_out_(\d+) regex doesn't match (no underscore after "to_out").
|
||
# This cleanup catches any lingering to_out_N patterns from edge cases.
|
||
new_key = re.sub(r'to_out_(\d+)', r'to_out.\1', new_key)
|
||
|
||
# Step 7c: Convert ff_net → ff.net (ComfyUI uses dots for hierarchy)
|
||
# In ComfyUI SD1.5, the feed-forward network module is named "ff.net"
|
||
# (with a dot), not "ff_net" (with underscore). The Diffusers format
|
||
# serializes this as "ff_net" which Step 6's _(\d+)_ regex can't convert
|
||
# because it only targets numeric underscores. Without this fix, keys
|
||
# like "ff_net.0.proj" fail to match the checkpoint's "ff.net.0.proj"
|
||
# in the reverse key map (Strategy 0), falling through to Energy Redirect.
|
||
# Reference: ComfyUI SD1.5 UNet — transformer_blocks.N.ff.net.0.proj
|
||
#
|
||
# 🔥 SUB-FIX: Handle ff_net_N (anime Diffusers second FF network)
|
||
# Some anime Diffusers models have a second FF network (ff_net_2) per
|
||
# transformer block. The key looks like ff_net_2.lora_down.weight — note
|
||
# the DOT after the numeral (not underscore), so Step 6's _(\d+)_ regex
|
||
# cannot convert it (requires trailing underscore). The boundary \b
|
||
# between 't' (word) and '2' (word) in ff_net_2 also prevents
|
||
# \bff_net\b from matching. This regex converts ff_net_N → ff.net.N.
|
||
new_key = re.sub(r'\bff_net_(\d+)\b', r'ff.net.\1', new_key)
|
||
# Secondary: handles ff_net (standard, followed by non-word like . or _)
|
||
new_key = re.sub(r'\bff_net\b', r'ff.net', new_key)
|
||
|
||
# Step 8: Convert lora_down/lora_up → lora_A/lora_B
|
||
new_key = (new_key
|
||
.replace("lora_down.weight", "lora_A.weight")
|
||
.replace("lora_up.weight", "lora_B.weight"))
|
||
|
||
# Step 9: Final cleanup — remove double dots
|
||
while ".." in new_key:
|
||
new_key = new_key.replace("..", ".")
|
||
# Safety: remove ._.
|
||
new_key = new_key.replace("._.", ".")
|
||
|
||
return new_key
|
||
|
||
|
||
def convert_sd15_diffusers_to_comfyui(state_dict):
|
||
"""
|
||
Convert SD1.5 Diffusers-format LoRA to ComfyUI-native format.
|
||
|
||
Diffusers SD1.5 keys use down_blocks/up_blocks/mid_block naming
|
||
(from HuggingFace diffusers library). ComfyUI uses input_blocks/
|
||
output_blocks/middle_block naming.
|
||
|
||
⚠️ KEY COLLISION HANDLING (CRITICAL):
|
||
SD1.5 Diffusers blocks can have MULTIPLE attention sub-modules
|
||
(attentions_0, attentions_1, etc.) that all map to the same ComfyUI
|
||
sub-block index .1. When two (or more) Diffusers keys map to the same
|
||
ComfyUI key, their tensors are SUMMED instead of dropped. This preserves
|
||
the total LoRA signal from all attention modules.
|
||
|
||
Previously, safe_insert() kept only the first occurrence and silently
|
||
dropped the second, causing ~41% key loss and weak baking.
|
||
|
||
Block index mapping (SD1.5 specific):
|
||
down_blocks.X → input_blocks.(X+1) (input_blocks.0 = conv_in)
|
||
up_blocks.X → output_blocks.X (1:1)
|
||
mid_block → middle_block (1:1)
|
||
|
||
Attention submodule mapping (same for all blocks that have attention):
|
||
down_blocks.X.attentions.N.transformer_blocks.M.*
|
||
→ input_blocks.(X+1).1.transformer_blocks.M.*
|
||
|
||
mid_block.attentions.N.transformer_blocks.M.*
|
||
→ middle_block.1.transformer_blocks.M.*
|
||
|
||
up_blocks.X.attentions.N.transformer_blocks.M.*
|
||
→ output_blocks.X.1.transformer_blocks.M.*
|
||
|
||
ResNet submodule mapping:
|
||
down_blocks.X.resnets.N.* → input_blocks.(X+1).0.*
|
||
up_blocks.X.resnets.N.* → output_blocks.X.0.*
|
||
"""
|
||
print(" 🔄 Converting SD1.5 Diffusers format to ComfyUI-native...")
|
||
|
||
# Phase 1: Convert all keys and group by output key
|
||
# Groups collect tensors that map to the same ComfyUI key (collisions).
|
||
groups: Dict[str, List[torch.Tensor]] = {}
|
||
for key, value in state_dict.items():
|
||
new_key = _convert_sd15_diffusers_key(key)
|
||
if new_key not in groups:
|
||
groups[new_key] = []
|
||
groups[new_key].append(value)
|
||
|
||
# Phase 2: Sum colliding groups, preserve unique entries
|
||
converted = {}
|
||
collision_count = 0
|
||
for new_key, tensors in groups.items():
|
||
if len(tensors) > 1:
|
||
# Sum colliding tensors — multiple Diffusers attention modules
|
||
# map to the same ComfyUI sub-block. Summing preserves the total
|
||
# LoRA signal.
|
||
converted[new_key] = sum(tensors[1:], tensors[0])
|
||
collision_count += len(tensors) - 1
|
||
else:
|
||
converted[new_key] = tensors[0]
|
||
|
||
total_original = len(state_dict)
|
||
total_converted = len(converted)
|
||
print(f" ✅ Converted {total_converted} keys to ComfyUI SD1.5 format "
|
||
f"({collision_count} collisions resolved via summation, "
|
||
f"{total_original - total_converted - collision_count} duplicates dropped)")
|
||
return converted
|
||
|
||
|
||
def convert_musubi_to_anima(state_dict):
|
||
"""Convert Musubi Anima format to standard.
|
||
|
||
Anima model uses MiniTrainDIT internally (self.blocks = nn.ModuleList).
|
||
In ComfyUI, BaseModel wraps the unet under self.diffusion_model, so
|
||
runtime state_dict keys are diffusion_model.blocks.0... NOT net.blocks.*.
|
||
The "net." prefix only exists in the trainer checkpoint file, not in the
|
||
model's internal module structure.
|
||
|
||
model_lora_keys_unet() builds key_map from diffusion_model.* keys, so
|
||
the LoRA keys must match: diffusion_model.blocks.0.adaln_modulation...
|
||
(without "net." prefix).
|
||
|
||
Mapping guide (trainer_key → model_module_name):
|
||
cross_attn_output_proj → cross_attn.output_proj (Attention.output_proj)
|
||
self_attn_output_proj → self_attn.output_proj (Attention.output_proj)
|
||
mlp_layer1 → mlp.layer1 (GPT2FeedForward.layer1)
|
||
mlp_layer2 → mlp.layer2 (GPT2FeedForward.layer2)
|
||
cross_attn_q_proj → cross_attn.q_proj (Attention.q_proj)
|
||
cross_attn_k_proj → cross_attn.k_proj (Attention.k_proj)
|
||
cross_attn_v_proj → cross_attn.v_proj (Attention.v_proj)
|
||
|
||
Input: lora_unet_blocks_0_adaln_modulation_cross_attn_1.lora_down.weight
|
||
Output: diffusion_model.blocks.0.adaln_modulation_cross_attn.1.lora_A.weight
|
||
"""
|
||
print("🔄 Converting Anima format...")
|
||
converted = {}
|
||
|
||
for key, value in state_dict.items():
|
||
new_key = key
|
||
|
||
# 0. Handle TE keys (Anima format: lora_te_layers_N_self_attn.X_proj.lora_A.weight)
|
||
# ComfyUI's model_lora_keys_clip() builds key_map from model CLIP state_dict keys.
|
||
# For SD1.5, the CLIP state dict has keys like:
|
||
# transformer.text_model.encoder.layers.N.self_attn.q_proj.weight
|
||
# model_lora_keys_clip() (line 101) creates:
|
||
# key_map["text_encoders.transformer.text_model.encoder.layers.N.self_attn.q_proj"]
|
||
# → "transformer.text_model.encoder.layers.N.self_attn.q_proj.weight"
|
||
# load_lora() then looks for "{key_map_key}.lora_B.weight" in the LoRA dict,
|
||
# so we must produce:
|
||
# text_encoders.transformer.text_model.encoder.layers.N.self_attn.X_proj.lora_A.weight
|
||
#
|
||
# The musubi format lora_te_layers_N_ is shorthand for the model's actual
|
||
# transformer.text_model.encoder.layers.N. structure, so we must produce:
|
||
# transformer.text_model.encoder.layers.N.self_attn.X_proj (NOT te.layers.N.)
|
||
if key.startswith("lora_te_"):
|
||
new_key = key
|
||
|
||
# a. Convert prefix: lora_te_layers → transformer.model.layers
|
||
# Anima TE uses Qwen3 (Llama2_ architecture), whose state dict keys are:
|
||
# qwen3_06b.transformer.model.layers.N.self_attn.q_proj.weight
|
||
# (via SD1ClipModel → SDClipModel → Qwen3_06BModel → Qwen3_06B → Llama2_)
|
||
# NOT SD1.5 CLIP which uses transformer.text_model.encoder.layers.N.*
|
||
# model_lora_keys_clip() strips the qwen3_06b. wrapper via .transformer. detection,
|
||
# producing: text_encoders.transformer.model.layers.N.self_attn.q_proj
|
||
new_key = new_key.replace("lora_te_layers", "transformer.model.layers", 1)
|
||
|
||
# b. Convert layers_N_ → layers.N. (now inside transformer.text_model.encoder.layers.N.)
|
||
new_key = re.sub(r'layers_(\d+)_', r'layers.\1.', new_key)
|
||
|
||
# c. Convert remaining underscore-to-dot for known submodule names
|
||
# o_proj is the short form of out_proj (self_attn_o_proj → self_attn.out_proj)
|
||
new_key = re.sub(
|
||
r'_(q_proj|k_proj|v_proj|output_proj|out_proj|o_proj|gate_proj|up_proj|down_proj|c_fc|c_proj)',
|
||
r'.\1',
|
||
new_key,
|
||
)
|
||
|
||
# d. Convert lora naming (master format: lora_A/lora_B)
|
||
new_key = (new_key
|
||
.replace("lora_down.weight", "lora_A.weight")
|
||
.replace("lora_up.weight", "lora_B.weight"))
|
||
|
||
# e. Final cleanup: remove double dots
|
||
while ".." in new_key:
|
||
new_key = new_key.replace("..", ".")
|
||
|
||
debug_key_change(key, new_key)
|
||
safe_insert(converted, new_key, value, key)
|
||
continue # Skip UNet handling for TE keys
|
||
|
||
# 1. Change prefix: lora_unet_ → diffusion_model.
|
||
# ComfyUI's model_lora_keys_unet() iterates model state_dict keys
|
||
# starting with "diffusion_model." and builds key_map entries.
|
||
# The Anima model (MiniTrainDIT) creates self.blocks internally,
|
||
# so the model's state_dict keys are:
|
||
# diffusion_model.blocks.0.adaln_modulation_cross_attn.1.weight
|
||
# NOT diffusion_model.net.blocks.* — "net." is only in the trainer
|
||
# checkpoint format, not in the model's internal module structure.
|
||
# load_lora() then looks for "{key}.lora_A.weight" matching our output.
|
||
new_key = new_key.replace("lora_unet_", "diffusion_model.", 1)
|
||
|
||
# 2. Convert blocks_N_ → blocks.N.
|
||
new_key = ANIMA_BLOCK_PATTERN.sub(r'blocks.\1.', new_key)
|
||
|
||
# 3. Convert remaining index patterns (_N. → .N. and .N_ → .N.)
|
||
new_key = re.sub(r'_(\d+)\.', r'.\1.', new_key)
|
||
new_key = re.sub(r'\.(\d+)_', r'.\1.', new_key)
|
||
|
||
# =====================================================================
|
||
# 🔥 FIX: Handle llm_adapter and final_layer naming (Anima-specific)
|
||
#
|
||
# The main DIT blocks (lora_unet_blocks_0_...) are handled by step 2
|
||
# (ANIMA_BLOCK_PATTERN). But Anima also has:
|
||
#
|
||
# 1. LLMAdapter (llm_adapter.blocks.N.*) — a separate ModuleList of
|
||
# TransformerBlocks for T5 cross-attention injection.
|
||
# Trainer format: lora_unet_llm_adapter_blocks_0_cross_attn_q_proj
|
||
# Model format: diffusion_model.llm_adapter.blocks.0.cross_attn.q_proj
|
||
#
|
||
# After step 2, llm_adapter_blocks_N_ becomes llm_adapter_blocks.N.
|
||
# (ANIMA_BLOCK_PATTERN matches blocks_0_ and converts to blocks.0.)
|
||
# We then need llm_adapter_blocks → llm_adapter.blocks
|
||
#
|
||
# 2. FinalLayer (net.final_layer.X) — single output projection + adaLN.
|
||
# Trainer format: lora_unet_final_layer_linear
|
||
# Model format: diffusion_model.final_layer.linear
|
||
#
|
||
# After step 1: diffusion_model.final_layer_linear
|
||
# Needs: diffusion_model.final_layer.linear
|
||
# =====================================================================
|
||
|
||
# 3a) llm_adapter_blocks → llm_adapter.blocks
|
||
new_key = new_key.replace('llm_adapter_blocks', 'llm_adapter.blocks')
|
||
|
||
# 3b) final_layer_N → final_layer.N (for numbered sub-parts like adaln_modulation.1)
|
||
new_key = re.sub(r'final_layer_(\d+)', r'final_layer.\1', new_key)
|
||
|
||
# 3c) final_layer_ → final_layer. (for non-digit suffixes like linear)
|
||
new_key = new_key.replace('final_layer_', 'final_layer.')
|
||
|
||
# 4. Map trainer-specific naming to model's actual module names.
|
||
# CRITICAL: Names must match predict2.py's module definitions:
|
||
# GPT2FeedForward uses layer1/layer2 (NOT c_fc/c_proj)
|
||
# Attention uses output_proj (NOT out_proj)
|
||
new_key = (new_key
|
||
.replace("mlp_layer1", "mlp.layer1")
|
||
.replace("mlp_layer2", "mlp.layer2"))
|
||
|
||
# 5. Convert structural underscores to dots for known submodule names.
|
||
# The trainer flattens ALL hierarchy separators to underscores, so
|
||
# 'self_attn.q_proj' becomes 'self_attn_q_proj' in trainer format.
|
||
# The _N. regex above only handles digit-indexed parts; for named
|
||
# submodules we need to convert the preceding underscore to a dot.
|
||
# NOTE: output_proj must appear BEFORE out_proj in the alternation
|
||
# because out_proj is a substring of output_proj.
|
||
new_key = re.sub(
|
||
r'_(q_proj|k_proj|v_proj|output_proj|out_proj|gate_proj|up_proj|down_proj|c_fc|c_proj)',
|
||
r'.\1',
|
||
new_key,
|
||
)
|
||
|
||
# 6. Convert lora naming (master format: lora_A/lora_B)
|
||
new_key = (new_key
|
||
.replace("lora_down.weight", "lora_A.weight")
|
||
.replace("lora_up.weight", "lora_B.weight"))
|
||
|
||
# 7. Final cleanup: remove double dots
|
||
while ".." in new_key:
|
||
new_key = new_key.replace("..", ".")
|
||
|
||
debug_key_change(key, new_key)
|
||
safe_insert(converted, new_key, value, key)
|
||
|
||
print(f"✅ Converted {len(converted)} Anima keys")
|
||
return converted
|
||
|
||
def convert_musubi_zimage_to_standard(state_dict):
|
||
"""Convert Musubi Z-Image format to standard with strict dot-notation."""
|
||
print("🔄 Converting Musubi Z-Image format...")
|
||
converted = {}
|
||
|
||
for key, value in state_dict.items():
|
||
# 1. Primary prefix change
|
||
new_key = key.replace("lora_unet_", "diffusion_model.")
|
||
|
||
# Repair hybrid separators (e.g., _0. -> .0.)
|
||
new_key = repair_hybrid_separators(new_key)
|
||
|
||
# 2. Convert underscores in layers and blocks
|
||
new_key = LAYER_PATTERN.sub(r'layers.\1.', new_key)
|
||
new_key = TO_OUT_PATTERN.sub(r'to_out.\1', new_key)
|
||
new_key = FEED_FORWARD_PATTERN.sub(r'feed_forward.w\1', new_key)
|
||
|
||
# 3. FIX: Specific Z-Image underscore-to-dot mapping
|
||
# This fixes the "attention_to_out" vs "attention.to_out" mismatch
|
||
new_key = new_key.replace("attention_to_", "attention.to_")
|
||
|
||
# 4. Standard LoRA naming conversion
|
||
new_key = (new_key
|
||
.replace("lora_down.weight", "lora_A.weight")
|
||
.replace("lora_up.weight", "lora_B.weight"))
|
||
|
||
new_key = fast_cleanup_key(new_key)
|
||
|
||
debug_key_change(key, new_key)
|
||
safe_insert(converted, new_key, value, key)
|
||
|
||
print(f"✅ Converted {len(converted)} Z-Image keys")
|
||
return converted
|
||
|
||
# ==================== NORMALIZERS ====================
|
||
|
||
def normalize_z_image_context_refiner(state_dict):
|
||
"""Convert Z-Image context_refiner format to standard."""
|
||
normalized = {}
|
||
for key, value in state_dict.items():
|
||
new_key = (key
|
||
.replace("context_refiner.", "diffusion_model.layers.")
|
||
.replace(".default", ""))
|
||
normalized[new_key] = value
|
||
return normalized
|
||
|
||
|
||
def normalize_sdxl_kohya_with_te(state_dict, include_te=True):
|
||
"""Convert SDXL Kohya format with optional TE normalization."""
|
||
normalized = {}
|
||
|
||
for key, value in state_dict.items():
|
||
new_key = key
|
||
|
||
# Handle TE keys
|
||
if include_te and new_key.startswith('lora_te'):
|
||
# SDXL TE format: lora_te1_text_model_encoder_layers_0_mlp_fc1.lora_down.weight
|
||
# Remove lora_ prefix
|
||
# ⚠️ ORDER MATTERS: lora_te2_ must be replaced BEFORE lora_te_ because
|
||
# 'lora_te_' is a SUBSTRING of 'lora_te2_'. If lora_te_ runs first,
|
||
# 'lora_te2_text_model...' becomes 'te.2_text_model...' (corrupted!).
|
||
new_key = new_key.replace('lora_te2_', 'te2.')
|
||
new_key = new_key.replace('lora_te1_', 'te1.')
|
||
new_key = new_key.replace('lora_te_', 'te.')
|
||
|
||
# Convert underscores to dots for TE structure
|
||
# text_model_encoder_layers_0_mlp_fc1 → text_model.encoder.layers.0.mlp.fc1
|
||
new_key = re.sub(r'text_model_encoder_layers_(\d+)_', r'text_model.encoder.layers.\1.', new_key)
|
||
new_key = re.sub(r'mlp_fc(\d)', r'mlp.fc\1', new_key)
|
||
new_key = re.sub(r'self_attn_', r'self_attn.', new_key)
|
||
new_key = re.sub(r'(k|q|v|out)_proj', r'\1_proj', new_key) # k_proj, q_proj etc.
|
||
|
||
# Convert lora_down/lora_up to lora_A/lora_B for consistency
|
||
new_key = new_key.replace('lora_down.', 'lora_A.')
|
||
new_key = new_key.replace('lora_up.', 'lora_B.')
|
||
|
||
# Handle UNet keys
|
||
elif new_key.startswith('lora_unet_'):
|
||
# Remove lora_unet_ prefix if present
|
||
new_key = new_key.replace('lora_unet_', '', 1)
|
||
|
||
# Convert underscores to dots in block patterns
|
||
patterns = [
|
||
(r'input_blocks_(\d+)_(\d+)_', r'input_blocks.\1.\2.'),
|
||
(r'middle_block_(\d+)_', r'middle_block.\1.'),
|
||
(r'output_blocks_(\d+)_(\d+)_', r'output_blocks.\1.\2.'),
|
||
(r'time_embed_', r'time_embed.'),
|
||
]
|
||
|
||
for pattern, replacement in patterns:
|
||
new_key = re.sub(pattern, replacement, new_key)
|
||
|
||
# Handle transformer blocks
|
||
if 'transformer_blocks' in new_key:
|
||
new_key = re.sub(r'transformer_blocks_(\d+)_', r'transformer_blocks.\1.', new_key)
|
||
|
||
# Convert lora_down/lora_up to lora_A/lora_B for consistency
|
||
# Use exact string replacement to catch all variations
|
||
new_key = new_key.replace('lora_down', 'lora_A')
|
||
new_key = new_key.replace('lora_up', 'lora_B')
|
||
|
||
# Ensure proper suffix
|
||
if not new_key.endswith(".weight") and ".alpha" not in new_key and "bias" not in new_key:
|
||
if any(x in new_key for x in ['lora_down', 'lora_up', 'lora_A', 'lora_B']):
|
||
new_key += ".weight"
|
||
|
||
normalized[new_key] = value
|
||
|
||
return normalized
|
||
|
||
def normalize_sdxl_diffusers(state_dict):
|
||
"""Convert SDXL diffusers format to standard."""
|
||
normalized = {}
|
||
|
||
for key, value in state_dict.items():
|
||
new_key = key
|
||
|
||
# Convert lora.down → lora_A, lora.up → lora_B
|
||
new_key = new_key.replace("lora.down", "lora_A").replace("lora.up", "lora_B")
|
||
|
||
# Convert to standard ComfyUI naming
|
||
if "to_k" in new_key or "to_v" in new_key or "to_q" in new_key or "to_out" in new_key:
|
||
new_key = new_key.replace("processor.", "").replace("to_out.0", "to_out")
|
||
|
||
# Add .weight suffix if missing
|
||
if not new_key.endswith(".weight") and ".alpha" not in new_key:
|
||
new_key += ".weight"
|
||
|
||
normalized[new_key] = value
|
||
|
||
return normalized
|
||
|
||
def normalize_z_image_ai_toolkit(state_dict):
|
||
"""Convert Z-Image AI-Toolkit format to standard."""
|
||
normalized = {}
|
||
|
||
for key, value in state_dict.items():
|
||
new_key = key
|
||
|
||
# Already uses lora_A/lora_B format
|
||
if not new_key.endswith(".weight") and ".alpha" not in new_key:
|
||
new_key += ".weight"
|
||
|
||
normalized[new_key] = value
|
||
|
||
return normalized
|
||
|
||
def normalize_z_image_default(state_dict):
|
||
"""Convert Z-Image .default.weight format to standard."""
|
||
normalized = {}
|
||
|
||
for key, value in state_dict.items():
|
||
new_key = key
|
||
|
||
# Remove .default from keys
|
||
new_key = (new_key
|
||
.replace(".lora_A.default", ".lora_A")
|
||
.replace(".lora_B.default", ".lora_B"))
|
||
|
||
# Ensure .weight suffix
|
||
if not new_key.endswith(".weight") and ".alpha" not in new_key:
|
||
new_key += ".weight"
|
||
|
||
normalized[new_key] = value
|
||
|
||
return normalized
|
||
|
||
def normalize_all_klein_formats(state_dict):
|
||
"""Standardizes Musubi and AI-Toolkit names into one format."""
|
||
normalized = {}
|
||
|
||
for key, value in state_dict.items():
|
||
new_key = key
|
||
new_key = repair_hybrid_separators(new_key)
|
||
|
||
# Klein format handling
|
||
new_key = re.sub(r'(img|txt)_attn_(proj|qkv)', r'\1_attn.\2', new_key)
|
||
new_key = re.sub(r'(img|txt)_mlp_', r'\1_mlp.', new_key)
|
||
|
||
if "single_blocks" in new_key:
|
||
new_key = re.sub(r'single_blocks_(\d+)_attn_(proj|qkv)', r'single_blocks.\1.attn.\2', new_key)
|
||
new_key = re.sub(r'single_blocks_(\d+)_mlp_', r'single_blocks.\1.mlp.', new_key)
|
||
|
||
if "lora_unet_" in new_key:
|
||
new_key = new_key.replace("lora_unet_", "diffusion_model.")
|
||
|
||
# Convert lora_down/lora_up to lora_A/lora_B
|
||
new_key = new_key.replace("lora_down", "lora_A").replace("lora_up", "lora_B")
|
||
new_key = new_key.replace("attn_proj", "attn.proj").replace("attn_qkv", "attn.qkv")
|
||
|
||
# Remove double dots
|
||
while ".." in new_key:
|
||
new_key = new_key.replace("..", ".")
|
||
|
||
# Add .weight suffix
|
||
if ("lora_A" in new_key or "lora_B" in new_key) and "weight" not in new_key and "alpha" not in new_key:
|
||
new_key += ".weight"
|
||
|
||
normalized[new_key] = value
|
||
|
||
return normalized
|
||
|
||
def normalize_key_for_flux(key: str) -> str:
|
||
"""Normalize Flux model key to standard format."""
|
||
normalized = key
|
||
patterns = [
|
||
(r'img_attn_proj', 'img_attn.proj'),
|
||
(r'img_attn_qkv', 'img_attn.qkv'),
|
||
(r'txt_attn_proj', 'txt_attn.proj'),
|
||
(r'txt_attn_qkv', 'txt_attn.qkv'),
|
||
(r'mlp_fc1', 'mlp.fc1'),
|
||
(r'mlp_fc2', 'mlp.fc2'),
|
||
]
|
||
|
||
for pattern, replacement in patterns:
|
||
normalized = re.sub(pattern, replacement, normalized)
|
||
|
||
# Remove double dots
|
||
while '..' in normalized:
|
||
normalized = normalized.replace('..', '.')
|
||
|
||
return normalized
|
||
|
||
# ==================== UNIVERSAL NORMALIZER ====================
|
||
|
||
def _compute_key_fingerprint(state_dict):
|
||
"""Compute a hashable fingerprint from sorted state dict keys for cache lookup."""
|
||
return tuple(sorted(state_dict.keys()))
|
||
|
||
|
||
def universal_normalize(state_dict, metadata=None, mapping_ref=None):
|
||
"""Universal normalizer with format preservation and metadata awareness."""
|
||
|
||
# Inject metadata into the object so get_alpha_value can find it
|
||
if metadata is not None:
|
||
# We use setattr because some state_dicts are custom objects
|
||
try:
|
||
state_dict.metadata = metadata
|
||
except:
|
||
# If it's a plain dict, we just attach it as a hidden property
|
||
pass
|
||
|
||
# B2: Skip format re-detection if already stored in mapping_ref
|
||
if mapping_ref is not None and "__format__" in mapping_ref:
|
||
format_type = mapping_ref["__format__"]
|
||
print(f"Detected format (from mapping_ref): {format_type}")
|
||
else:
|
||
format_type = detect_lora_format(state_dict, mapping_ref)
|
||
print(f"Detected format: {format_type}")
|
||
# Store detected format in mapping_ref for later reuse (B2)
|
||
if mapping_ref is not None:
|
||
mapping_ref["__format__"] = format_type
|
||
|
||
# Set global mapping reference if provided
|
||
global _current_mapping_ref
|
||
previous_mapping = _current_mapping_ref
|
||
if mapping_ref is not None:
|
||
_current_mapping_ref = mapping_ref
|
||
else:
|
||
_current_mapping_ref = None
|
||
|
||
|
||
# 1. Run the Auto-Repair/Acceleration pipeline FIRST
|
||
# This ensures Alphas are calculated before keys get renamed
|
||
state_dict = accelerate_lora_load(state_dict)
|
||
|
||
# B1: Check conversion cache — if we've already converted a LoRA with the
|
||
# same key structure (same format + same set of keys), reuse the key mapping
|
||
# to avoid redundant string processing and tensor iteration.
|
||
fingerprint = _compute_key_fingerprint(state_dict)
|
||
cache_key = (format_type, fingerprint, NORMALIZE_CACHE_VERSION)
|
||
if cache_key in _NORMALIZE_CACHE:
|
||
cached_mapping = _NORMALIZE_CACHE[cache_key]
|
||
print(f" [Cache] Reusing cached key mapping for format '{format_type}' ({len(cached_mapping)} keys)")
|
||
result = {}
|
||
for new_key, old_key in cached_mapping.items():
|
||
result[new_key] = state_dict[old_key]
|
||
return result
|
||
|
||
# Bypass standard SD formats
|
||
bypass_formats = ["pony", "illustrious", "sd15", "sd1.5", "sdxl_with_te", "illustrious_sdxl", "sdxl_kohya", "sdxl_standard", "sdxl_diffusers"]
|
||
bypass_match = any(fmt in format_type.lower() for fmt in bypass_formats)
|
||
|
||
# Special handling: SDXL formats should still be normalized (remove from bypass)
|
||
if "sdxl" in format_type.lower():
|
||
bypass_match = False
|
||
|
||
# CRITICAL: "needs_conversion" formats MUST NEVER be bypassed — they require
|
||
# active transformation to produce ComfyUI-native key format. A format like
|
||
# "sd15_diffusers_needs_conversion" would otherwise match "sd15" in the
|
||
# bypass list (substring match) and skip conversion entirely.
|
||
if "needs_conversion" in format_type.lower():
|
||
bypass_match = False
|
||
|
||
if bypass_match:
|
||
print(f"OK: Keeping {format_type} format as-is (Standard SD format)")
|
||
return state_dict
|
||
|
||
# Auto-convert Anima format (must be checked before Musubi catch-all)
|
||
if format_type == "anima_needs_conversion":
|
||
print("Converting Anima format...")
|
||
result = convert_musubi_to_anima(state_dict)
|
||
# Cache the key mapping (new_key → old_key): post-conversion → pre-conversion
|
||
_NORMALIZE_CACHE[cache_key] = {new: old for new, old in zip(result.keys(), state_dict.keys())}
|
||
return result
|
||
|
||
# Auto-convert SD1.5 Diffusers format (must be checked before Musubi catch-all)
|
||
if format_type == "sd15_diffusers_needs_conversion":
|
||
print("Converting SD1.5 Diffusers to ComfyUI format...")
|
||
result = convert_sd15_diffusers_to_comfyui(state_dict)
|
||
# Cache the key mapping (new_key → old_key): post-conversion → pre-conversion
|
||
_NORMALIZE_CACHE[cache_key] = {new: old for new, old in zip(result.keys(), state_dict.keys())}
|
||
return result
|
||
|
||
# Auto-convert Musubi formats
|
||
if "musubi" in format_type.lower() and "needs_conversion" in format_type:
|
||
print("Converting Musubi to standard format...")
|
||
|
||
if "zimage" in format_type:
|
||
result = convert_musubi_zimage_to_standard(state_dict)
|
||
elif "flux" in format_type or "klein" in format_type:
|
||
result = convert_musubi_to_standard(state_dict)
|
||
elif "te" in format_type:
|
||
result = convert_musubi_te_to_standard(state_dict)
|
||
else:
|
||
result = convert_musubi_to_standard(state_dict)
|
||
# Cache the key mapping (new_key → old_key): post-conversion → pre-conversion
|
||
_NORMALIZE_CACHE[cache_key] = {new: old for new, old in zip(result.keys(), state_dict.keys())}
|
||
return result
|
||
|
||
# Already converted Musubi
|
||
if "musubi" in format_type.lower() and "converted" in format_type:
|
||
print(" Musubi already converted, using as-is")
|
||
return state_dict
|
||
|
||
# Keep Pony as-is
|
||
if "pony" in format_type.lower():
|
||
print(" Keeping Pony Diffusion format as-is")
|
||
return state_dict
|
||
|
||
# Apply format-specific normalization
|
||
normalizers = {
|
||
"sdxl_kohya": lambda sd: normalize_sdxl_kohya_with_te(sd, include_te=False),
|
||
"sdxl_standard": lambda sd: normalize_sdxl_kohya_with_te(sd, include_te=False),
|
||
"sdxl_with_te_lora_down_up": lambda sd: normalize_sdxl_kohya_with_te(sd, include_te=True),
|
||
"sdxl_with_te_lora_a_b": lambda sd: normalize_sdxl_kohya_with_te(sd, include_te=True),
|
||
"illustrious_sdxl": lambda sd: normalize_sdxl_kohya_with_te(sd, include_te=False),
|
||
"sdxl_diffusers": normalize_sdxl_diffusers,
|
||
"z_image_ai_toolkit": normalize_z_image_ai_toolkit,
|
||
"z_image_default": normalize_z_image_default,
|
||
"z_image_context_refiner": normalize_z_image_context_refiner,
|
||
"z_image_lora_down_up": normalize_all_klein_formats,
|
||
"z_image_lora_a_b": normalize_all_klein_formats,
|
||
"flux_klein_9b": normalize_all_klein_formats,
|
||
"flux_klein_4b": normalize_all_klein_formats,
|
||
"flux_klein_lora_a_b": normalize_all_klein_formats,
|
||
"flux_klein_lora_down_up": normalize_all_klein_formats,
|
||
"ai_toolkit_klein": normalize_all_klein_formats,
|
||
}
|
||
|
||
if format_type in normalizers:
|
||
normalized = normalizers[format_type](state_dict)
|
||
# Apply Flux key normalization if needed
|
||
if "flux" in format_type or "klein" in format_type:
|
||
final = {}
|
||
for key, value in normalized.items():
|
||
new_key = normalize_key_for_flux(key)
|
||
final[new_key] = value
|
||
# Also cache the key mapping for normalizers
|
||
_NORMALIZE_CACHE[cache_key] = {new: old for new, old in zip(final.keys(), normalized.keys())}
|
||
return final
|
||
return normalized
|
||
|
||
|
||
|
||
# Unknown format - try Klein normalization
|
||
print(f"WARNING: Unknown format, trying Klein normalization")
|
||
try:
|
||
normalized = normalize_all_klein_formats(state_dict)
|
||
final = {}
|
||
for key, value in normalized.items():
|
||
new_key = normalize_key_for_flux(key)
|
||
final[new_key] = value
|
||
return final
|
||
except:
|
||
print(f"WARNING: Could not normalize, using raw format")
|
||
return state_dict
|
||
|
||
# ==================== FORMAT INFO ====================
|
||
|
||
_FORMAT_STYLE_MAP = {
|
||
# SDXL formats (underscore separator)
|
||
"sdxl_standard": ("underscore", "lora_down_up"),
|
||
"sdxl_kohya": ("underscore", "lora_down_up"),
|
||
"sdxl_with_te_lora_down_up": ("underscore", "lora_down_up"),
|
||
"sdxl_with_te_lora_a_b": ("underscore", "lora_a_b"),
|
||
"illustrious_sdxl": ("underscore", "lora_down_up"),
|
||
"sdxl_diffusers": ("underscore", "lora_down_up"),
|
||
# SD1.5 Kohya
|
||
"sd15_kohya": ("underscore", "lora_a_b"),
|
||
# Z-Image formats (dot separator)
|
||
"z_image_lora_a_b": ("dot", "lora_a_b"),
|
||
"z_image_lora_down_up": ("dot", "lora_down_up"),
|
||
"z_image_ai_toolkit": ("dot", "lora_a_b"),
|
||
"z_image_default": ("dot", "lora_a_b"),
|
||
"z_image_context_refiner": ("dot", "lora_a_b"),
|
||
"z_image_other": ("dot", "lora_a_b"),
|
||
# Flux Klein formats (dot separator)
|
||
"flux_klein_lora_a_b": ("dot", "lora_a_b"),
|
||
"flux_klein_lora_down_up": ("dot", "lora_down_up"),
|
||
"flux_klein_9b": ("dot", "lora_a_b"),
|
||
"flux_klein_4b": ("dot", "lora_a_b"),
|
||
"ai_toolkit_klein": ("dot", "lora_a_b"),
|
||
# Musubi formats (dot after conversion)
|
||
"musubi_zimage_needs_conversion": ("dot", "lora_a_b"),
|
||
"musubi_zimage_te_needs_conversion": ("dot", "lora_a_b"),
|
||
"musubi_flux_needs_conversion": ("dot", "lora_a_b"),
|
||
"musubi_te_needs_conversion": ("dot", "lora_a_b"),
|
||
"musubi_partially_converted": ("dot", "lora_a_b"),
|
||
"musubi_other_needs_conversion": ("dot", "lora_a_b"),
|
||
# Anima format (dot separator, lora_down_up naming to match model keys)
|
||
"anima_needs_conversion": ("dot", "lora_down_up"),
|
||
# Pony Diffusion (underscore)
|
||
"pony_diffusion_lora_down_up": ("underscore", "lora_down_up"),
|
||
"pony_diffusion_lora_a_b": ("underscore", "lora_a_b"),
|
||
# LyCORIS (unknown, default underscore)
|
||
"lycoris": ("underscore", "lora_down_up"),
|
||
# Unknown fallback
|
||
"standard_webui": ("underscore", "lora_down_up"),
|
||
"comfy_native": ("dot", "lora_a_b"),
|
||
"forge_optimized": ("underscore", "lora_down_up"),
|
||
"unknown": ("underscore", "lora_down_up"),
|
||
}
|
||
|
||
def get_format_style(format_str: str) -> tuple[str, str]:
|
||
"""
|
||
Return (separator_style, naming_style) for a given format string.
|
||
separator_style: "dot" or "underscore"
|
||
naming_style: "lora_a_b" or "lora_down_up"
|
||
"""
|
||
return _FORMAT_STYLE_MAP.get(format_str, ("underscore", "lora_down_up"))
|
||
|
||
|
||
|
||
|