[christopher5106]

What does this PR do?

Fixes #14876

A LoRA trained on Z-Image's original module names (musubi-tuner and LyCORIS emit lora_unet_<block>_attention_qkv and lora_unet_<block>_attention_out, and nothing else for attention) is currently emptied by _convert_non_diffusers_z_image_lora_to_diffusers: the converter drops the fused qkv key and the bare out key on the assumption that split to.q/k/v and to_out.0 keys are also present (the Anime-Z layout), then raises ValueError: state_dict should be empty on the leftover .alpha keys.

This PR keeps the Anime-Z behaviour (skip the redundant keys when split ones exist) and otherwise:

• splits the fused qkv LoRA into to_q / to_k / to_v, chunking the up weight along dim 0 like convert_z_image_fused_attention does for base weights and repeating the shared down weight;
• maps the bare out key to to_out.0, consuming the alpha that normalize_out_key had already renamed.

Two tests in tests/lora/test_lora_conversion_z_image.py: the attention-only layout converts to the expected 16 tensors with the right shapes and the k chunk lands unscaled when alpha equals rank; a state dict that carries both fused and split keys still uses the split ones.

Found by generating synthetic LoRAs from every layout we see in the wild and running them through a deployed Z-Image pipeline; the fused-only layout was the one that failed.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? Yes, Claude Code.
  • Read the Coding with AI agents guide.
  • Self-review notes, final round:
    • The change is confined to the dotted-key path (.lora.down.weight). The underscore-key path (diffusion_model.*.lora_down.weight, ComfyUI exports) still passes a fused qkv key through unsplit; left alone because I have not seen such a file, so there is nothing to test it against.
    • When a file carries both bare out and to_out.0 keys, normalize_out_key renames the bare alpha to to_out.0.alpha, which can collide with the real one. Pre-existing and unchanged here.
    • Tests use synthetic tensors rather than a checkpoint fixture, matching the existing conversion tests.
• Did you read the contributor guideline?
• Was this discussed/approved via a GitHub issue? Z-Image LoRAs with only fused attention_qkv + attention_out keys are dropped, then fail on leftover alpha keys #14876, with a CPU-only reproduction.
• Did you write any new necessary tests?

Who can review?

@sayakpaul @BenjaminBossan

[github-actions]

Hi @christopher5106, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[sayakpaul]

Thanks for your contribution. I left two questions.

[sayakpaul]

Not needed. We don't test private methods and in general checkpoints, that are beyond our control.

[christopher5106]

Removed in 1476f6b. On our side the layout stays covered end to end by a synthetic-LoRA run against the deployed pipeline, so nothing is lost by dropping the unit test here.

[sayakpaul]

Could we simplify this statement and provide an example of where this occurs?

[christopher5106]

Simplified in 1476f6b. Where it occurs: kohya-style trainers that wrap Z-Image's original module tree, whose attention is a fused qkv linear plus out, emit lora_unet_layers_0_attention_qkv.lora_down.weight / lora_unet_layers_0_attention_out.lora_down.weight and no split q/k/v keys. musubi-tuner's current zimage_model.py splits attention.qkv.weight into to_q/to_k/to_v before wrapping, so its recent checkpoints use the split names, but files from the fused layout circulate: the utensils/mold Z-Image loader documents it as the "Kohya / sd-scripts" layout and maps attention_qkv onto to_q/to_k/to_v and attention_out onto to_out.0 (https://github.com/utensils/mold/blob/main/crates/mold-inference/src/zimage/lora.rs), which is what this hunk does too, with the same chunk order as convert_z_image_fused_attention. The comment now says that in five lines and links that reference.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.