""" Key matching: reverse key map, multi-strategy cascade. Extracted from baking_processor.py. Contains all key matching logic: - ComfyUI-native reverse key map generation - 8 matching strategies (reverse map, exact, prefix, underscore-to-dot, TE2, suffix, shape+block, shape match) - Fused QKV resolution, structural orphan detection - Component categorization and diagnostics """ import torch import re import math from typing import Any, Callable, Dict, List, Optional, Tuple, Set from ..utils import ProgressTracker from .baking_processor_constants import ( UNDERSCORE_TO_DOT_PATTERNS, LORA_TO_CHECKPOINT_PREFIXES, ) from .baking_processor_baking import _check_shape_compatible from .baking_processor_delta import _LazyDeltaDict, _MatchedDeltas, _ShapeInfo from .key_mapper import build_lora_key_variants def _strip_tensor_suffix(key: str) -> str: """Strip .weight, .bias, or .alpha suffix from a tensor key. Shared helper used by _build_reverse_key_map(), _global_search_by_shape_block(), and _find_matching_keys() to normalize checkpoint keys to their base form. """ for suffix in ('.weight', '.bias', '.alpha'): if key.endswith(suffix): return key[:-len(suffix)] return key # =================================================================== # Strategy 2: Prefix-Aware Conversion # =================================================================== def _lora_key_to_checkpoint_key(lora_base: str) -> Optional[str]: """ Convert a LoRA key to a checkpoint key by trying known prefix mappings. Iterates through LORA_TO_CHECKPOINT_PREFIXES and tries each mapping. Returns the converted key or None if no mapping applies. Example: lora_unet_input_blocks.3.1.attn1.to_k → model.diffusion_model.input_blocks.3.1.attn1.to_k """ for lora_prefix, ckpt_prefix, _desc in LORA_TO_CHECKPOINT_PREFIXES: if lora_base.startswith(lora_prefix): rest = lora_base[len(lora_prefix):] return ckpt_prefix + rest return None # =================================================================== # Strategy 2.5a: Underscore-to-Dot Fuzzy Conversion # =================================================================== def _try_underscore_to_dot_conversion( lora_base: str, ckpt_sd: Dict[str, torch.Tensor], # kept for API consistency with other strategies ckpt_base_lookup: Dict[str, str], ) -> Optional[str]: """ Try fuzzy underscore-to-dot conversion patterns. Applies each pattern from UNDERSCORE_TO_DOT_PATTERNS to the LoRA base key. If the converted key exists in the checkpoint base lookup, returns the full checkpoint key. Catches common naming inconsistencies: attn1_to_k → attn1.to.k ff_net_0 → ff.net.0 etc. """ for pattern, replacement, _desc in UNDERSCORE_TO_DOT_PATTERNS: converted = re.sub(pattern, replacement, lora_base) if converted != lora_base: if converted in ckpt_base_lookup: return ckpt_base_lookup[converted] return None # =================================================================== # Strategy 2.5b: SDXL TE Conditioner Prefix (handles both te1. and te2.) # =================================================================== def _try_te_conditioner_prefix( lora_base: str, ckpt_sd: Dict[str, torch.Tensor], ckpt_base_lookup: Dict[str, str], lora_deltas: Dict[str, torch.Tensor], ) -> Optional[str]: """ Try to match SDXL TE1/TE2 keys by testing the conditioner.embedders prefix. SDXL uses conditioner.embedders.{0|1}.transformer.text_model for its text encoders, while Kohya-style LoRAs may use te1. or te2. prefix. Uses shape-based disambiguation to choose between embedders.0 and embedders.1 when both exist. Returns: Full checkpoint key if matched, None otherwise. """ # Handle both te1. and te2. prefix keys if not lora_base.startswith(('te1.', 'te2.')): return None # Determine embedder hint: te1. prefers embedders.0, te2. prefers embedders.1 if lora_base.startswith('te1.'): prefix = 'te1.' embedder_hint = 0 else: prefix = 'te2.' embedder_hint = 1 # Extract the inner key (everything after the prefix) inner_key = lora_base[len(prefix):] # e.g., 'text_model.encoder.layers.0.self_attn.k_proj' # Try hinted embedder first with shape check candidate_hint = f"conditioner.embedders.{embedder_hint}.transformer.{inner_key}" if candidate_hint in ckpt_base_lookup: full_key = ckpt_base_lookup[candidate_hint] if full_key in ckpt_sd: target_delta = lora_deltas.get(lora_base) if target_delta is None or _check_shape_compatible(target_delta, ckpt_sd[full_key]): return full_key else: return full_key # Try the other embedder with shape check other = 1 - embedder_hint candidate_other = f"conditioner.embedders.{other}.transformer.{inner_key}" if candidate_other in ckpt_base_lookup: full_key = ckpt_base_lookup[candidate_other] if full_key in ckpt_sd: target_delta = lora_deltas.get(lora_base) if target_delta is None or _check_shape_compatible(target_delta, ckpt_sd[full_key]): return full_key else: return full_key # Shape-based disambiguation: find which embedder has matching shapes target_delta = lora_deltas.get(lora_base) if target_delta is None: return None for embedder_idx in range(2): # Try embedders.0 and embedders.1 check_key = f"conditioner.embedders.{embedder_idx}.transformer.{inner_key}" full_key = ckpt_base_lookup.get(check_key) if full_key is not None and full_key in ckpt_sd: ckpt_tensor = ckpt_sd[full_key] if isinstance(ckpt_tensor, torch.Tensor): # Check if shapes are compatible if ckpt_tensor.shape == target_delta.shape or \ (ckpt_tensor.dim() == target_delta.dim() and ckpt_tensor.shape[0] == target_delta.shape[0]): return full_key return None # =================================================================== # Strategy 2.5d: Deterministic TE Direct-Mapping (Kohya-inspired) # =================================================================== def _try_deterministic_te_mapping( lora_base: str, ckpt_sd: Dict[str, torch.Tensor], # kept for API consistency with other strategies ckpt_base_lookup: Dict[str, str], delta: torch.Tensor, is_shape_compatible_fn: Callable[[str, torch.Tensor], bool], verbose: bool = False, ) -> Optional[str]: """ Deterministic TE key mapping (Kohya-inspired, Strategy 2.5d). Maps LoRA TE keys to checkpoint keys using known structural patterns rather than relying on the reverse_map or shape-based fallbacks. Mirrors Kohya's approach of knowing the architecture layout a priori. Handles: - SD1.5 TE: lora_te_text_model.encoder.layers.N.* → cond_stage_model.transformer.text_model.encoder.layers.N.* - SDXL TE1 (lora_te1_): lora_te1_text_model.encoder.layers.N.* → cond_stage_model.transformer.text_model.encoder.layers.N.* - SDXL TE1 (normalized te1.): te1.text_model.encoder.layers.N.* → conditioner.embedders.{0|1}.transformer.text_model.encoder.layers.N.* - SDXL TE2 (lora_te2_): lora_te2_text_model.encoder.layers.N.* → conditioner.embedders.{0|1}.transformer.text_model.encoder.layers.N.* Returns: Full checkpoint key if matched, None otherwise. """ # Handle both raw LoRA keys (lora_te_*) and normalized keys (text_model.*, te1.) # After identity_normalize() → convert_sd15_diffusers_to_comfyui(), the # lora_te_ prefix is stripped and keys use text_model.encoder.layers.N.* format. # After normalize_sdxl_kohya_with_te(), keys use te1.text_model.encoder.layers.N.* format. if not lora_base.startswith(('lora_te_', 'lora_te1_', 'lora_te2_', 'te1.', 'te2.', 'text_model.')): return None # Extract the suffix after the prefix # Raw: lora_te_text_model.encoder.layers.0.self_attn.k_proj # → suffix = "text_model.encoder.layers.0.self_attn.k_proj" # Norm: text_model.encoder.layers.0.self_attn.k_proj # → suffix = "encoder.layers.0.self_attn.k_proj" # Post-norm: te1.text_model.encoder.layers.0.self_attn.k_proj # → suffix = "text_model.encoder.layers.0.self_attn.k_proj" for prefix in ('lora_te2_text_model.', 'lora_te1_text_model.', 'lora_te_text_model.', 'text_model.'): if lora_base.startswith(prefix): suffix = lora_base[len(prefix):] break else: # Check for post-normalization te1./te2. prefix (from normalize_sdxl_kohya_with_te) if lora_base.startswith('te1.'): suffix = lora_base[4:] # Strip 'te1.' → 'text_model.encoder.layers.N.*' elif lora_base.startswith('te2.'): suffix = lora_base[4:] # Strip 'te2.' → 'text_model.encoder.layers.N.*' else: return None # --- SDXL TE2 (lora_te2_) --- if lora_base.startswith('lora_te2_'): # Try embedders.0 first, then embedders.1 for embedder_idx in ('0', '1'): ckpt_candidate = ( f"conditioner.embedders.{embedder_idx}." f"transformer.text_model.{suffix}" ) if ckpt_candidate in ckpt_base_lookup: full_key = ckpt_base_lookup[ckpt_candidate] if is_shape_compatible_fn(full_key, delta): if verbose: print(f" 🎯 Deterministic TE2 (embedder {embedder_idx}): " f"{lora_base} → {full_key}") return full_key return None # --- SDXL TE1 (normalized te1. prefix) → conditioner.embedders --- # After normalize_sdxl_kohya_with_te(), the key uses te1. prefix. # In SDXL, both CLIP-L (TE1) and CLIP-G (TE2) live under conditioner.embedders.*. # Try embedders.0 first (primary CLIP-L placement), then embedders.1. if lora_base.startswith('te1.'): for embedder_idx in ('0', '1'): ckpt_candidate = ( f"conditioner.embedders.{embedder_idx}." f"transformer.text_model.{suffix}" ) if ckpt_candidate in ckpt_base_lookup: full_key = ckpt_base_lookup[ckpt_candidate] if is_shape_compatible_fn(full_key, delta): if verbose: print(f" 🎯 Deterministic TE1 (embedder {embedder_idx}): " f"{lora_base} → {full_key}") return full_key return None # --- SDXL TE2 (normalized te2. prefix) → conditioner.embedders --- if lora_base.startswith('te2.'): for embedder_idx in ('0', '1'): ckpt_candidate = ( f"conditioner.embedders.{embedder_idx}." f"transformer.text_model.{suffix}" ) if ckpt_candidate in ckpt_base_lookup: full_key = ckpt_base_lookup[ckpt_candidate] if is_shape_compatible_fn(full_key, delta): if verbose: print(f" 🎯 Deterministic TE2 (embedder {embedder_idx}): " f"{lora_base} → {full_key}") return full_key return None # --- SD1.5 TE / SDXL TE1 (lora_te_ or lora_te1_) --- ckpt_candidate = f"cond_stage_model.transformer.text_model.{suffix}" if ckpt_candidate in ckpt_base_lookup: full_key = ckpt_base_lookup[ckpt_candidate] if is_shape_compatible_fn(full_key, delta): if verbose: print(f" 🎯 Deterministic TE: {lora_base} → {full_key}") return full_key return None # =================================================================== # Strategy 0: Reverse Key Map (ComfyUI-native, exhaustive) # =================================================================== # Delegates to engine.key_mapper.build_lora_key_variants() — the single # authoritative source for ALL checkpoint → LoRA key variant generation. # This replaces the previous ~350-line manual implementation. def _build_reverse_key_map(ckpt_sd: Dict[str, torch.Tensor]) -> Dict[str, str]: """ Build a reverse key map using the shared :func:`build_lora_key_variants`. For each unique checkpoint base key, generates ALL known LoRA naming variants via :func:`~engine.key_mapper.build_lora_key_variants` and registers them in the map:: lora_key_variant → ckpt_base_key (without .weight/.bias/.alpha suffix) Args: ckpt_sd: Checkpoint state dict. Returns: Dict mapping lora_key_variant → ckpt_base_key. """ reverse_map: Dict[str, str] = {} # Collect unique base keys (without .weight/.bias/.alpha suffix) for k in ckpt_sd: base = _strip_tensor_suffix(k) if base in reverse_map: continue # already processed # Generate all known LoRA variants for this checkpoint key variants = build_lora_key_variants(base) for variant in variants: reverse_map[variant] = base return reverse_map # =================================================================== # Block ID Extraction (used by multiple strategies) # =================================================================== def _is_spatial_depth_block(block_id: str) -> bool: """ Check if a block ID represents a UNet spatial depth level. Spatial depth blocks (input_blocks, output_blocks) have resolution levels where adjacent blocks represent different spatial scales. Cross-depth redirection between non-adjacent levels causes Picasso effect (pattern from one resolution appears at wrong resolution). NOT spatial depth blocks: - double_blocks (Flux) — same resolution throughout - single_blocks (Flux) — same resolution throughout - encoder.layers (TE) — not resolution-based - middle_block — single block, no depth variation """ return block_id.startswith('input_blocks.') or block_id.startswith('output_blocks.') def _extract_block_id(key: str) -> Optional[str]: """ Extract a canonical 'block ID' from a model key. The block ID identifies which structural block a tensor belongs to, abstracting away the specific tensor function (weight, bias, etc.). For UNet keys (input_blocks/output_blocks/middle_block): input_blocks.3.1.attn1 → input_blocks.3.1.attn1 For TE encoder layers: encoder.layers.0.self_attn → encoder.layers.0.self_attn For Flux blocks: double_blocks.0.attn → double_blocks.0.attn Strips model prefixes (model.diffusion_model., model.text_model., etc.) and tensor suffixes (.weight, .bias). """ if key is None: return None # Strip model prefixes k = key if k.startswith('model.diffusion_model.'): k = k[len('model.diffusion_model.'):] elif k.startswith('model.text_model.'): k = k[len('model.text_model.'):] elif k.startswith('model.'): k = k[len('model.'):] elif k.startswith('diffusion_model.'): k = k[len('diffusion_model.'):] # For Flux transformer blocks if k.startswith('transformer.'): k = k[len('transformer.'):] # For conditioner embedders if k.startswith('conditioner.embedders.'): parts = k.split('.') if len(parts) >= 4 and parts[2].isdigit(): k = '.'.join(parts[3:]) # For conditioner transformer.text_model prefix if k.startswith('transformer.text_model.'): k = k[len('transformer.text_model.'):] # For text_model prefix (from conditioner) if k.startswith('text_model.'): k = k[len('text_model.'):] parts = k.split('.') # For input_blocks/output_blocks/middle_block (UNet): if parts[0] in ('input_blocks', 'output_blocks', 'middle_block'): if len(parts) >= 3 and parts[1].isdigit(): block_type = parts[0] block_num = parts[1] rest = parts[2:] sub_block = parts[2] if parts[2].isdigit() else '' component = '' for p in rest: if p in ('attn1', 'attn2', 'mlp', 'fc', 'proj', 'conv'): component = p break if component: if sub_block: return f"{block_type}.{block_num}.{sub_block}.{component}" return f"{block_type}.{block_num}.{component}" return f"{block_type}.{block_num}" # For TE encoder layers: if parts[0] == 'encoder' and len(parts) >= 3: if parts[1] == 'layers' and parts[2].isdigit(): layer_num = parts[2] component = parts[3] if len(parts) > 3 else '' return f"encoder.layers.{layer_num}.{component}" # For Flux double/single blocks: if parts[0] in ('double_blocks', 'single_blocks') and len(parts) >= 2: if parts[1].isdigit(): block_num = parts[1] component = parts[2] if len(parts) > 2 else '' return f"{parts[0]}.{block_num}.{component}" # Fallback: just return first 3 segments if len(parts) >= 3: return '.'.join(parts[:3]) return '.'.join(parts) if parts else None # =================================================================== # Shape+Block Index Builder (O(1) lookup for Strategy 4) # =================================================================== def _build_shape_block_index( ckpt_sd: Dict[str, torch.Tensor], ) -> Dict[Tuple[Tuple[int, ...], str], List[str]]: """ Build index: (shape_tuple, block_type) -> [full_ckpt_keys]. Enables O(1) candidate lookup in _global_search_by_shape_block() instead of O(N) scan over all checkpoint keys. Only indexes keys where _extract_block_id() returns a valid block ID. Keys with unrecognizable structure (e.g., embeddings, scalars) are excluded — they would never match any LoRA key via shape+block anyway. Model support: - SDXL: input_blocks.*, output_blocks.*, middle_block.* - SD1.5: input_blocks.*, output_blocks.*, middle_block.* (same format) - Flux: double_blocks.*, single_blocks.* - Z-Image: layers.* (via fallback block_id) - Anima: custom prefixes (via fallback block_id) """ index: Dict[Tuple[Tuple[int, ...], str], List[str]] = {} for ckpt_key, tensor in ckpt_sd.items(): if not isinstance(tensor, torch.Tensor): continue shape = tuple(tensor.shape) base = _strip_tensor_suffix(ckpt_key) block_id = _extract_block_id(base) if block_id is None: continue block_type = block_id.split('.')[0] # 'input_blocks', 'double_blocks', 'layers', etc. key = (shape, block_type) index.setdefault(key, []).append(ckpt_key) # store FULL key (with suffix) return index def _build_shape_block_index_from_header( header: Dict[str, Any], ) -> Dict[Tuple[Tuple[int, ...], str], List[str]]: """ Build shape+block index from safetensors header dict (no tensor loading). Reads shapes from the header JSON instead of loading tensor data. Memory: ~50 KB for 425 keys vs 17+ GiB for full tensor load. """ index: Dict[Tuple[Tuple[int, ...], str], List[str]] = {} for ckpt_key, info in header.items(): if not isinstance(info, dict) or 'shape' not in info: continue shape = tuple(info['shape']) base = _strip_tensor_suffix(ckpt_key) block_id = _extract_block_id(base) if block_id is None: continue block_type = block_id.split('.')[0] key = (shape, block_type) index.setdefault(key, []).append(ckpt_key) return index def _build_shape_index( ckpt_sd: Dict[str, torch.Tensor], ) -> Dict[Tuple[int, ...], List[str]]: """ Build index: shape_tuple -> [full_ckpt_keys]. Used by Strategy 5 (shape match last resort) for O(1) lookup. NOTE: When ckpt_sd is a _LazyCheckpointMapping, this calls .items() which loads ALL tensors. For mmap mode, use _build_shape_index_from_header() instead. """ index: Dict[Tuple[int, ...], List[str]] = {} for ckpt_key, tensor in ckpt_sd.items(): if not isinstance(tensor, torch.Tensor): continue shape = tuple(tensor.shape) index.setdefault(shape, []).append(ckpt_key) return index def _build_shape_index_from_header( header: Dict[str, Any], ) -> Dict[Tuple[int, ...], List[str]]: """ Build shape index from safetensors header dict (no tensor loading). The header maps key -> {dtype, shape, data_offsets}. This reads shapes directly from the header without loading any tensor data. Memory: ~50 KB for 425 keys vs 17+ GiB for full tensor load. """ index: Dict[Tuple[int, ...], List[str]] = {} for ckpt_key, info in header.items(): if not isinstance(info, dict) or 'shape' not in info: continue shape = tuple(info['shape']) index.setdefault(shape, []).append(ckpt_key) return index # =================================================================== # Strategy 4: Global Search by Shape + Block ID # =================================================================== def _global_search_by_shape_block( lora_base: str, delta: torch.Tensor, ckpt_sd: Dict[str, torch.Tensor], ckpt_base_lookup: Dict[str, str], matched_keys: Set[str], shape_block_index: Dict[Tuple[Tuple[int, ...], str], List[str]], # NEW ) -> Optional[str]: """ Global search: match an orphan LoRA key to a checkpoint key by tensor shape + block ID similarity. Uses pre-computed shape_block_index for O(1) candidate filtering instead of scanning all checkpoint keys. Pipeline: 1. Extract the block ID from the LoRA key 2. Look up candidates by (delta_shape, block_type) in index 3. Among filtered candidates, rank by block ID similarity 4. Return the best match not already used Matching logic is IDENTICAL to the original — only the lookup changes. """ delta_shape = delta.shape if delta.dim() < 2: return None lora_block_id = _extract_block_id(lora_base) if lora_block_id is None: return None block_type = lora_block_id.split('.')[0] candidates: List[Tuple[str, float]] = [] # O(1) lookup instead of O(N) scan over all checkpoint keys candidate_keys = shape_block_index.get((delta_shape, block_type), []) for ckpt_key in candidate_keys: ckpt_tensor = ckpt_sd[ckpt_key] if not isinstance(ckpt_tensor, torch.Tensor): continue if ckpt_key in matched_keys: continue # Shape already matches by index — no need to check again ckpt_base = _strip_tensor_suffix(ckpt_key) ckpt_block_id = _extract_block_id(ckpt_base) if ckpt_block_id is None: continue # Simple string overlap score lora_parts = set(lora_block_id.split('.')) ckpt_parts = set(ckpt_block_id.split('.')) overlap = len(lora_parts & ckpt_parts) total = len(lora_parts | ckpt_parts) similarity = overlap / total if total > 0 else 0.0 # Bonus: same block type prefix if lora_block_id.split('.')[0] == ckpt_block_id.split('.')[0]: similarity += 0.3 # Bonus: same numeric block index lora_nums = [p for p in lora_block_id.split('.') if p.isdigit()] ckpt_nums = [p for p in ckpt_block_id.split('.') if p.isdigit()] if lora_nums and ckpt_nums and lora_nums[0] == ckpt_nums[0]: similarity += 0.2 # 🚫 Architecture-gated cross-depth rejection (F1) if _is_spatial_depth_block(lora_block_id): if lora_nums and ckpt_nums: offset = abs(int(ckpt_nums[0]) - int(lora_nums[0])) if offset >= 2: continue elif offset == 1: similarity -= 0.3 # 🔧 Shape proportion compatibility check (F2) if delta.dim() == 2 and ckpt_tensor.dim() == 2: out_ratio = max(delta.shape[0], ckpt_tensor.shape[0]) / max(min(delta.shape[0], ckpt_tensor.shape[0]), 1) if out_ratio > 2.0: continue # 🔧 VRAM-1: Early-exit at >0.8 similarity — sufficiently good match found, # no need to scan remaining checkpoint keys. if similarity > 0.8: ckpt_base = _strip_tensor_suffix(ckpt_key) if ckpt_base in ckpt_base_lookup: return ckpt_base_lookup[ckpt_base] return ckpt_key candidates.append((ckpt_key, similarity)) if not candidates: return None candidates.sort(key=lambda x: x[1], reverse=True) best_key, best_score = candidates[0] if best_score > 0.3: ckpt_base = _strip_tensor_suffix(best_key) if ckpt_base in ckpt_base_lookup: return ckpt_base_lookup[ckpt_base] return best_key return None # =================================================================== # Key Categorization # =================================================================== def _categorize_lora_key(lora_base: str) -> str: """ Categorize a normalized LoRA base key into a component. Delegates to :func:`engine.key_utils.categorize_key`. Returns one of: ``'unet'``, ``'te'``, ``'clip'``, ``'vae'`` (falls back to ``'unet'`` if :func:`categorize_key` returns ``'other'``). """ from ..engine.key_utils import categorize_key as _cat comp = _cat(lora_base) return comp if comp != 'other' else 'unet' # =================================================================== # Diagnostic: Effect Bar Visualization # =================================================================== def _create_effect_bar(mean_percent: float, max_percent: float, width: int = 20) -> str: """Build a compact visual bar representing the effect magnitude.""" if not math.isfinite(mean_percent) or not math.isfinite(max_percent): return "████████████████████ mean=NaN max=NaN" filled = int((mean_percent / 5.0) * width) filled = max(0, min(filled, width)) remaining = width - filled bar = '█' * filled + '▁' * remaining return f"{bar} mean={mean_percent}% max={max_percent}%" # =================================================================== # Main Matching Entry Point (instance method) # =================================================================== def _find_matching_keys( self, ckpt_sd: Dict[str, torch.Tensor], lora_deltas: Dict[str, torch.Tensor], return_stats: bool = False, ckpt_header: Optional[Dict[str, Any]] = None, ) -> Dict[str, torch.Tensor]: """ Match LoRA delta keys to checkpoint keys using ComfyUI-native reverse key map as the primary strategy, with multiple fallback layers. After identity_normalize() converts LoRA keys to master format, this method tries a cascade of strategies: PRIMARY: 0. **Reverse Key Map** — Build a map from each checkpoint key to ALL possible LoRA naming variants. FALLBACK (kept for edge cases): 1. Exact match (normalized lookup without .weight suffix) 2. Prefix-aware conversion 2.5a. Underscore-to-dot fuzzy conversion 2.5b. SDXL TE2 conditioner prefix 2.5d. Deterministic TE Direct-Mapping (Kohya-inspired) 2.5e. Post-normalization TE Direct Mapping 3. Suffix match (lowercase, normalized) 4. Global Search by Shape + Block ID 5. Shape match (last resort) NOTE: This function receives `self` for access to self.device, self._verbose, etc. When `ckpt_header` is provided, shape/block indices are built from the header JSON (no tensors loaded) instead of from the full state dict. This enables the Survey→Match→Weave pipeline where matching happens after only header reading (~50 KB instead of 17+ GiB for Klein 9B). """ print(" 🔗 Matching LoRA keys to checkpoint keys...") # Quick diagnostic: count keys by prefix pattern sorted_ckpt_keys = sorted(ckpt_sd.keys()) model_prefix = sum(1 for k in sorted_ckpt_keys if k.startswith('model.')) diff_model_prefix = sum(1 for k in sorted_ckpt_keys if k.startswith('diffusion_model.')) layers_prefix = sum(1 for k in sorted_ckpt_keys if 'layers.' in k) print(f" 📊 Checkpoint key prefixes: model.*={model_prefix}, diffusion_model.*={diff_model_prefix}, layers.*={layers_prefix}") matched: Dict[str, torch.Tensor] = {} unmatched_lora: List[str] = [] # Build normalized lookup: key_without_suffix -> original_key ckpt_base_lookup: Dict[str, str] = {} for k in ckpt_sd.keys(): base = _strip_tensor_suffix(k) suffix_found = next((s for s in ('.weight', '.bias', '.alpha') if k.endswith(s)), None) if base in ckpt_base_lookup: existing_key = ckpt_base_lookup[base] if suffix_found == '.weight' and not existing_key.endswith('.weight'): ckpt_base_lookup[base] = k else: ckpt_base_lookup[base] = k # Build lowercase lookup for suffix matching ckpt_base_lower: Dict[str, str] = {} for base, orig_key in ckpt_base_lookup.items(): ckpt_base_lower[base.lower()] = orig_key # Build suffix index for O(1) candidate filtering in Strategy 3 # Indexes checkpoint base keys by their last path segment (after final '.') # so that suffix matching only iterates relevant candidates instead of # scanning all checkpoint keys per unmatched LoRA key. ckpt_suffix_index: Dict[str, List[str]] = {} for base_lower in ckpt_base_lower: suffix_key = base_lower.rsplit('.', 1)[-1] ckpt_suffix_index.setdefault(suffix_key, []).append(base_lower) # Build shape+block index for O(1) Strategy 4 lookup # Use header-based indices when available (no tensor loading needed) if ckpt_header is not None: shape_block_index = _build_shape_block_index_from_header(ckpt_header) else: shape_block_index = _build_shape_block_index(ckpt_sd) # Build shape index for O(1) Strategy 5 lookup if ckpt_header is not None: shape_index = _build_shape_index_from_header(ckpt_header) else: shape_index = _build_shape_index(ckpt_sd) # ===================================================================== # PHASE 2: Build reverse key map (ComfyUI-native, exhaustive) # ===================================================================== reverse_map = _build_reverse_key_map(ckpt_sd) # ===================================================================== # 🔍 DIAGNOSTIC: Comprehensive reverse_map coverage (runs always) # ===================================================================== print(f" 📊 Reverse map entries: {len(reverse_map)}") if lora_deltas: # Count ALL keys NOT in reverse_map rmap_misses_total = sum(1 for k in lora_deltas if k not in reverse_map) rmap_hits_total = len(lora_deltas) - rmap_misses_total print(f" 📊 Reverse map coverage: {rmap_hits_total}/{len(lora_deltas)} keys in map ({rmap_misses_total} misses)") # Special focus on diffusion_model.* keys (Flux Klein normalized format) diffusion_keys = [k for k in lora_deltas if k.startswith('diffusion_model.')] if diffusion_keys: diffusion_misses = [k for k in diffusion_keys if k not in reverse_map] diffusion_hits = len(diffusion_keys) - len(diffusion_misses) rmap_iter_msg = "❌" if diffusion_misses else "✅" print(f" {rmap_iter_msg} diffusion_model.* keys: {diffusion_hits}/{len(diffusion_keys)} in reverse_map") if diffusion_misses: print(f" Sample misses (first 5): {diffusion_misses[:5]}") print(f" ⚠️ {len(diffusion_misses)} diffusion_model.* keys will use fallback strategies") # Show first 5 misses for any key type all_misses = [k for k in lora_deltas if k not in reverse_map] if all_misses: for k in all_misses[:5]: print(f" ❌ Not in reverse_map: '{k}'") else: print(f" ✅ All {len(lora_deltas)} LoRA keys in reverse_map") # ===================================================================== # 🔍 DIAGNOSTIC: Comprehensive key format trace (verbose-only) # ===================================================================== if self._verbose: has_unet_deltas = any('input_blocks.' in k or 'output_blocks.' in k or 'middle_block.' in k for k in lora_deltas) if has_unet_deltas: ckpt_unet_bases = [] for k in sorted(ckpt_sd.keys()): if 'input_blocks' in k or 'output_blocks' in k or 'middle_block' in k: base = _strip_tensor_suffix(k) if base not in ckpt_unet_bases: ckpt_unet_bases.append(base) print(f" 🔍 [DIAG] Sample checkpoint UNet bases (first 5):") for b in ckpt_unet_bases[:5]: print(f" '{b}'") lora_unet_deltas = sorted([k for k in lora_deltas if 'input_blocks.' in k or 'output_blocks.' in k or 'middle_block.' in k]) print(f" 🔍 [DIAG] Sample LoRA UNet delta bases (first 5):") for b in lora_unet_deltas[:5]: in_rmap = b in reverse_map if not in_rmap and b.startswith('diffusion_model.'): mdl_version = f"model.{b}" mdl_in_rmap = mdl_version in reverse_map print(f" '{b}' → in reverse_map: {in_rmap} → model. version: {mdl_in_rmap}") else: print(f" '{b}' → in reverse_map: {in_rmap}") # Detailed mismatch analysis unmatched = [b for b in lora_unet_deltas if b not in reverse_map] if unmatched: print(f" 🔍 [DIAG] === KEY-BY-KEY MISMATCH ANALYSIS ({len(unmatched)} unmatched) ===") pattern_counts = {} for b in unmatched[:30]: segs = b.split('.') block_prefix = '.'.join(segs[:4]) if len(segs) >= 4 else b ckpt_prefixes = [block_prefix] if block_prefix.startswith('diffusion_model.'): ckpt_prefixes.append(f'model.{block_prefix}') matching_ckpt = [ck for ck in ckpt_unet_bases if any(ck.startswith(p) for p in ckpt_prefixes)] if matching_ckpt: pattern_counts.setdefault("block_exists", 0) pattern_counts["block_exists"] += 1 if pattern_counts["block_exists"] <= 5: ckpt_suffix = matching_ckpt[0][len(block_prefix):] lora_suffix = b[len(block_prefix):] print(f" BLOCK MATCH: '{block_prefix}'") print(f" LoRA suffix: '{lora_suffix}'") print(f" Ckpt suffix: '{ckpt_suffix}'") print(f" LoRA key: '{b}'") print(f" Ckpt key: '{matching_ckpt[0]}'") else: pattern_counts.setdefault("block_missing", 0) pattern_counts["block_missing"] += 1 if len(segs) >= 4: for offset in range(-1, 2): try: idx = int(segs[3]) + offset except (ValueError, IndexError): continue alt_prefix = '.'.join(segs[:3] + [str(idx)]) alt_prefixes = [alt_prefix] if alt_prefix.startswith('diffusion_model.'): alt_prefixes.append(f'model.{alt_prefix}') alt_match = [ck for ck in ckpt_unet_bases if any(ck.startswith(p) for p in alt_prefixes)] if alt_match: if pattern_counts.get("block_missing", 0) <= 3: ckpt_label = alt_match[0] print(f" BLOCK MISMATCH: LoRA '{block_prefix}' → Ckpt has '{alt_prefix}' (offset {offset:+d})") print(f" LoRA key: '{b}'") print(f" Ckpt key: '{ckpt_label}'") pattern_counts.setdefault("block_shift", 0) pattern_counts["block_shift"] += 1 break else: dotted_variant = b.replace('ff_net.', 'ff.net.') if dotted_variant != b and dotted_variant in reverse_map: if pattern_counts.get("block_missing", 0) <= 3: print(f" NAMING MISMATCH (ff_net → ff.net):") print(f" LoRA key: '{b}'") print(f" Dotted in RM: '{dotted_variant}'") pattern_counts.setdefault("naming_mismatch_ff_net", 0) pattern_counts["naming_mismatch_ff_net"] += 1 print(f" 🔍 [DIAG] Mismatch categories:") for cat, count in sorted(pattern_counts.items()): print(f" {cat}: {count}") else: print(f" 🔍 [DIAG] All UNet deltas matched via reverse map! ✅") unet_in_rmap = len(lora_unet_deltas) - len(unmatched) print(f" 🔍 [DIAG] UNet delta keys in reverse map: {unet_in_rmap}/{len(lora_unet_deltas)} (total deltas: {len(lora_deltas)})") else: print(f" 🔍 [DIAG] No UNet-style keys in lora deltas (total deltas: {len(lora_deltas)})") matched_ckpt_keys: Set[str] = set() # Track ckpt_key → lora_base for _LazyDeltaDict consumers. # When lora_deltas is a lazy dict, `matched` stores _ShapeInfo values # (not tensors), so the bake loop needs a reverse mapping to reconstruct # deltas on demand from the stored A/B components. _ckpt_to_lora: Dict[str, str] = {} # Pre-computed tensors (fused QKV, to_out.0 fallback) that cannot be # reconstructed from _LazyDeltaDict components alone. _precomputed: Dict[str, torch.Tensor] = {} # ===================================================================== # FUSED QKV DETECTION # ===================================================================== fused_qkv_heads: Dict[str, torch.Tensor] = {} _head_re = re.compile( r'(?:diffusion_model\.|model\.diffusion_model\.)?' r'layers\.(\d+)\.attention\.(to_q|to_k|to_v)$' ) has_fused_qkv = any( re.search(r'layers\.\d+\.attention\.qkv\.weight$', k) for k in ckpt_sd ) if has_fused_qkv: for lora_base in list(lora_deltas.keys()): if _head_re.match(lora_base): if isinstance(lora_deltas, _LazyDeltaDict): # pop from _LazyDeltaDict reconstructs the delta fused_qkv_heads[lora_base] = lora_deltas.pop(lora_base) else: fused_qkv_heads[lora_base] = lora_deltas.pop(lora_base) total_lora_keys = len(lora_deltas) + len(fused_qkv_heads) # ===================================================================== # COMPONENT MATCHING ANALYSIS # ===================================================================== self._component_breakdown = { "unet": {"lora_keys": 0, "matched": 0}, "te": {"lora_keys": 0, "matched": 0}, "clip": {"lora_keys": 0, "matched": 0}, "vae": {"lora_keys": 0, "matched": 0}, } lora_component_map: Dict[str, str] = {} all_lora_bases: Set[str] = set() def _count_lora_key(lb: str) -> None: all_lora_bases.add(lb) comp = _categorize_lora_key(lb) lora_component_map[lb] = comp self._component_breakdown[comp]["lora_keys"] += 1 for lb in lora_deltas: _count_lora_key(lb) for lb in fused_qkv_heads: _count_lora_key(lb) with ProgressTracker(total=len(lora_deltas), desc="Matching LoRA keys") as match_progress: strategy_counts: Dict[str, int] = {} def _track_strategy(name: str) -> None: strategy_counts[name] = strategy_counts.get(name, 0) + 1 def _is_shape_compatible(full_key: str, delta_tensor: torch.Tensor) -> bool: if full_key not in ckpt_sd: return False ckpt_t = ckpt_sd[full_key] if not isinstance(ckpt_t, torch.Tensor): return False return _check_shape_compatible(delta_tensor, ckpt_t) for lora_base, delta in lora_deltas.items(): if not isinstance(delta, (torch.Tensor, _ShapeInfo)): continue # ============================================================= # Strategy 0 (PRIMARY): Reverse Key Map lookup # ============================================================= if lora_base in reverse_map: ckpt_base = reverse_map[lora_base] if ckpt_base in ckpt_base_lookup: full_key = ckpt_base_lookup[ckpt_base] if _is_shape_compatible(full_key, delta): matched[full_key] = delta matched_ckpt_keys.add(full_key) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[full_key] = lora_base _track_strategy("reverse_map") continue # ============================================================= # Strategy 1: exact match # ============================================================= if lora_base in ckpt_base_lookup: full_key = ckpt_base_lookup[lora_base] if _is_shape_compatible(full_key, delta): matched[full_key] = delta matched_ckpt_keys.add(full_key) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[full_key] = lora_base _track_strategy("exact") continue # ============================================================= # Strategy 2: prefix-aware conversion # ============================================================= ckpt_candidate = _lora_key_to_checkpoint_key(lora_base) if ckpt_candidate is not None: if ckpt_candidate in ckpt_base_lookup: full_key = ckpt_base_lookup[ckpt_candidate] if _is_shape_compatible(full_key, delta): matched[full_key] = delta matched_ckpt_keys.add(full_key) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[full_key] = lora_base _track_strategy("prefix") continue # ============================================================= # Strategy 2.5a: Underscore-to-dot fuzzy conversion # ============================================================= fuzzy_match = _try_underscore_to_dot_conversion(lora_base, ckpt_sd, ckpt_base_lookup) if fuzzy_match is not None: if _is_shape_compatible(fuzzy_match, delta): matched[fuzzy_match] = delta matched_ckpt_keys.add(fuzzy_match) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[fuzzy_match] = lora_base _track_strategy("underscore_to_dot") continue # ============================================================= # Strategy 2.5b: SDXL TE conditioner prefix (handles te1. and te2.) # ============================================================= te_cond_match = _try_te_conditioner_prefix(lora_base, ckpt_sd, ckpt_base_lookup, lora_deltas) if te_cond_match is not None: matched[te_cond_match] = delta matched_ckpt_keys.add(te_cond_match) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[te_cond_match] = lora_base _track_strategy("te_conditioner") continue # ============================================================= # Strategy 2.5d: Deterministic TE Direct-Mapping (Kohya-inspired) # ============================================================= te_direct_match = _try_deterministic_te_mapping( lora_base, ckpt_sd, ckpt_base_lookup, delta, _is_shape_compatible, verbose=getattr(self, '_verbose', False), ) if te_direct_match is not None: matched[te_direct_match] = delta matched_ckpt_keys.add(te_direct_match) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[te_direct_match] = lora_base _track_strategy("deterministic_te") continue # ============================================================= # Strategy 2.5e: Post-normalization TE Direct Mapping # ============================================================= # After identity_normalize() → convert_sd15_diffusers_to_comfyui() # converts lora_te_* to text_model.* format, the reverse_map # (Section 2a) should already register these via pure_te entries. # But as a safety net, directly construct the checkpoint key for # any text_model.* keys that slipped through the above strategies. # # Mirrors Kohya's approach: know the architecture layout a priori # and construct the checkpoint key directly from the LoRA key. if lora_base.startswith('text_model.'): # 🔧 Strip trailing dots from lora_base — delta keys from # _reconstruct_lora_delta may have trailing '.' artifacts from # the SD1.5 Diffusers→ComfyUI conversion that prevent exact # key matching against ckpt_base_lookup (which has clean keys). clean_base = lora_base.rstrip('.') ckpt_candidate = f"cond_stage_model.transformer.{clean_base}" # 🔍 PHASE A2 DIAGNOSTIC: Log first 3 TE keys hitting Strategy 2.5e _te_25e_count = getattr(self, '_te_25e_count', 0) if _te_25e_count < 3: self._te_25e_count = _te_25e_count + 1 in_lookup = ckpt_candidate in ckpt_base_lookup print(f" 🔍 [A2] Strategy 2.5e check: '{lora_base}'") print(f" candidate (clean): '{ckpt_candidate}' (in ckpt_base_lookup: {in_lookup})") if not in_lookup: similar = [k for k in list(ckpt_base_lookup.keys())[:30] if 'text_model' in k and 'encoder' in k and 'layers' in k][:3] if similar: for sk in similar: print(f" similar ckpt key: '{sk}'") print(f" delta shape: {list(delta.shape)}") if ckpt_candidate in ckpt_base_lookup: full_key = ckpt_base_lookup[ckpt_candidate] if _is_shape_compatible(full_key, delta): matched[full_key] = delta matched_ckpt_keys.add(full_key) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[full_key] = lora_base _track_strategy("normalized_te_direct") if getattr(self, '_verbose', False): print(f" ✅ Strategy 2.5e (normalized TE direct): " f"{lora_base} → {full_key}") continue # 🔥 TE-key early exit: skip unmatched TE keys (not in checkpoint) # Uses key_utils substring matching (broader & more correct than startswith) from ..engine.key_utils import is_te_key as _is_te_key if _is_te_key(lora_base): # 🔧 SILENT-2: Track skipped TE keys for diagnostic summary _te_skip_count = getattr(self, '_te_skip_count', 0) self._te_skip_count = _te_skip_count + 1 if getattr(self, '_verbose', False): print(f" ⏭️ Skipping TE key (not in checkpoint): {lora_base}") unmatched_lora.append(lora_base) continue # ============================================================= # Strategy 3: suffix match (lowercase, normalized) # ============================================================= lora_base_lower = lora_base.lower() matched_suffix = False # Quick guard: skip suffix scan if lora key lacks known structural blocks. # All supported architectures (Flux, SDXL, SD1.5, Z-Image, Anima) # use one of these block identifiers in their key hierarchy. # If none match, the key cannot suffix-match any checkpoint key. _KNOWN_BLOCKS = ('double_blocks', 'single_blocks', 'input_blocks', 'middle_blocks', 'output_blocks', 'layers.') if not any(b in lora_base for b in _KNOWN_BLOCKS): unmatched_lora.append(lora_base) continue # 🔧 VRAM-2: Use suffix index for O(1) candidate filtering instead of # scanning ALL ckpt_base_lower items per unmatched key. # Only checkpoint keys whose last path segment matches the LoRA key's # last segment are checked, reducing iteration from O(N) to O(M) where # M << N for most architectures. lora_last_seg = lora_base_lower.rsplit('.', 1)[-1] suffix_candidates = ckpt_suffix_index.get(lora_last_seg, []) for ckpt_lower in suffix_candidates: ckpt_key = ckpt_base_lower[ckpt_lower] if ckpt_lower.endswith(lora_base_lower): if _is_shape_compatible(ckpt_key, delta): matched[ckpt_key] = delta matched_ckpt_keys.add(ckpt_key) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[ckpt_key] = lora_base matched_suffix = True break if matched_suffix: _track_strategy("suffix") continue # Also try suffix match with underscore-to-dot converted lora_base for pattern, replacement, _desc in UNDERSCORE_TO_DOT_PATTERNS: converted = re.sub(pattern, replacement, lora_base) if converted != lora_base: converted_lower = converted.lower() converted_last_seg = converted_lower.rsplit('.', 1)[-1] conv_candidates = ckpt_suffix_index.get(converted_last_seg, []) for ckpt_lower in conv_candidates: ckpt_key = ckpt_base_lower[ckpt_lower] if ckpt_lower.endswith(converted_lower): if _is_shape_compatible(ckpt_key, delta): matched[ckpt_key] = delta matched_ckpt_keys.add(ckpt_key) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[ckpt_key] = lora_base matched_suffix = True break if matched_suffix: break if matched_suffix: _track_strategy("suffix_underscore_to_dot") continue # ============================================================= # Strategy 4: Global Search by Shape + Block ID # ============================================================= shape_block_match = _global_search_by_shape_block( lora_base, delta, ckpt_sd, ckpt_base_lookup, matched_ckpt_keys, shape_block_index, ) if shape_block_match is not None: matched[shape_block_match] = delta matched_ckpt_keys.add(shape_block_match) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[shape_block_match] = lora_base if getattr(self, '_verbose', False): print(f" 🔍 Shape+Block match: {lora_base} → {shape_block_match}") _track_strategy("shape_block") continue # ============================================================= # Strategy 6: shape match as last resort (O(1) indexed) # ============================================================= shape_candidates = shape_index.get(delta.shape, []) for ckpt_key in shape_candidates: if ckpt_key in matched_ckpt_keys: continue if delta.dim() >= 2: matched[ckpt_key] = delta matched_ckpt_keys.add(ckpt_key) if isinstance(lora_deltas, _LazyDeltaDict): _ckpt_to_lora[ckpt_key] = lora_base _track_strategy("shape_match") break else: _track_strategy("unmatched") unmatched_lora.append(lora_base) match_progress += 1 # ===================================================================== # POST-PROCESSING: Fused QKV Resolution # ===================================================================== fused_qkv_resolved = 0 if has_fused_qkv and fused_qkv_heads: fused_pending: Dict[str, Dict[str, str]] = {} for lora_base in fused_qkv_heads: m = _head_re.match(lora_base) if m: layer_num = m.group(1) head = m.group(2) group_key = f"layers.{layer_num}" if group_key not in fused_pending: fused_pending[group_key] = {} fused_pending[group_key][head] = lora_base for group_key, heads in fused_pending.items(): if len(heads) >= 3: layer_num = group_key.split('.')[1] ckpt_qkv_base = f"layers.{layer_num}.attention.qkv" qkv_full_key = None for ckpt_key in ckpt_sd: ckpt_base = _strip_tensor_suffix(ckpt_key) if ckpt_base == ckpt_qkv_base: qkv_full_key = ckpt_key break if qkv_full_key and qkv_full_key not in matched_ckpt_keys: delta_q = fused_qkv_heads[heads['to_q']] delta_k = fused_qkv_heads[heads['to_k']] delta_v = fused_qkv_heads[heads['to_v']] qkv_weight = ckpt_sd[qkv_full_key] if delta_q.dim() >= 2 and qkv_weight.dim() >= 2: if delta_q.shape[0] * 3 == qkv_weight.shape[0]: cat_dim = 0 elif delta_q.shape[1] * 3 == qkv_weight.shape[1]: cat_dim = 1 else: cat_dim = 0 # Ensure dtype consistency before cat target_dtype = delta_q.dtype fused_delta = torch.cat([ delta_q.to(target_dtype), delta_k.to(target_dtype), delta_v.to(target_dtype), ], dim=cat_dim) matched[qkv_full_key] = fused_delta matched_ckpt_keys.add(qkv_full_key) if isinstance(lora_deltas, _LazyDeltaDict): # QKV fused delta cannot be reconstructed from # a single _LazyDeltaDict entry — store precomputed. _precomputed[qkv_full_key] = fused_delta fused_qkv_resolved += 3 if getattr(self, '_verbose', False): print(f" 🔄 Fused QKV: layers.{layer_num}.attention " f"to_q+to_k+to_v → {qkv_full_key} " f"(cat_dim={cat_dim}, shapes: {delta_q.shape}*3 → {fused_delta.shape})") continue else: for head_name, lora_key in heads.items(): unmatched_lora.append(lora_key) continue else: for head_name, lora_key in heads.items(): unmatched_lora.append(lora_key) continue for head_name, lora_key in heads.items(): unmatched_lora.append(lora_key) # --- Phase B: Handle to_out.0 fallback from unmatched --- if unmatched_lora: remaining_unmatched = [] for lora_base in unmatched_lora: m_out = re.match( r'(?:diffusion_model\.|model\.diffusion_model\.)?' r'layers\.(\d+)\.attention\.to_out\.0$', lora_base ) if m_out: layer_num = m_out.group(1) ckpt_out_base = f"layers.{layer_num}.attention.out" out_full_key = None for ckpt_key in ckpt_sd: ckpt_base = _strip_tensor_suffix(ckpt_key) if ckpt_base == ckpt_out_base: out_full_key = ckpt_key break if out_full_key and out_full_key not in matched_ckpt_keys: delta = lora_deltas[lora_base] matched[out_full_key] = delta matched_ckpt_keys.add(out_full_key) if isinstance(lora_deltas, _LazyDeltaDict): # Already reconstructed by __getitem__ above — store precomputed. _precomputed[out_full_key] = delta fused_qkv_resolved += 1 if getattr(self, '_verbose', False): print(f" 🔄 to_out.0→out: {lora_base} → {out_full_key}") continue remaining_unmatched.append(lora_base) unmatched_lora = remaining_unmatched # --- Compute component breakdown from matched vs unmatched --- matched_lora_bases = all_lora_bases - set(unmatched_lora) for lb in matched_lora_bases: comp = lora_component_map.get(lb, 'unet') self._component_breakdown[comp]["matched"] += 1 for comp, data in self._component_breakdown.items(): if data["lora_keys"] > 0 and data["matched"] == 0: data["reason"] = f"no matching {comp} keys found in checkpoint" # --- Print results --- if unmatched_lora: print(f" ⚠️ Could not match {len(unmatched_lora)} LoRA keys:") for uk in unmatched_lora[:10]: print(f" - {uk}") in_rmap = uk in reverse_map in_ckpt_lookup = uk in ckpt_base_lookup in_ckpt_lower = uk.lower() in ckpt_base_lower suffix = uk.split('.')[-1] if '.' in uk else uk close_matches = [] for ckpt_lower_key, ckpt_orig_key in ckpt_base_lower.items(): if ckpt_lower_key.endswith('.' + suffix.lower()): close_matches.append(ckpt_orig_key) elif ckpt_lower_key.endswith(suffix.lower()): close_matches.append(ckpt_orig_key) print(f" In reverse_map: {in_rmap} | In ckpt_lookup: {in_ckpt_lookup} | In ckpt_lower: {in_ckpt_lower}") if close_matches: seen = set() unique_matches = [] for cm in close_matches: if cm not in seen: seen.add(cm) unique_matches.append(cm) print(f" 📋 Close ckpt matches (suffix='{suffix}', showing {min(3, len(unique_matches))}):") for cm in unique_matches[:3]: print(f" - {cm}") print(f" 💡 Suggestion: Check if suffix '{suffix}' needs a new reverse_map entry") if len(unmatched_lora) > 10: print(f" ... and {len(unmatched_lora) - 10} more unmatched keys (see above for truncated list)") diffusers_hints = [uk for uk in unmatched_lora if "down_blocks_" in uk or "up_blocks_" in uk or "mid_block_" in uk] if diffusers_hints: print(f" 💡 TIP: {len(diffusers_hints)} unmatched key(s) contain SD1.5 Diffusers naming") print(f" (down_blocks_/up_blocks_/mid_block_). This suggests the LoRA wasn't") print(f" properly detected as SD1.5 Diffusers format during normalization.") else: print(f" ✅ All {total_lora_keys} LoRA keys matched successfully!") if fused_qkv_resolved > 0: print(f" 🔄 Fused QKV resolution: {fused_qkv_resolved} attention heads combined into fused deltas") # 🔧 SILENT-2: Report TE keys skipped due to absence in checkpoint _te_skip_total = getattr(self, '_te_skip_count', 0) if _te_skip_total > 0: print(f" ⏭️ Skipped {_te_skip_total} TE keys (not present in checkpoint — expected for SDXL/Flux LoRA on mismatched base)") print(f" ✅ Matched {len(matched)} LoRA deltas to checkpoint keys") # ── Wrap matched in _MatchedDeltas for on-demand reconstruction ── # When lora_deltas is a _LazyDeltaDict, matched contains _ShapeInfo # values instead of real tensors. Wrap it in _MatchedDeltas so that # the bake loop can pop() one delta at a time, reconstructing on demand. use_lazy = isinstance(lora_deltas, _LazyDeltaDict) if use_lazy and _precomputed: # Some matched entries (fused QKV, to_out.0) have precomputed # tensors that must be stored directly. Create the wrapper. lazy_matched = _MatchedDeltas(lora_deltas, _ckpt_to_lora, _precomputed) elif use_lazy: lazy_matched = _MatchedDeltas(lora_deltas, _ckpt_to_lora) else: lazy_matched = matched if return_stats: return lazy_matched, strategy_counts return lazy_matched