321 lines
14 KiB
Python
321 lines
14 KiB
Python
"""
|
||
Serialization Factory for Ariadne Project.
|
||
Restores original trainer keys using master key map, ensuring perfect compatibility.
|
||
"""
|
||
import torch
|
||
from typing import Dict, Optional, Tuple
|
||
try:
|
||
from .klein_normalizer import get_format_style # optional
|
||
except ImportError:
|
||
from klein_normalizer import get_format_style # optional
|
||
|
||
from ..utils import ProgressTracker
|
||
|
||
|
||
def _clean_double_underscores(s: str) -> str:
|
||
"""Collapse repeated underscores into single underscores."""
|
||
while '__' in s:
|
||
s = s.replace('__', '_')
|
||
return s
|
||
|
||
|
||
# Patterns to match suffixes (in order of priority, shared by TE and non-TE branches)
|
||
_SUFFIX_PATTERNS = [
|
||
(r'\.lora_A\.weight$', '.lora_A.weight'),
|
||
(r'\.lora_B\.weight$', '.lora_B.weight'),
|
||
(r'\.lora_down\.weight$', '.lora_down.weight'),
|
||
(r'\.lora_up\.weight$', '.lora_up.weight'),
|
||
(r'\.alpha$', '.alpha'),
|
||
(r'\.weight$', '.weight'), # Generic weight (should come after lora patterns)
|
||
(r'\.bias$', '.bias'),
|
||
]
|
||
|
||
|
||
def _is_te_key(key: str) -> bool:
|
||
"""Return True if key matches a text encoder (CLIP/Llama/T5) pattern.
|
||
|
||
Delegates to :func:`engine.key_utils.is_te_key`.
|
||
"""
|
||
from ..engine.key_utils import is_te_key as _is_te
|
||
return _is_te(key)
|
||
|
||
|
||
def _match_suffix(key: str, patterns):
|
||
"""Find the first matching suffix pattern for a key.
|
||
|
||
Args:
|
||
key: The key string to search in.
|
||
patterns: List of (regex_pattern, suffix_string) tuples to match.
|
||
|
||
Returns:
|
||
Tuple of (matched_pattern, matched_suffix) or (None, "") if no match.
|
||
"""
|
||
import re
|
||
for pattern, suffix in patterns:
|
||
if re.search(pattern, key):
|
||
return pattern, suffix
|
||
return None, ""
|
||
|
||
|
||
def _convert_to_target_format(key: str, separator_style: str, naming_style: str) -> str:
|
||
"""
|
||
Convert normalized key to target format using specified separator and naming style.
|
||
|
||
Args:
|
||
key: Normalized key (dot notation, lora_A/lora_B naming).
|
||
separator_style: Either "dot" (keep dot notation) or "underscore" (convert dot‑index patterns to underscores).
|
||
naming_style: Either "lora_a_b" (keep lora_A/lora_B) or "lora_down_up" (convert to lora_down/lora_up).
|
||
|
||
Returns:
|
||
Key in target format.
|
||
"""
|
||
import re
|
||
|
||
new_key = key
|
||
is_te = _is_te_key(key)
|
||
is_flux = 'double_blocks' in key or 'single_blocks' in key
|
||
|
||
if is_te:
|
||
# Handle TE prefix based on separator style
|
||
if separator_style in ("underscore", "hybrid"):
|
||
if key.startswith("te1."):
|
||
new_key = "lora_te1_" + key[4:]
|
||
elif key.startswith("te."):
|
||
new_key = "lora_te_" + key[3:]
|
||
elif key.startswith("transformer."):
|
||
# transformer.text_model.encoder.layers.N.xxx → lora_te_text_model.encoder.layers.N.xxx
|
||
# (dots will be converted to underscores later for underscore style)
|
||
new_key = "lora_te_" + key[len("transformer."):]
|
||
elif key.startswith("text_model.encoder."):
|
||
# text_model.encoder.layers.N.xxx → lora_te_layers.N.xxx
|
||
# (dots will be converted to underscores later for underscore style)
|
||
new_key = "lora_te_" + key[len("text_model.encoder."):]
|
||
# If already has lora_te prefix, keep as is
|
||
elif key.startswith("lora_te"):
|
||
new_key = key
|
||
else: # dot style
|
||
# For dot style (comfy_native), add text_encoders. prefix.
|
||
# ComfyUI's model_lora_keys_clip() builds key_map entries as:
|
||
# "text_encoders.{model_key[:-len('.weight')]}" → model_key
|
||
# The LoRA adapter then looks for "{key_map_key}.lora_B.weight" in
|
||
# the merged LoRA, so we must match the text_encoders. prefix.
|
||
# This is added AFTER pattern extraction below, not here.
|
||
pass
|
||
|
||
# Find matching suffix pattern using shared helper
|
||
matched_pattern, matched_suffix = _match_suffix(new_key, _SUFFIX_PATTERNS)
|
||
|
||
if matched_pattern:
|
||
# Extract base part (everything before the matched pattern)
|
||
base = re.sub(matched_pattern, '', new_key)
|
||
# Apply separator conversion to base
|
||
if separator_style in ("underscore", "hybrid"):
|
||
# Convert all dots in base to underscores, then clean
|
||
base = _clean_double_underscores(base.replace(".", "_"))
|
||
else: # dot style
|
||
# Add text_encoders. prefix for ComfyUI's model_lora_keys_clip matching.
|
||
# The key_map expects: "text_encoders.{model_key}" which means the LoRA
|
||
# must have "{text_encoders.{model_key}}.lora_A.weight".
|
||
# Skip if already has text_encoders. or lora_te prefix.
|
||
if not base.startswith("text_encoders.") and not base.startswith("lora_te"):
|
||
base = "text_encoders." + base
|
||
# Reconstruct with proper suffix
|
||
new_key = base + matched_suffix
|
||
else:
|
||
# No pattern matched, apply separator conversion to whole key
|
||
if separator_style in ("underscore", "hybrid"):
|
||
new_key = _clean_double_underscores(new_key.replace(".", "_"))
|
||
else: # dot style
|
||
# Add text_encoders. prefix for TE keys in dot style
|
||
if not new_key.startswith("text_encoders.") and not new_key.startswith("lora_te"):
|
||
if "te." in new_key or "te1." in new_key or new_key.startswith("transformer.") or new_key.startswith("text_model.encoder."):
|
||
new_key = "text_encoders." + new_key
|
||
else:
|
||
# Non‑TE keys (UNet and other)
|
||
unet_patterns = ["input_blocks", "middle_block", "output_blocks", "time_embed"]
|
||
# Add lora_unet_ prefix only for underscore style (SDXL expectation)
|
||
if separator_style in ("underscore", "hybrid") and any(p in key for p in unet_patterns) and not key.startswith("lora_unet_"):
|
||
# Replace diffusion_model. with lora_unet_ if present, otherwise prepend
|
||
if key.startswith("diffusion_model."):
|
||
new_key = "lora_unet_" + key[len("diffusion_model."):]
|
||
else:
|
||
new_key = "lora_unet_" + key
|
||
|
||
# Find matching suffix pattern using shared helper
|
||
matched_pattern, matched_suffix = _match_suffix(new_key, _SUFFIX_PATTERNS)
|
||
|
||
if matched_pattern:
|
||
base = re.sub(matched_pattern, '', new_key)
|
||
# Apply dot‑index conversion for underscore style
|
||
if separator_style in ("underscore", "hybrid") and not is_flux:
|
||
# Convert all dots in base to underscores, then clean
|
||
base = _clean_double_underscores(base.replace(".", "_"))
|
||
# else dot style: keep base unchanged
|
||
new_key = base + matched_suffix
|
||
else:
|
||
# No pattern matched, apply dot‑index conversion for underscore style
|
||
if separator_style in ("underscore", "hybrid") and not is_flux:
|
||
# Convert all dots in key to underscores, then clean
|
||
new_key = _clean_double_underscores(new_key.replace(".", "_"))
|
||
# else dot style: keep unchanged
|
||
|
||
# Apply naming style conversion
|
||
if naming_style == "lora_down_up" and not is_flux:
|
||
if 'lora_A.weight' in new_key:
|
||
new_key = new_key.replace('lora_A.weight', 'lora_down.weight')
|
||
if 'lora_B.weight' in new_key:
|
||
new_key = new_key.replace('lora_B.weight', 'lora_up.weight')
|
||
# else lora_a_b: keep as is
|
||
|
||
# Remove any double underscores introduced by the above steps
|
||
new_key = _clean_double_underscores(new_key)
|
||
|
||
return new_key
|
||
|
||
|
||
# Moved from klein_normalizer.py
|
||
def _is_musubi_key(key: str) -> bool:
|
||
"""Return True if key matches Musubi native format (includes Anima)."""
|
||
musubi_patterns = [
|
||
"lora_unet_layers",
|
||
"lora_unet_double_blocks",
|
||
"lora_unet_blocks", # Anima: lora_unet_blocks_N_... (trainer format)
|
||
"lora_te_res_blocks",
|
||
"lora_te_layers", # Anima: lora_te_layers_N_... (TE trainer format)
|
||
]
|
||
return any(pattern in key for pattern in musubi_patterns)
|
||
|
||
|
||
def _ensure_kohya_naming(key: str, naming_style: str = "lora_down_up") -> str:
|
||
"""
|
||
Convert lora_A/lora_B naming to Kohya standard lora_down/lora_up,
|
||
unless naming_style is "lora_a_b" (keep original naming).
|
||
Leaves other keys unchanged.
|
||
"""
|
||
is_flux = 'double_blocks' in key or 'single_blocks' in key
|
||
if naming_style == "lora_down_up" and not is_flux:
|
||
if 'lora_A.weight' in key:
|
||
key = key.replace('lora_A.weight', 'lora_down.weight')
|
||
if 'lora_B.weight' in key:
|
||
key = key.replace('lora_B.weight', 'lora_up.weight')
|
||
# If naming_style is "lora_a_b", do nothing.
|
||
return key
|
||
|
||
|
||
|
||
|
||
def guess_original_key(norm_key: str, master_key_map: Dict[str, str], target_format: Optional[str] = None) -> str:
|
||
"""
|
||
Heuristic fallback for normalized keys not present in master_key_map.
|
||
Attempts to convert to a plausible original format based on target_format.
|
||
If target_format is None, falls back to SDXL underscore style.
|
||
"""
|
||
# If mapping exists, use it (should have been caught earlier)
|
||
if norm_key in master_key_map:
|
||
return master_key_map[norm_key]
|
||
|
||
# Determine separator and naming style based on target_format
|
||
if target_format is None:
|
||
separator_style = "underscore"
|
||
naming_style = "lora_down_up"
|
||
else:
|
||
separator_style, naming_style = get_format_style(target_format)
|
||
|
||
guessed = _convert_to_target_format(norm_key, separator_style, naming_style)
|
||
print(f"WARNING: No mapping for normalized key {norm_key}, guessing original as {guessed}")
|
||
return guessed
|
||
|
||
|
||
def finalize_for_save(merged_tensors: Dict[str, torch.Tensor],
|
||
master_key_map: Dict[str, str],
|
||
target_format: Optional[str] = None,
|
||
meta_a: Optional[Dict[str, str]] = None,
|
||
meta_b: Optional[Dict[str, str]] = None) -> Tuple[Dict[str, torch.Tensor], Dict[str, str]]:
|
||
"""
|
||
Convert merged normalized tensors back to original trainer keys and produce merged metadata.
|
||
|
||
Args:
|
||
merged_tensors: State dict with normalized keys (mathematical keys).
|
||
master_key_map: Mapping from normalized key -> original key.
|
||
target_format: Detected format of primary LoRA (LoRA A) used to decide separator and naming style.
|
||
meta_a: Metadata from LoRA A (optional).
|
||
meta_b: Metadata from LoRA B (optional).
|
||
|
||
Returns:
|
||
restored_dict: State dict with original keys, ready for saving.
|
||
final_metadata: Merged metadata to embed in the saved file.
|
||
"""
|
||
print("🔄 Restoring original trainer keys...")
|
||
# Determine separator and naming style based on target format
|
||
if target_format is None:
|
||
# fallback to SDXL underscore style for backward compatibility
|
||
separator_style = "underscore"
|
||
naming_style = "lora_down_up"
|
||
else:
|
||
separator_style, naming_style = get_format_style(target_format)
|
||
restored = {}
|
||
missing_map = 0
|
||
guessed_map = 0
|
||
overridden_count = 0
|
||
|
||
total_keys = len(merged_tensors)
|
||
with ProgressTracker(total=total_keys, desc="Restoring keys") as restore_progress:
|
||
for norm_key, tensor in merged_tensors.items():
|
||
if norm_key in master_key_map:
|
||
orig_key = master_key_map[norm_key]
|
||
if _is_musubi_key(orig_key):
|
||
# Musubi/Anima override: use converted target‑format nomenclature instead
|
||
# of restoring the original trainer key. The normalized key (norm_key) is
|
||
# already in the correct format (e.g. diffusion_model.blocks.N.* for Anima).
|
||
if separator_style == "underscore" and naming_style == "lora_down_up":
|
||
# SDXL-style musubi: convert norm_key to underscore/lora_down_up
|
||
guessed = _convert_to_target_format(norm_key, separator_style, naming_style)
|
||
else:
|
||
# Dot-style musubi (Anima): keep the normalized key format
|
||
# diffusion_model.blocks.N.cross_attn.out_proj.lora_A.weight
|
||
# with lora_A/lora_B naming (required by LoRAAdapter.load() diffusers2 format)
|
||
# The naming_style may be "lora_down_up" which would convert lora_A→lora_down,
|
||
# but LoRAAdapter.load() needs lora_A/lora_B (diffusers2 format), so force lora_a_b.
|
||
guessed = _ensure_kohya_naming(norm_key, "lora_a_b")
|
||
restored[guessed] = tensor
|
||
guessed_map += 1
|
||
overridden_count += 1
|
||
else:
|
||
# Priority 1 (Standard): use Identity Map to restore 1:1 original keys.
|
||
# Apply naming style conversion (lora_A/lora_B → lora_down/lora_up) if needed.
|
||
orig_key = _ensure_kohya_naming(orig_key, naming_style)
|
||
restored[orig_key] = tensor
|
||
else:
|
||
# Fallback: guess original key using heuristic, respecting target format
|
||
guessed = guess_original_key(norm_key, master_key_map, target_format)
|
||
guessed = _ensure_kohya_naming(guessed, naming_style)
|
||
restored[guessed] = tensor
|
||
missing_map += 1
|
||
if guessed != norm_key:
|
||
guessed_map += 1
|
||
|
||
restore_progress += 1
|
||
|
||
if missing_map or overridden_count:
|
||
print(f"⚠️ {missing_map} keys missing mapping ({guessed_map} guessed), {overridden_count} Musubi keys overridden")
|
||
|
||
# Ensure alpha keys are present (if any)
|
||
# We rely on alpha keys already being in merged_tensors (if they were present).
|
||
|
||
# Merge metadata using MetadataMerger, then sign via factory
|
||
from ..validation import MetadataMerger
|
||
from .metadata_factory import finalize_metadata
|
||
meta_a_safe = meta_a or {}
|
||
meta_b_safe = meta_b or {}
|
||
merged_meta = MetadataMerger.merge(meta_a_safe, meta_b_safe, "preserve_b")
|
||
final_metadata = finalize_metadata(
|
||
metadata=merged_meta,
|
||
mode="preserve_b",
|
||
component="merger",
|
||
extra_fields={"merged_key_count": str(len(merged_tensors))},
|
||
)
|
||
|
||
print(f"✅ Restored {len(restored)} keys to original format")
|
||
return restored, final_metadata
|
||
|