221 lines
8.7 KiB
Python
221 lines
8.7 KiB
Python
"""
|
||
Identity Normalizer for Ariadne Project.
|
||
Maps original LoRA keys to mathematical keys using universal_normalize with mapping capture.
|
||
"""
|
||
import re
|
||
import torch
|
||
from typing import Dict, List, Tuple, Optional
|
||
|
||
try:
|
||
from .klein_normalizer import universal_normalize
|
||
except ImportError:
|
||
from klein_normalizer import universal_normalize
|
||
from .scale_utils import build_alpha_mapping
|
||
|
||
def identity_normalize(state_dict: Dict[str, torch.Tensor], metadata=None) -> Tuple[Dict[str, torch.Tensor], Dict[str, str]]:
|
||
"""
|
||
Normalize LoRA keys while preserving a mapping from normalized key back to original key.
|
||
Includes alpha keys in the mapping.
|
||
|
||
Returns:
|
||
normalized_dict: State dict with normalized keys (mathematical keys).
|
||
key_map: Mapping from normalized_key -> original_key (includes weight and alpha keys).
|
||
"""
|
||
# Create mapping dict
|
||
mapping = {}
|
||
# Call universal_normalize with mapping_ref
|
||
normalized = universal_normalize(state_dict, metadata=metadata, mapping_ref=mapping)
|
||
|
||
# Ensure mapping is populated (some normalizers may not use safe_insert).
|
||
# Fallback: compute mapping by matching tensor ids.
|
||
# Preserve special __format__ entry if present
|
||
saved_format = mapping.get("__format__")
|
||
|
||
if not mapping:
|
||
mapping = compute_mapping_by_tensor_id(state_dict, normalized)
|
||
else:
|
||
# mapping may be incomplete (some keys not captured). Fill gaps.
|
||
missing = set(normalized.keys()) - set(mapping.keys())
|
||
if missing:
|
||
extra_mapping = compute_mapping_by_tensor_id(state_dict, normalized, missing_keys=missing)
|
||
mapping.update(extra_mapping)
|
||
|
||
# Restore saved format if we had one
|
||
if saved_format is not None:
|
||
mapping["__format__"] = saved_format
|
||
|
||
# Add explicit alpha‑key mapping (in case tensor‑id matching failed for alpha tensors)
|
||
alpha_map = build_alpha_mapping(state_dict, mapping)
|
||
mapping.update(alpha_map) # alpha_map entries may already be present; that's fine
|
||
|
||
# Verify mapping is bijective (no duplicate normalized keys)
|
||
# This is already enforced by safe_insert collisions, but double-check.
|
||
norm_to_orig = {}
|
||
duplicates = []
|
||
for norm_key, orig_key in mapping.items():
|
||
if norm_key in norm_to_orig:
|
||
duplicates.append((norm_key, orig_key, norm_to_orig[norm_key]))
|
||
else:
|
||
norm_to_orig[norm_key] = orig_key
|
||
if duplicates:
|
||
print(f"WARNING: Duplicate normalized keys in mapping (should not happen):")
|
||
for norm_key, orig1, orig2 in duplicates[:5]:
|
||
print(f" {norm_key} -> {orig1} vs {orig2}")
|
||
|
||
# Ensure every normalized key has a mapping (should be true)
|
||
missing_mappings = set(normalized.keys()) - set(mapping.keys())
|
||
if missing_mappings:
|
||
print(f"WARNING: {len(missing_mappings)} normalized keys missing mapping; adding identity mapping")
|
||
for k in missing_mappings:
|
||
mapping[k] = k
|
||
|
||
# Count alpha keys in mapping (original keys containing .alpha)
|
||
alpha_count = sum(1 for orig_key in mapping.values() if '.alpha' in orig_key)
|
||
print(f"[INFO] Identity mapping captured {len(mapping)} key pairs (including {alpha_count} alpha keys)")
|
||
return normalized, mapping
|
||
|
||
def _extract_key_signature(key: str) -> Tuple[str, ...]:
|
||
"""
|
||
Extract a structural signature from a key for fuzzy matching.
|
||
|
||
Strips common prefixes and normalizes separators, then extracts
|
||
the sequence of (type, value) pairs — layer numbers, operation names, etc.
|
||
|
||
Example:
|
||
"model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_q.lora_A.weight"
|
||
→ ("input", "blocks", "NUM", 1, "NUM", 1, "transformer", "blocks", "NUM", 0, "attn1", "to", "q", "lora", "A", "weight")
|
||
|
||
"lora_unet_input_blocks_1_1_transformer_blocks_0_attn1_to_q.lora_down.weight"
|
||
→ ("input", "blocks", "NUM", 1, "NUM", 1, "transformer", "blocks", "NUM", 0, "attn1", "to", "q", "lora", "down", "weight")
|
||
"""
|
||
# Strip known prefixes to get at the structural core
|
||
for prefix in ('model.diffusion_model.', 'diffusion_model.', 'lora_unet_',
|
||
'lora_te_', 'lora_te1_', 'lora_te2_',
|
||
'cond_stage_model.transformer.', 'cond_stage_model.'):
|
||
if key.startswith(prefix):
|
||
key = key[len(prefix):]
|
||
break
|
||
|
||
# Normalize: split on both '.' and '_'
|
||
parts = re.split(r'[._]+', key)
|
||
|
||
# Build signature: keep strings as-is, mark digits as ('NUM', int)
|
||
sig_parts: List[Tuple[str, ...]] = []
|
||
for p in parts:
|
||
if p.isdigit():
|
||
sig_parts.append(('NUM', p))
|
||
else:
|
||
sig_parts.append((p,))
|
||
return tuple(sig_parts)
|
||
|
||
|
||
def _score_key_similarity(norm_sig: Tuple, orig_sig: Tuple) -> int:
|
||
"""
|
||
Score similarity between two key signatures.
|
||
Higher score = more similar. Considers matching structural elements
|
||
in order, with exact string matches weighted higher than numeric matches.
|
||
"""
|
||
score = 0
|
||
# Compare element by element up to min length
|
||
min_len = min(len(norm_sig), len(orig_sig))
|
||
for i in range(min_len):
|
||
n = norm_sig[i]
|
||
o = orig_sig[i]
|
||
if n == o:
|
||
score += 2 # exact string match
|
||
elif len(n) == 2 and len(o) == 2 and n[0] == 'NUM' and o[0] == 'NUM':
|
||
# Both are numeric — check if close values
|
||
if n[1] == o[1]:
|
||
score += 1 # exact numeric value
|
||
elif abs(int(n[1]) - int(o[1])) <= 1:
|
||
score += 0 # adjacent layer (not a penalty)
|
||
# Penalize length mismatch
|
||
score -= abs(len(norm_sig) - len(orig_sig))
|
||
return score
|
||
|
||
|
||
def _find_matching_original_key(
|
||
norm_key: str,
|
||
original: Dict[str, torch.Tensor],
|
||
norm_tensor_shape: torch.Size,
|
||
) -> Optional[str]:
|
||
"""
|
||
Find the best-matching original key for a normalized key by structural similarity.
|
||
|
||
Used as fallback when tensor-ID matching fails (e.g., after collision resolution
|
||
in convert_sd15_diffusers_to_comfyui which sums tensors).
|
||
"""
|
||
norm_sig = _extract_key_signature(norm_key)
|
||
best_score = -999
|
||
best_key: Optional[str] = None
|
||
|
||
for orig_key, orig_tensor in original.items():
|
||
orig_sig = _extract_key_signature(orig_key)
|
||
score = _score_key_similarity(norm_sig, orig_sig)
|
||
|
||
# Boost score if tensor shapes match (strong signal)
|
||
if norm_tensor_shape == orig_tensor.shape:
|
||
score += 5
|
||
|
||
if score > best_score:
|
||
best_score = score
|
||
best_key = orig_key
|
||
|
||
# Only return if we have a reasonable match
|
||
if best_score >= 0 and best_key is not None:
|
||
return best_key
|
||
return None
|
||
|
||
|
||
def compute_mapping_by_tensor_id(original: Dict[str, torch.Tensor],
|
||
normalized: Dict[str, torch.Tensor],
|
||
missing_keys: Optional[set] = None) -> Dict[str, str]:
|
||
"""
|
||
Compute mapping by matching tensor objects via id().
|
||
For each normalized key, find the original key with the same tensor object.
|
||
|
||
When tensor-ID matching fails (e.g., after collision resolution that creates
|
||
new tensor objects via sum()), falls back to key-signature matching.
|
||
"""
|
||
# Build a map from tensor id to original key (handle duplicates by preferring first)
|
||
id_to_original = {}
|
||
for k, v in original.items():
|
||
tid = id(v)
|
||
# If duplicate tensor ids (rare), keep first.
|
||
if tid not in id_to_original:
|
||
id_to_original[tid] = k
|
||
|
||
mapping = {}
|
||
unmatched_keys: List[str] = []
|
||
target_keys = missing_keys if missing_keys is not None else normalized.keys()
|
||
|
||
for nk in target_keys:
|
||
if nk not in normalized:
|
||
continue
|
||
v = normalized[nk]
|
||
tid = id(v)
|
||
if tid in id_to_original:
|
||
mapping[nk] = id_to_original[tid]
|
||
else:
|
||
# Tensor may have been copied (e.g., after collision resolution in
|
||
# convert_sd15_diffusers_to_comfyui which sums tensors).
|
||
unmatched_keys.append(nk)
|
||
|
||
if unmatched_keys:
|
||
print(f" [INFO] Tensor-ID matching failed for {len(unmatched_keys)} keys — "
|
||
f"trying key-signature fallback...")
|
||
matched_by_sig = 0
|
||
for nk in unmatched_keys:
|
||
orig_match = _find_matching_original_key(nk, original, normalized[nk].shape)
|
||
if orig_match is not None:
|
||
mapping[nk] = orig_match
|
||
matched_by_sig += 1
|
||
else:
|
||
mapping[nk] = nk
|
||
print(f"WARNING: Could not map normalized key {nk} to any original tensor")
|
||
|
||
if matched_by_sig > 0:
|
||
print(f" ✅ Key-signature fallback matched {matched_by_sig}/{len(unmatched_keys)} keys")
|
||
|
||
return mapping
|