Skip to content

Fix Qwen Image DreamBooth prior-preservation batch ordering#13441

Open
azolotenkov wants to merge 1 commit intohuggingface:mainfrom
azolotenkov:fix-qwen-prior-preservation-batch-order
Open

Fix Qwen Image DreamBooth prior-preservation batch ordering#13441
azolotenkov wants to merge 1 commit intohuggingface:mainfrom
azolotenkov:fix-qwen-prior-preservation-batch-order

Conversation

@azolotenkov
Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes the prior-preservation batching issue in the Qwen Image DreamBooth LoRA script.

PR #13396 corrected the repeat count under --with_prior_preservation, but the script still repeats prompt embeddings with repeat(...), which interleaves the batch as [inst, class, inst, class, ...]. The dataloader and later torch.chunk(..., 2, dim=0) logic assume grouped ordering [inst1..instB, class1..classB].

In addition, weighting was not chunked alongside model_pred and target, which causes a batch mismatch for train_batch_size > 1.

This update:

  • uses repeat_interleave(..., dim=0) for prompt_embeds and prompt_embeds_mask
  • chunks weighting in the prior-preservation loss path

Reproduction:

accelerate launch \
  examples/dreambooth/train_dreambooth_lora_qwen_image.py \
  --pretrained_model_name_or_path hf-internal-testing/tiny-qwenimage-pipe \
  --instance_data_dir /dataset/instance_nocaptions \
  --instance_prompt "photo" \
  --with_prior_preservation \
  --class_data_dir /dataset/class \
  --class_prompt "class photo" \
  --num_class_images 2 \
  --resolution 64 \
  --train_batch_size 2 \
  --max_train_steps 1 \
  --learning_rate 5e-4 \
  --scale_lr \
  --output_dir /tmp/qwen-prior-b2-smoke

This fails with: RuntimeError: The size of tensor a (4) must match the size of tensor b (2)

Before submitting

Who can review?

@sayakpaul

@github-actions github-actions bot added examples size/S PR with diff < 50 LOC labels Apr 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

examples size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant