Skip to content

Commit 4174214

Browse files
fix(model-manager): add Z-Image LoRA/DoRA detection support
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 (including DoRA format) use keys like: - diffusion_model.layers.X.attention.to_k.lora_down.weight - diffusion_model.layers.X.attention.to_k.dora_scale The base LyCORIS config only checked for lora_A.weight/lora_B.weight suffixes, missing the lora_down.weight/lora_up.weight and dora_scale patterns used by Z-Image LoRAs.
1 parent 0b1befa commit 4174214

File tree

1 file changed

+36
-0
lines changed
  • invokeai/backend/model_manager/configs

1 file changed

+36
-0
lines changed

invokeai/backend/model_manager/configs/lora.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,42 @@ class LoRA_LyCORIS_ZImage_Config(LoRA_LyCORIS_Config_Base, Config_Base):
227227

228228
base: Literal[BaseModelType.ZImage] = Field(default=BaseModelType.ZImage)
229229

230+
@classmethod
231+
def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None:
232+
"""Z-Image LoRAs have different key patterns than SD/SDXL LoRAs.
233+
234+
Z-Image LoRAs use keys like:
235+
- diffusion_model.layers.X.attention.to_k.lora_down.weight (DoRA format)
236+
- diffusion_model.layers.X.attention.to_k.lora_A.weight (PEFT format)
237+
- diffusion_model.layers.X.attention.to_k.dora_scale (DoRA scale)
238+
"""
239+
state_dict = mod.load_state_dict()
240+
241+
# Check for Z-Image specific LoRA patterns
242+
has_z_image_lora_keys = state_dict_has_any_keys_starting_with(
243+
state_dict,
244+
{
245+
"diffusion_model.layers.", # Z-Image S3-DiT layer pattern
246+
},
247+
)
248+
249+
# Also check for LoRA weight suffixes (various formats)
250+
has_lora_suffix = state_dict_has_any_keys_ending_with(
251+
state_dict,
252+
{
253+
"lora_A.weight",
254+
"lora_B.weight",
255+
"lora_down.weight",
256+
"lora_up.weight",
257+
"dora_scale",
258+
},
259+
)
260+
261+
if has_z_image_lora_keys and has_lora_suffix:
262+
return
263+
264+
raise NotAMatchError("model does not match Z-Image LoRA heuristics")
265+
230266
@classmethod
231267
def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
232268
"""Z-Image LoRAs are identified by their diffusion_model.layers structure.

0 commit comments

Comments
 (0)