Add ZImage LoRA support and integrate into ZImagePipeline

[CalamitousFelicitousness]

What does this PR do?

Adds support for LoRA for ZImagePipeline and adds a conversion script into Diffusers format. Resolves issue #12745

Before submitting

• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul

[CalamitousFelicitousness]

Tested using a simple script.

Testing script

#!/usr/bin/env python"""Test script for ZImage LoRA support."""importsyssys.path.insert(0, '/home/ohiom/diffusers/src')
importtorchfromdiffusersimportZImagePipeline# PathsMODEL_PATH="database/models/huggingface/models--Tongyi-MAI--Z-Image-Turbo/snapshots/78771b7e11b922c868dd766476bda1f4fc6bfc96"LORA_PATH="TechnicallyColorZ_V1.safetensors"print("Loading ZImagePipeline...")
pipe=ZImagePipeline.from_pretrained(
MODEL_PATH,
torch_dtype=torch.bfloat16,
local_files_only=True,
)
pipe.to("cuda")
print(f"Pipeline loaded. Has load_lora_weights: {hasattr(pipe, 'load_lora_weights')}")
print(f"\nLoading LoRA from {LORA_PATH}...")
pipe.load_lora_weights(LORA_PATH)
print("LoRA loaded successfully!")
# Generate an imageprompt="t3chnic4lly vibrant 1960s close-up of a woman sitting under a tree in a blue skit and white blouse, she has blonde wavy short hair and a smile with green eyes lake scene by a garden with flowers in the foreground 1960s styl;e film She's holding her hand out there is a small smooth frog in her palm, she's making eye contact with the toad."print(f"\nGenerating image with prompt: {prompt}")
image=pipe(
prompt=prompt,
num_inference_steps=8,
guidance_scale=1.0,
height=1024,
width=1024,
generator=torch.Generator(device="cuda").manual_seed(42),
).images[0]
output_path="test_zimage_lora_output.png"image.save(output_path)
print(f"\nImage saved to {output_path}")

[sayakpaul]

@asomoza cc

[asomoza]

thanks for the contribution, it works ok without scale but we need to add the transformer model here

With that:

without lora	lora 1.0	lora 0.5

Also I did a quick run on the tests to check if this was detected but they all fail with:

ValueError: too many values to unpack (expected 3)

[sayakpaul]

What a clean PR this is! Thanks so much for contributing!

