Skip to content

Add Flux2 DreamBooth prior preservation tests#13857

Open
azolotenkov wants to merge 2 commits into
huggingface:mainfrom
azolotenkov:test-flux2-dreambooth-prior-batch-size
Open

Add Flux2 DreamBooth prior preservation tests#13857
azolotenkov wants to merge 2 commits into
huggingface:mainfrom
azolotenkov:test-flux2-dreambooth-prior-batch-size

Conversation

@azolotenkov
Copy link
Copy Markdown
Contributor

What does this PR do?

Adds DreamBooth LoRA smoke coverage for Flux2 and Flux2 Klein prior-preservation training with train_batch_size=2.

This covers the prior-preservation path where the effective batch is grouped as [instance samples..., class samples...]. The test helps guard against regressions in prompt/text embedding repetition and prior-loss splitting when the per-device train batch size is greater than 1.

Before submitting

Who can review?

@sayakpaul

Copilot AI review requested due to automatic review settings June 3, 2026 14:24
@github-actions github-actions Bot added examples size/M PR with diff < 200 LOC labels Jun 3, 2026
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds regression tests to ensure DreamBooth LoRA training for Flux2 (and Flux2 Klein) works with prior preservation when train_batch_size=2.

Changes:

  • Added a new prior-preservation test case with train_batch_size=2 to test_dreambooth_lora_flux2.py.
  • Added the equivalent prior-preservation test case to test_dreambooth_lora_flux2_klein.py.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File Description
examples/dreambooth/test_dreambooth_lora_flux2.py Adds a prior-preservation training test using batch size 2 and validates LoRA weights are produced.
examples/dreambooth/test_dreambooth_lora_flux2_klein.py Adds the same prior-preservation batch size 2 test for the Flux2 Klein variant.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +121 to +124
--with_prior_preservation
--class_data_dir {self.instance_data_dir}
--class_prompt dog
--num_class_images 2
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I updated both tests to use a dedicated temporary class_data directory and copy two fixture images into it, so the tests no longer reuse the instance data directory for class images.

Comment thread examples/dreambooth/test_dreambooth_lora_flux2_klein.py
Comment thread examples/dreambooth/test_dreambooth_lora_flux2.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

examples size/M PR with diff < 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants