[Z-Image] joint_attention_kwargs={"scale": ...} is silently ignored when loading LoRA weights #15007

Describe the bug

When a LoRA is loaded into a Z-Image pipeline, joint_attention_kwargs={"scale": ...} has no effect.

The argument is accepted, documented and stored in self._joint_attention_kwargs, but it is never forwarded to self.transformer(...) (pipeline_z_image.py#L533). ZImageTransformer2DModel.forward (transformer_z_image.py#L894) has neither a joint_attention_kwargs argument nor the @apply_lora_scale("joint_attention_kwargs") decorator. No error or warning is raised: the LoRA is silently applied at full strength whatever scale is requested.

Other pipelines such as Flux handle this correctly: FluxTransformer2DModel.forward is decorated with @apply_lora_scale("joint_attention_kwargs") and FluxPipeline passes joint_attention_kwargs=self.joint_attention_kwargs to the transformer.

LoRA loading itself works, and pipe.set_adapters(name, weight) does change the strength, because it sets the weights directly on the LoRA layers. What is missing is the per-call scale.

All six Z-Image pipelines have the same pattern: ZImagePipeline, ZImageControlNetPipeline, ZImageControlNetInpaintPipeline, ZImageImg2ImgPipeline, ZImageInpaintPipeline and ZImageOmniPipeline.

Reproduction

import numpy as np
import torch
from diffusers import ZImagePipeline

pipe = ZImagePipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()
pipe.load_lora_weights("artificialguybr/PIXELART-REDMOND-ZIMAGETURBO")


def generate(**kwargs):
    return pipe(
        "Pixel Art, PixArFK, a red fox sitting in the snow",
        height=1024,
        width=1024,
        num_inference_steps=8,
        guidance_scale=0.0,
        generator=torch.Generator("cpu").manual_seed(0),
        **kwargs,
    ).images[0]


with_lora = generate()  # LoRA at full strength
scale_zero = generate(joint_attention_kwargs={"scale": 0.0})  # expected: the LoRA has no effect
pipe.disable_lora()
without_lora = generate()  # reference: no LoRA at all


def same(a, b):
    return np.array_equal(np.asarray(a), np.asarray(b))


print("scale=0.0 identical to LoRA at full strength:", same(scale_zero, with_lora))
print("scale=0.0 identical to no LoRA:", same(scale_zero, without_lora))

Output on current main:

scale=0.0 identical to LoRA at full strength: True
scale=0.0 identical to no LoRA: False

Expected: False / True, since scale=0.0 should cancel the LoRA's contribution. Actual: True / False, so the scale is ignored.

The image below comes from a longer comparison script. The model, LoRA, prompt and seed are the same in every tile; only the scaling method changes. From left to right: no LoRA, LoRA at full strength, joint_attention_kwargs scale 0.0, joint_attention_kwargs scale 0.5, set_adapters weight 0.0, set_adapters weight 0.5.

• Both joint_attention_kwargs tiles are pixel-identical to full strength (mean absolute difference 0.0).
• set_adapters(0.0) is identical to no LoRA.
• set_adapters(0.5) differs from full strength by 0.194 (LoRA vs no LoRA: 0.263).

Proposed implementation/fix

1. ZImageTransformer2DModel.forward: add joint_attention_kwargs: dict[str, Any] | None = None and decorate it with @apply_lora_scale("joint_attention_kwargs"). The decorator pops "scale" and applies and restores the LoRA scaling around the forward pass, so no manual scale = ... handling is needed in the transformer.
2. In each of the six pipelines, pass joint_attention_kwargs=self.joint_attention_kwargs to self.transformer(...). It has to be passed by keyword, because the decorator reads it from kwargs.
3. Remove the two skip overrides so the shared LoRA tests run for Z-Image.

If the proposed fix is welcome I would be happy to open a PR.

Logs

No error or warning is raised; see the reproduction output above.

System Info

• 🤗 Diffusers version: 0.41.0.dev0 (installed from source, main @ 85c9fa4)
• Platform: Linux-7.0.0-38-generic-x86_64-with-glibc2.39
• Running on Google Colab?: No
• Python version: 3.12.3
• PyTorch version (GPU?): 2.14.1+cu130 (True)
• Huggingface_hub version: 1.33.0
• Transformers version: 5.19.0
• Accelerate version: 1.15.0
• PEFT version: 0.21.2
• Safetensors version: 0.8.0
• xFormers version: not installed
• Accelerator: NVIDIA GeForce RTX 5070 Ti, 16303 MiB
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

Who can help?

@sayakpaul @yiyixuxu