mirror of
https://github.com/invoke-ai/InvokeAI
synced 2026-03-02 13:09:06 +01:00
fix(model-manager): add Z-Image LoRA/DoRA detection and loading support
Two fixes for Z-Image LoRA support: 1. Override _validate_looks_like_lora in LoRA_LyCORIS_ZImage_Config to recognize Z-Image specific LoRA formats that use different key patterns than SD/SDXL LoRAs. Z-Image LoRAs use lora_down.weight/lora_up.weight and dora_scale suffixes instead of lora_A.weight/lora_B.weight. 2. Fix _group_by_layer in z_image_lora_conversion_utils.py to correctly group LoRA keys by layer name. The previous logic used rsplit with maxsplit=2 which incorrectly grouped keys like: - "to_k.alpha" -> layer "diffusion_model.layers.17.attention" - "lora_down.weight" -> layer "diffusion_model.layers.17.attention.to_k" Now uses suffix matching to ensure all keys for a layer are grouped together (alpha, dora_scale, lora_down.weight, lora_up.weight).
This commit is contained in:
parent
0b1befa9ab
commit
355c985cc3
@ -227,6 +227,42 @@ class LoRA_LyCORIS_ZImage_Config(LoRA_LyCORIS_Config_Base, Config_Base):
|
||||
|
||||
base: Literal[BaseModelType.ZImage] = Field(default=BaseModelType.ZImage)
|
||||
|
||||
@classmethod
|
||||
def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None:
|
||||
"""Z-Image LoRAs have different key patterns than SD/SDXL LoRAs.
|
||||
|
||||
Z-Image LoRAs use keys like:
|
||||
- diffusion_model.layers.X.attention.to_k.lora_down.weight (DoRA format)
|
||||
- diffusion_model.layers.X.attention.to_k.lora_A.weight (PEFT format)
|
||||
- diffusion_model.layers.X.attention.to_k.dora_scale (DoRA scale)
|
||||
"""
|
||||
state_dict = mod.load_state_dict()
|
||||
|
||||
# Check for Z-Image specific LoRA patterns
|
||||
has_z_image_lora_keys = state_dict_has_any_keys_starting_with(
|
||||
state_dict,
|
||||
{
|
||||
"diffusion_model.layers.", # Z-Image S3-DiT layer pattern
|
||||
},
|
||||
)
|
||||
|
||||
# Also check for LoRA weight suffixes (various formats)
|
||||
has_lora_suffix = state_dict_has_any_keys_ending_with(
|
||||
state_dict,
|
||||
{
|
||||
"lora_A.weight",
|
||||
"lora_B.weight",
|
||||
"lora_down.weight",
|
||||
"lora_up.weight",
|
||||
"dora_scale",
|
||||
},
|
||||
)
|
||||
|
||||
if has_z_image_lora_keys and has_lora_suffix:
|
||||
return
|
||||
|
||||
raise NotAMatchError("model does not match Z-Image LoRA heuristics")
|
||||
|
||||
@classmethod
|
||||
def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
|
||||
"""Z-Image LoRAs are identified by their diffusion_model.layers structure.
|
||||
|
||||
@ -140,16 +140,50 @@ def _get_lora_layer_values(layer_dict: dict[str, torch.Tensor], alpha: float | N
|
||||
|
||||
|
||||
def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
|
||||
"""Groups the keys in the state dict by layer."""
|
||||
"""Groups the keys in the state dict by layer.
|
||||
|
||||
Z-Image LoRAs have keys like:
|
||||
- diffusion_model.layers.17.attention.to_k.alpha
|
||||
- diffusion_model.layers.17.attention.to_k.dora_scale
|
||||
- diffusion_model.layers.17.attention.to_k.lora_down.weight
|
||||
- diffusion_model.layers.17.attention.to_k.lora_up.weight
|
||||
|
||||
We need to group these by the full layer path (e.g., diffusion_model.layers.17.attention.to_k)
|
||||
and extract the suffix (alpha, dora_scale, lora_down.weight, lora_up.weight).
|
||||
"""
|
||||
layer_dict: dict[str, dict[str, torch.Tensor]] = {}
|
||||
|
||||
# Known suffixes that indicate the end of a layer name
|
||||
known_suffixes = [
|
||||
".lora_A.weight",
|
||||
".lora_B.weight",
|
||||
".lora_down.weight",
|
||||
".lora_up.weight",
|
||||
".dora_scale",
|
||||
".alpha",
|
||||
]
|
||||
|
||||
for key in state_dict:
|
||||
if not isinstance(key, str):
|
||||
continue
|
||||
# Split the 'lora_A.weight' or 'lora_B.weight' suffix from the layer name.
|
||||
parts = key.rsplit(".", maxsplit=2)
|
||||
layer_name = parts[0]
|
||||
key_name = ".".join(parts[1:])
|
||||
|
||||
# Try to find a known suffix
|
||||
layer_name = None
|
||||
key_name = None
|
||||
for suffix in known_suffixes:
|
||||
if key.endswith(suffix):
|
||||
layer_name = key[: -len(suffix)]
|
||||
key_name = suffix[1:] # Remove leading dot
|
||||
break
|
||||
|
||||
if layer_name is None:
|
||||
# Fallback to original logic for unknown formats
|
||||
parts = key.rsplit(".", maxsplit=2)
|
||||
layer_name = parts[0]
|
||||
key_name = ".".join(parts[1:])
|
||||
|
||||
if layer_name not in layer_dict:
|
||||
layer_dict[layer_name] = {}
|
||||
layer_dict[layer_name][key_name] = state_dict[key]
|
||||
|
||||
return layer_dict
|
||||
|
||||
Loading…
Reference in New Issue
Block a user