[LoRA] add LoRA training for Qwen Image 2.1

[linoytsaban]

add t2i and i2i lora training scripts for Qwen Image 2.1

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

[Copilot]

🟡 Changes recommended

One or more issues must be addressed before approval.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds DreamBooth LoRA training for Qwen-Image 2.1 text-to-image and image-to-image workflows.

Changes:

• Added both training scripts with caching, bucketing, quantization, validation, and checkpointing.
• Added smoke tests for LoRA output, metadata, caching, and aspect-ratio buckets.
• Added usage documentation.

File summaries

File	Description
examples/dreambooth/train_dreambooth_lora_qwenimage21.py	Updated as part of this pull request.
examples/dreambooth/train_dreambooth_lora_qwenimage21_img2img.py	Updated as part of this pull request.
examples/dreambooth/test_dreambooth_lora_qwenimage21.py	Updated as part of this pull request.
examples/dreambooth/test_dreambooth_lora_qwenimage21_img2img.py	Updated as part of this pull request.
examples/dreambooth/README_qwenimage21.md	Updated as part of this pull request.

Review details

Suppressed comments (3)

examples/dreambooth/train_dreambooth_lora_qwenimage21.py:1772

• When --cache_latents is not enabled, pixel_values remains on the CPU, while vae has been moved to accelerator.device (and offload_models restores it there). This makes vae.encode(pixel_values) fail with a device-mismatch error in the default training path. Move the batch to accelerator.device here, as is already done for the cached path.

                    with offload_models(vae, device=accelerator.device, offload=args.offload):
                        pixel_values = batch["pixel_values"].to(dtype=vae.dtype)
                    model_input = vae.encode(pixel_values).latent_dist.sample()

examples/dreambooth/train_dreambooth_lora_qwenimage21_img2img.py:1852

• The target tensor is not moved to accelerator.device, unlike the condition tensor below it. With the normal GPU setup (and no latent cache), vae is on the accelerator device and vae.encode(pixel_values) therefore fails due to the CPU/GPU mismatch. Add device=accelerator.device to this transfer.

                    with offload_models(vae, device=accelerator.device, offload=args.offload):
                        pixel_values = batch["pixel_values"].to(dtype=vae.dtype)
                        cond_pixel_values = batch["cond_pixel_values"].to(device=accelerator.device, dtype=vae.dtype)
                        model_input = vae.encode(pixel_values).latent_dist.sample()
                        cond_model_input = vae.encode(cond_pixel_values).latent_dist.mode()

examples/dreambooth/train_dreambooth_lora_qwenimage21_img2img.py:1657

• The final-only validation option is not encoded here: this block runs only when --validation_prompt is set. If a user supplies only --final_validation_prompt, the later final-inference path still runs but passes no prompt embeddings to a pipeline whose text encoder is None, so final validation fails instead of generating an image.

    validation_pipeline_args = {}
    validation_image_pad_mask = None
    if args.validation_prompt is not None:
        if args.validation_image is None:
            raise ValueError("`--validation_prompt` needs `--validation_image`, the image the edit is applied to.")
        validation_image = load_image(args.validation_image)
        # Encoded at the size the pipeline will resize to, so vision tokens and latents line up.
        width, height, _ = calculate_dimensions(
            args.resolution * args.resolution, validation_image.size[0] / validation_image.size[1]
        )
        resized_validation_image = validation_image.resize((width, height))
        with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
            embeds, embeds_mask, image_pad_mask = compute_text_embeddings(
                args.validation_prompt, text_encoding_pipeline, resized_validation_image
            )
        validation_pipeline_args = {
            "prompt_embeds": embeds,
            "prompt_embeds_mask": embeds_mask,
            "image": validation_image,

• Files reviewed: 5/5 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[sayakpaul]

Can we add:

• LoRA tests, e.g.

  diffusers/tests/pipelines/flux/test_pipeline_flux.py

  Line 242 in 344d6e7

  class TestFluxPipelineLoRA(FluxPipelineTesterConfig, LoraTesterMixin):

  
  

  diffusers/tests/pipelines/flux/test_pipeline_flux.py

  Line 356 in 344d6e7

  class TestFluxPipelineLoRAMemory(FluxPipelineTesterConfig, LoraMemoryTesterMixin):

• Consolidate i2i and t2i into the same script (optional but desirable)
• Show one success run for T2I and one for I2I

[sayakpaul]

Need to also fix the failing tests (for examples).

[linoytsaban]

@sayakpaul re:the failing checks:

LoRA tests with PEFT main - TestQwenImage21PipelineLoRA::test_simple_inference_with_text_denoiser_lora_and_scale - the scale path and the dummy config both come from #14804.
Fails on CPU, passes on GPU. The assertion is not torch.allclose(..., atol=1e-3, rtol=1e-3), so the bound is atol + rtol*|b| ~= 1.65e-3. On CPU the 2.1 dummy's max diff is 0.001563 - under the bound at every element, so allclose is True and the assert fails. The qwenimage 2.1 tester uses num_attention_heads=2 and the qwenimage tester uses 3, raising it to 3 makes it pass. That's a fixture change to clear a tolerance, should we do that or adjust the shared mixin?

Example tests (CPU + CUDA) - hf-internal-testing/tiny-qwenimage21-pipe is private, so CI gets a 401, the other test is unrelated.

[sayakpaul]

Can we fix the failing tests in https://github.com/huggingface/diffusers/actions/runs/35349539904/job/105614536678?pr=14808 and https://github.com/huggingface/diffusers/actions/runs/35349539897/job/105614547820?pr=14808? For failing LoRA test, does using a higher tolerance help?

I think we also need to verify the effectiveness with one successful run for each of T2I and I2I?