Skip to content

Fix prior preservation weighting in advanced DreamBooth LoRA - #14982

Open
Subodh-17 wants to merge 1 commit into
huggingface:mainfrom
Subodh-17:fix/dreambooth-prior-preservation-weighting
Open

Subodh-17 wants to merge 1 commit into
huggingface:mainfrom
Subodh-17:fix/dreambooth-prior-preservation-weighting

Conversation

@Subodh-17

Copy link
Copy Markdown

Description

Fixes the missing prior-preservation weighting split in the advanced DreamBooth LoRA training scripts.

When prior preservation is enabled, model_pred and target are split into instance and prior batches, but weighting was left unsplit. The prior loss therefore attempted to apply a weighting tensor covering both batches to the prior batch only, resulting in a shape mismatch.

This change splits weighting in the same way as model_pred and target and uses weighting_prior for the prior loss.

Changes

  • Chunk weighting into instance and prior portions in the advanced Flux DreamBooth LoRA script.
  • Chunk weighting into instance and prior portions in the advanced SDXL DreamBooth LoRA script.
  • Use weighting_prior when calculating prior-preservation loss.

Scope

This PR is intentionally limited to the missing weighting chunking described in #11489.

Testing

  • python -m py_compile examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py
  • ruff check examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py
  • git diff --check

The full prior-preservation training regression path could not be isolated into a dedicated test because the existing advanced Flux script has a separate pre-existing text_ids issue that occurs earlier in the training path. That unrelated issue is intentionally not modified by this PR.

Fixes #11489

Signed-off-by: Subodh Bansode <subodh171011@gmail.com>
@github-actions github-actions Bot added fixes-issue examples size/S PR with diff < 50 LOC labels Oct 7, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Error when I'm trying to train a Flux lora with train_dreambooth_lora_flux_advanced

1 participant