[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.

[CalamitousFelicitousness]

@sayakpaul You are too kind.
@asomoza Thank you so much for your review, and apologies for my mistake, ironically, in the script testing the testing scripts. Now I have updated peft.py and fixes all the errors I can manage.

As for the remaining failures, as far as I can tell, due to ZImage requiring torch.use_deterministic_algorithms(False) those tests cannot be completed. Please advise if I am correct, and those tests can be skipped, or are there any remaining steps for me to do.
From my understanding the changes needed to make those tests pass, like using a single pipeline(?), would be rather intensive, so I'm afraid I will have to defer to your judgement.

[asomoza]

thanks a lot @CalamitousFelicitousness, can you please add the ZImageLoraLoaderMixin to the lora docs here

for the tests, I leave it to @sayakpaul to make the decision

[sayakpaul]

Let's go!

[sayakpaul]

Hmm, the failing tests -- are these because of the numerical instabilities arising from the use of torch.empty within the Z-Image DiT?

[CalamitousFelicitousness]

Squashed one more naming convention issue with an override.

My best somewhat-educated 2AM guess is that it lies somewhere between the Complex64 RoPE and torch.empty, unfortunately here I am beginning to approach the functional limits of my practical knowledge, so if anything comes to mind later on, I will bring it up.

[sayakpaul]

@CalamitousFelicitousness that is understandable. I am facing something similar in #12741

[sayakpaul]

Let's try to add a is_flaky decorator to the LoRA tester class, @CalamitousFelicitousness and see if that helps.

[CalamitousFelicitousness]

They are really flaky. I'm doing 10 retries right now, and it whittled down the failures to 6 at the lowest, some tests are failing one run and pass on another. I will just set it to 100 and see if they pass.

100 retries

==================================================================================================================================== warnings summary ====================================================================================================================================
tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_get_adapters
tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_get_list_adapters
tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_simple_inference_with_text_denoiser_multi_adapter
tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_simple_inference_with_text_denoiser_multi_adapter_delete_adapter
tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_simple_inference_with_text_denoiser_multi_adapter_weighted
tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_simple_inference_with_text_lora_denoiser_fused_multi
/home/ohiom/diffusers/venv/lib/python3.13/site-packages/peft/tuners/tuners_utils.py:282: UserWarning: Already found a `peft_config` attribute in the model. This will lead to having multiple adapters in the model. Make sure to know what you are doing!
warnings.warn(
tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_lora_B_bias
/home/ohiom/diffusers/venv/lib/python3.13/site-packages/peft/tuners/lora/layer.py:170: PeftWarning: `lora_bias=True` was passed but the targeted layer of type Linear has no bias. This means that merging LoRA weights won't be possible.
warnings.warn(
tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_lora_fuse_nan
/home/ohiom/diffusers/venv/lib/python3.13/site-packages/peft/tuners/tuners_utils.py:1817: UserWarning: All adapters are already merged, nothing to do.
warnings.warn("All adapters are already merged, nothing to do.")
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================================================================================================================================ short test summary info =================================================================================================================================
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_logs_info_when_no_lora_keys_found - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_lora_scale_kwargs_match_fusion - AssertionError: False is not true : Fused lora should not change the output
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_set_adapters_match_attention_kwargs - AssertionError: False is not true : Lora + scale should match the output of `set_adapters()`.
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_simple_inference_with_text_denoiser_lora_and_scale - AssertionError: False is not true : Lora + scale should change the output
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_simple_inference_with_text_denoiser_multi_adapter - AssertionError: False is not true : output with no lora and output with lora disabled should give same results
=========================================================================================================== 5 failed, 36 passed, 10 skipped, 8 warnings in 1253.64s (0:20:53) ============================================================================================================

Another 100 retries

================================================================================================================================ short test summary info =================================================================================================================================
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_lora_scale_kwargs_match_fusion - AssertionError: False is not true : Fused lora should not change the output
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_set_adapters_match_attention_kwargs - AssertionError: False is not true : Lora + scale should match the output of `set_adapters()`.
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_simple_inference_with_text_denoiser_lora_and_scale - AssertionError: False is not true : Lora + scale should change the output
FAILED tests/lora/test_lora_layers_z_image.py::ZImageLoRATests::test_simple_inference_with_text_lora_unloaded - AssertionError: False is not true : Fused lora should change the output
=========================================================================================================== 4 failed, 37 passed, 10 skipped, 8 warnings in 2019.36s (0:33:39) ============================================================================================================

Status: We have three remaining tests that haven't passed yet:

1. test_lora_scale_kwargs_match_fusion
2. test_set_adapters_match_attention_kwargs
3. test_simple_inference_with_text_denoiser_lora_and_scale

[CalamitousFelicitousness]

@sayakpaul The last three tests fail, even at 250 retries, I tried to initialise the padding tokens to try and get them to pass, but no dice.

[sayakpaul]

Okay, we can rest this case. Let's skip the LoRA testing class entirely (unittest.skip) with a valid reason and merge. No reason for us to block this PR.

[sayakpaul]

@asomoza since I will be off mostly today, could you take care of merging the PR once the test class has been skipped?

[CalamitousFelicitousness]

Added the skip and cleaned up the accumulated overrides and workarounds.

[asomoza]

thanks a lot!, the failed test now are unrelated to this PR

[vagitablebirdcode]

@asomoza Hi,
I am trying to use the LoRA loading feature with Z-Image-Turbo in diffusers 0.36.0, but all LoRA-related APIs appear to be in an inconsistent or broken state. Below is a minimal reproduction and the observed behavior.

Environment

• diffusers: 0.36.0
• peft: 0.17.0
• Model: Z-Image-Turbo
• Test LoRA: https://huggingface.co/tarn59/pixel_art_style_lora_z_image_turbo

Minimal Reproduction

fromdiffusersimportDiffusionPipelinemodel=DiffusionPipeline.from_pretrained("/weights/Z-Image-Turbo")
lora_path="/path/to/lora"model.load_lora_weights(lora_path, adapter_name="xxx")
model.get_list_adapters()
model.set_adapters(adapter_names="xxx")

Observed Behavior

1. Adapter is not registered

   • After calling load_lora_weights(..., adapter_name="xxx"),

     model.get_list_adapters()

     returns {}.

   • The adapter name is not visible to the adapter management system.

2. set_adapters fails

   • Calling:

     model.set_adapters(adapter_names="xxx")

     raises:

     ValueError: Adapter name(s) {'xxx'} not in the list of present adapters: set().
     

   • This indicates the adapter was never properly registered, despite being loaded.

3. unload_lora_weights does not fully unload

   • After calling:

     model.unload_lora_weights()

     the internal LoRA state is not fully cleared.

4. Reloading with the same adapter name fails

   • Calling load_lora_weights again with the same adapter name raises:

     ValueError: Adapter name xxx already in use in the model
     

   • This contradicts the fact that the adapter does not appear in get_list_adapters().

5. Reloading without adapter_name or with a different name

   • Loading again without specifying adapter_name, or using a different name, produces:

     UserWarning: Already found a `peft_config` attribute in the model.
     This will lead to having multiple adapters in the model.
     

   • This suggests that LoRA weights remain attached to the base model even after unload_lora_weights().

Summary of Problems

• load_lora_weights does not correctly register adapters.
• get_list_adapters does not reflect loaded adapters.
• set_adapters is unusable due to missing registration.
• unload_lora_weights does not fully clean LoRA-related state.
• Adapter name collisions occur even when adapters are supposedly unloaded.

Overall, the LoRA adapter lifecycle (load → register → activate → unload → reload) appears inconsistent for Z-Image-Turbo in diffusers 0.36.0.

Please let me know if I can provide additional debugging information or test patches.

In [1]: fromdiffusersimportDiffusionPipelineIn [2]: model=DiffusionPipeline.from_pretrained("/weights/Z-Image-Turbo")
In [3]: lora_path="/path/to/lora"In [4]: model.load_lora_weights(lora_path, adapter_name="xxx")
In [5]: model.get_list_adapters()
Out[5]: {}
In [6]: model.set_adapters(adapter_names="xxx")
---------------------------------------------------------------------------ValueErrorTraceback (mostrecentcalllast)
CellIn[6], line1---->1model.set_adapters(adapter_names="xxx")
File/usr/local/lib/python3.11/site-packages/diffusers/loaders/lora_base.py:736, inLoraBaseMixin.set_adapters(self, adapter_names, adapter_weights)
734missing_adapters=set(adapter_names) -all_adapters735iflen(missing_adapters) >0:
-->736raiseValueError(
737f"Adapter name(s) {missing_adapters} not in the list of present adapters: {all_adapters}."738 )
740# eg {"adapter1": ["unet"], "adapter2": ["unet", "text_encoder"]}741invert_list_adapters= {
742adapter: [partforpart, adaptersinlist_adapters.items() ifadapterinadapters]
743foradapterinall_adapters744 }
ValueError: Adaptername(s) {'xxx'} notinthelistofpresentadapters: set().
In [7]: model.unload_lora_weights()
In [8]: model.load_lora_weights(lora_path, adapter_name="xxx")
---------------------------------------------------------------------------ValueErrorTraceback (mostrecentcalllast)
CellIn[8], line1---->1model.load_lora_weights(lora_path, adapter_name="xxx")
File/usr/local/lib/python3.11/site-packages/diffusers/loaders/lora_pipeline.py:5190, inZImageLoraLoaderMixin.load_lora_weights(self, pretrained_model_name_or_path_or_dict, adapter_name, hotswap, **kwargs)
5187ifnotis_correct_format:
5188raiseValueError("Invalid LoRA checkpoint.")
->5190self.load_lora_into_transformer(
5191state_dict,
5192transformer=getattr(self, self.transformer_name) ifnothasattr(self, "transformer") elseself.transformer,
5193adapter_name=adapter_name,
5194metadata=metadata,
5195_pipeline=self,
5196low_cpu_mem_usage=low_cpu_mem_usage,
5197hotswap=hotswap,
5198 )
File/usr/local/lib/python3.11/site-packages/diffusers/loaders/lora_pipeline.py:5222, inZImageLoraLoaderMixin.load_lora_into_transformer(cls, state_dict, transformer, adapter_name, _pipeline, low_cpu_mem_usage, hotswap, metadata)
5220# Load the layers corresponding to transformer.5221logger.info(f"Loading {cls.transformer_name}.")
->5222transformer.load_lora_adapter(
5223state_dict,
5224network_alphas=None,
5225adapter_name=adapter_name,
5226metadata=metadata,
5227_pipeline=_pipeline,
5228low_cpu_mem_usage=low_cpu_mem_usage,
5229hotswap=hotswap,
5230 )
File/usr/local/lib/python3.11/site-packages/diffusers/loaders/peft.py:221, inPeftAdapterMixin.load_lora_adapter(self, pretrained_model_name_or_path_or_dict, prefix, hotswap, **kwargs)
219iflen(state_dict) >0:
220ifadapter_nameingetattr(self, "peft_config", {}) andnothotswap:
-->221raiseValueError(
222f"Adapter name {adapter_name} already in use in the model - please select a new adapter name."223 )
224elifadapter_namenotingetattr(self, "peft_config", {}) andhotswap:
225raiseValueError(
226f"Trying to hotswap LoRA adapter '{adapter_name}' but there is no existing adapter by that name. "227"Please choose an existing adapter name or set `hotswap=False` to prevent hotswapping."228 )
ValueError: Adapternamexxxalreadyinuseinthemodel-pleaseselectanewadaptername.
In [9]: model.load_lora_weights(lora_path)
/usr/local/lib/python3.11/site-packages/peft/tuners/tuners_utils.py:196: UserWarning: Alreadyfounda`peft_config`attributeinthemodel. Thiswillleadtohavingmultipleadaptersinthemodel. Makesuretoknowwhatyouaredoing!
warnings.warn(
In [10]: model.unload_lora_weights()
IIn [11]: model.load_lora_weights(lora_path)
/usr/local/lib/python3.11/site-packages/peft/tuners/tuners_utils.py:196: UserWarning: Alreadyfounda`peft_config`attributeinthemodel. Thiswillleadtohavingmultipleadaptersinthemodel. Makesuretoknowwhatyouaredoing!
warnings.warn(
In [12]: model.unload_lora_weights()
In [13]: model.load_lora_weights(lora_path)
/usr/local/lib/python3.11/site-packages/peft/tuners/tuners_utils.py:196: UserWarning: Alreadyfounda`peft_config`attributeinthemodel. Thiswillleadtohavingmultipleadaptersinthemodel. Makesuretoknowwhatyouaredoing!
warnings.warn(

[asomoza]

Hi @vagitablebirdcode , with that lora it works without any problems:

no lora	lora

since you're using an LLM for this issue, I don't really know what you're doing, but if you're using plain python you need to print the return to see the values, ie:

pipe.load_lora_weights("tarn59/pixel_art_style_lora_z_image_turbo", adapter_name="pixel_art")
print(pipe.get_list_adapters())

output:

{'transformer': ['pixel_art']}