Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion invokeai/backend/model_manager/configs/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,24 +114,43 @@ def _has_main_keys(state_dict: dict[str | int, Any]) -> bool:


def _has_z_image_keys(state_dict: dict[str | int, Any]) -> bool:
"""Check if state dict contains Z-Image S3-DiT transformer keys."""
"""Check if state dict contains Z-Image S3-DiT transformer keys.

This function returns True only for Z-Image main models, not LoRAs.
LoRAs are excluded by checking for LoRA-specific weight suffixes.
"""
# Z-Image specific keys that distinguish it from other models
z_image_specific_keys = {
"cap_embedder", # Caption embedder - unique to Z-Image
"context_refiner", # Context refiner blocks
"cap_pad_token", # Caption padding token
}

# LoRA-specific suffixes - if present, this is a LoRA not a main model
lora_suffixes = (
".lora_down.weight",
".lora_up.weight",
".lora_A.weight",
".lora_B.weight",
".dora_scale",
)

for key in state_dict.keys():
if isinstance(key, int):
continue

# If we find any LoRA-specific keys, this is not a main model
if key.endswith(lora_suffixes):
return False

# Check for Z-Image specific key prefixes
# Handle both direct keys (cap_embedder.0.weight) and
# ComfyUI-style keys (model.diffusion_model.cap_embedder.0.weight)
key_parts = key.split(".")
for part in key_parts:
if part in z_image_specific_keys:
return True

return False


Expand Down