From 5bd3e729ee1b84130f7d8c0e4e990e8105443bdd Mon Sep 17 00:00:00 2001 From: christopher5106 Date: Fri, 25 Sep 2026 19:52:08 +0200 Subject: [PATCH 1/2] [LoRA] convert Z-Image LoRAs that only carry the fused `qkv` and bare `out` attention keys A LoRA trained on Z-Image's original module names (musubi-tuner, LyCORIS) has `attention.qkv` and `attention.out` and nothing else. The non-diffusers converter dropped both on the assumption that split `to.q/k/v` and `to_out.0` keys are also present (the Anime-Z layout), which emptied such a LoRA and then raised on its leftover `.alpha` keys. Split the fused key into to_q/to_k/to_v like the single-file converter does and map bare `out` to `to_out.0` when no split keys exist; the Anime-Z case keeps skipping the redundant keys. Co-Authored-By: Claude Fable 5.1 --- .../loaders/lora_conversion_utils.py | 37 +++++++++--- tests/lora/test_lora_conversion_z_image.py | 57 +++++++++++++++++++ 2 files changed, 87 insertions(+), 7 deletions(-) create mode 100644 tests/lora/test_lora_conversion_z_image.py diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 1b7bcc795d6b..fd02d64b64fe 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2900,6 +2900,7 @@ def get_alpha_scales(down_weight, alpha_key): if has_lora_dot_format: dot_keys = list(state_dict.keys()) + dot_key_set = set(dot_keys) for k in dot_keys: if lora_dot_down_key not in k: continue @@ -2908,17 +2909,39 @@ def get_alpha_scales(down_weight, alpha_key): base = k[: -len(lora_dot_down_key)] - # Skip combined "qkv" projection — individual to.q/k/v keys are also present. + # Combined "qkv" projection. Next to split to.q/k/v keys (Anime-Z) it is + # redundant; a LoRA trained on Z-Image's original module names carries ONLY + # the fused key, so split it like the single-file converter does (q, k, v + # chunks along dim 0, the shared down weight repeated). if base.endswith(".qkv"): - state_dict.pop(k) - state_dict.pop(k.replace(lora_dot_down_key, lora_dot_up_key), None) - state_dict.pop(base + ".alpha", None) + block = base[: -len(".qkv")] + down_weight = state_dict.pop(k) + up_weight = state_dict.pop(k.replace(lora_dot_down_key, lora_dot_up_key), None) + has_split_keys = any( + f"{block}.{form}{lora_dot_down_key}" in dot_key_set + for proj in "qkv" + for form in (f"to.{proj}", f"to_{proj}") + ) + if has_split_keys or up_weight is None: + state_dict.pop(base + ".alpha", None) + continue + scale_down, scale_up = get_alpha_scales(down_weight, base + ".alpha") + for proj, up_chunk in zip("qkv", up_weight.chunk(3, dim=0)): + converted_state_dict[f"{block}.to_{proj}.lora_A.weight"] = down_weight * scale_down + converted_state_dict[f"{block}.to_{proj}.lora_B.weight"] = up_chunk * scale_up continue - # Skip bare "out.lora.*" — "to_out.0.lora.*" covers the same projection. + # Bare "out": redundant next to "to_out.0" keys, otherwise it IS to_out.0 + # (normalize_out_key already renamed its alpha to "to_out.0.alpha"). if re.search(r"\.out$", base) and ".to_out" not in base: - state_dict.pop(k) - state_dict.pop(k.replace(lora_dot_down_key, lora_dot_up_key), None) + block = base[: -len(".out")] + down_weight = state_dict.pop(k) + up_weight = state_dict.pop(k.replace(lora_dot_down_key, lora_dot_up_key), None) + if f"{block}.to_out.0{lora_dot_down_key}" in dot_key_set or up_weight is None: + continue + scale_down, scale_up = get_alpha_scales(down_weight, f"{block}.to_out.0.alpha") + converted_state_dict[f"{block}.to_out.0.lora_A.weight"] = down_weight * scale_down + converted_state_dict[f"{block}.to_out.0.lora_B.weight"] = up_weight * scale_up continue # Normalise "to.q/k/v" → "to_q/k/v" for the diffusers output key. diff --git a/tests/lora/test_lora_conversion_z_image.py b/tests/lora/test_lora_conversion_z_image.py new file mode 100644 index 000000000000..6267ea358078 --- /dev/null +++ b/tests/lora/test_lora_conversion_z_image.py @@ -0,0 +1,57 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch + +from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_z_image_lora_to_diffusers + + +def _z_image_attention_only_state_dict(rank=2, dim=8, blocks=("layers.0", "context_refiner.0")): + """musubi-tuner LoRA on Z-Image's original attention names: fused qkv + out, nothing else.""" + sd = {} + for block in blocks: + flat = block.replace(".", "_") + sd[f"lora_unet_{flat}_attention_qkv.lora_down.weight"] = torch.randn(rank, dim) + sd[f"lora_unet_{flat}_attention_qkv.lora_up.weight"] = torch.randn(3 * dim, rank) + sd[f"lora_unet_{flat}_attention_qkv.alpha"] = torch.tensor(float(rank)) + sd[f"lora_unet_{flat}_attention_out.lora_down.weight"] = torch.randn(rank, dim) + sd[f"lora_unet_{flat}_attention_out.lora_up.weight"] = torch.randn(dim, rank) + sd[f"lora_unet_{flat}_attention_out.alpha"] = torch.tensor(float(rank)) + return sd + + +def test_z_image_attention_only_lora_splits_fused_qkv_and_maps_out(): + sd = _z_image_attention_only_state_dict() + converted = _convert_non_diffusers_z_image_lora_to_diffusers(dict(sd)) + for block in ("layers.0", "context_refiner.0"): + for proj in "qkv": + assert converted[f"transformer.{block}.attention.to_{proj}.lora_A.weight"].shape == (2, 8) + assert converted[f"transformer.{block}.attention.to_{proj}.lora_B.weight"].shape == (8, 2) + assert converted[f"transformer.{block}.attention.to_out.0.lora_B.weight"].shape == (8, 2) + # alpha == rank, so the k chunk of the fused up weight comes through unscaled + up = sd["lora_unet_layers_0_attention_qkv.lora_up.weight"] + torch.testing.assert_close(converted["transformer.layers.0.attention.to_k.lora_B.weight"], up[8:16]) + assert len(converted) == 2 * 4 * 2 + + +def test_z_image_fused_qkv_is_skipped_next_to_split_keys(): + sd = _z_image_attention_only_state_dict(blocks=("layers.0",)) + for proj in "qkv": + sd[f"lora_unet_layers_0_attention_to_{proj}.lora_down.weight"] = torch.randn(2, 8) + sd[f"lora_unet_layers_0_attention_to_{proj}.lora_up.weight"] = torch.randn(8, 2) + converted = _convert_non_diffusers_z_image_lora_to_diffusers(dict(sd)) + torch.testing.assert_close( + converted["transformer.layers.0.attention.to_q.lora_B.weight"], + sd["lora_unet_layers_0_attention_to_q.lora_up.weight"], + ) From 1476f6b53ba1391cea065e05cf44f2643657cdd5 Mon Sep 17 00:00:00 2001 From: christopher5106 Date: Sun, 27 Sep 2026 17:34:33 +0200 Subject: [PATCH 2/2] [LoRA] review follow-ups: shorter comment with the layout's origin, drop the private-method test Co-Authored-By: Claude Fable 5.1 --- .../loaders/lora_conversion_utils.py | 9 +-- tests/lora/test_lora_conversion_z_image.py | 57 ------------------- 2 files changed, 5 insertions(+), 61 deletions(-) delete mode 100644 tests/lora/test_lora_conversion_z_image.py diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index fd02d64b64fe..cef2e88454a8 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2909,10 +2909,11 @@ def get_alpha_scales(down_weight, alpha_key): base = k[: -len(lora_dot_down_key)] - # Combined "qkv" projection. Next to split to.q/k/v keys (Anime-Z) it is - # redundant; a LoRA trained on Z-Image's original module names carries ONLY - # the fused key, so split it like the single-file converter does (q, k, v - # chunks along dim 0, the shared down weight repeated). + # Fused "qkv" LoRA, e.g. `lora_unet_layers_0_attention_qkv.lora_down.weight` from + # kohya-style trainers that wrap Z-Image's original module tree (the layout + # https://github.com/utensils/mold/blob/main/crates/mold-inference/src/zimage/lora.rs + # documents as "Kohya / sd-scripts"). Split it into to_q/to_k/to_v like + # `convert_z_image_fused_attention`; next to split keys (Anime-Z) it is redundant. if base.endswith(".qkv"): block = base[: -len(".qkv")] down_weight = state_dict.pop(k) diff --git a/tests/lora/test_lora_conversion_z_image.py b/tests/lora/test_lora_conversion_z_image.py deleted file mode 100644 index 6267ea358078..000000000000 --- a/tests/lora/test_lora_conversion_z_image.py +++ /dev/null @@ -1,57 +0,0 @@ -# coding=utf-8 -# Copyright 2026 HuggingFace Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import torch - -from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_z_image_lora_to_diffusers - - -def _z_image_attention_only_state_dict(rank=2, dim=8, blocks=("layers.0", "context_refiner.0")): - """musubi-tuner LoRA on Z-Image's original attention names: fused qkv + out, nothing else.""" - sd = {} - for block in blocks: - flat = block.replace(".", "_") - sd[f"lora_unet_{flat}_attention_qkv.lora_down.weight"] = torch.randn(rank, dim) - sd[f"lora_unet_{flat}_attention_qkv.lora_up.weight"] = torch.randn(3 * dim, rank) - sd[f"lora_unet_{flat}_attention_qkv.alpha"] = torch.tensor(float(rank)) - sd[f"lora_unet_{flat}_attention_out.lora_down.weight"] = torch.randn(rank, dim) - sd[f"lora_unet_{flat}_attention_out.lora_up.weight"] = torch.randn(dim, rank) - sd[f"lora_unet_{flat}_attention_out.alpha"] = torch.tensor(float(rank)) - return sd - - -def test_z_image_attention_only_lora_splits_fused_qkv_and_maps_out(): - sd = _z_image_attention_only_state_dict() - converted = _convert_non_diffusers_z_image_lora_to_diffusers(dict(sd)) - for block in ("layers.0", "context_refiner.0"): - for proj in "qkv": - assert converted[f"transformer.{block}.attention.to_{proj}.lora_A.weight"].shape == (2, 8) - assert converted[f"transformer.{block}.attention.to_{proj}.lora_B.weight"].shape == (8, 2) - assert converted[f"transformer.{block}.attention.to_out.0.lora_B.weight"].shape == (8, 2) - # alpha == rank, so the k chunk of the fused up weight comes through unscaled - up = sd["lora_unet_layers_0_attention_qkv.lora_up.weight"] - torch.testing.assert_close(converted["transformer.layers.0.attention.to_k.lora_B.weight"], up[8:16]) - assert len(converted) == 2 * 4 * 2 - - -def test_z_image_fused_qkv_is_skipped_next_to_split_keys(): - sd = _z_image_attention_only_state_dict(blocks=("layers.0",)) - for proj in "qkv": - sd[f"lora_unet_layers_0_attention_to_{proj}.lora_down.weight"] = torch.randn(2, 8) - sd[f"lora_unet_layers_0_attention_to_{proj}.lora_up.weight"] = torch.randn(8, 2) - converted = _convert_non_diffusers_z_image_lora_to_diffusers(dict(sd)) - torch.testing.assert_close( - converted["transformer.layers.0.attention.to_q.lora_B.weight"], - sd["lora_unet_layers_0_attention_to_q.lora_up.weight"], - )