ZImageModel extension for Ostris ai-toolkit. Fixes: - LoRA key conversion (Ostris → ZImage transformer layers) - attn_qkv/out, mlp_w1/w2 targets - Assistant LoRA merge support Install: Replace ai-toolkit/extensions_built_in/diffusion_models/z_image/z_image.py Compatible: Z-Image-Turbo, rank 8-32 LoRA.
Description
Comments (4)
What was the need for this, and what does it fix?
Need: Ostris ai-toolkit LoRA weights don't match ZImageTransformer2DModel layer names. Fixes: - Converts Ostris keys (e.g. "diffusion_model.attention.to_q.lora_A") → ZImage ("transformer.attn_qkv.lora_down") - MLP: "ff.net.proj" → "mlp_w1/w2" - Save/load bidirectional conversion - Target modules: "ZImageTransformer2DModel" Without it, LoRA training/merge fails on Z-Image-Turbo. Also optimizes device/load stability.
The many errors in the log have been eliminated, and it seems as if the model is learning a little better.
"Due to elevated AIToolKit, Z-Image, Wan, and Hunyuan training failure rates we've temporarily disabled these options while we investigate!" (CivitaI Trainer) ↑ Civitai's administrators seem to be having trouble; they might need to download this. 😅