Skip to content

[DreamBooth SDXL LoRA] Apply SNR min-weighting to prior preservation loss when snr_gamma is set#13797

Open
Dev-X25874 wants to merge 2 commits into
huggingface:mainfrom
Dev-X25874:fix/dreambooth-sdxl-snr-prior-loss
Open

[DreamBooth SDXL LoRA] Apply SNR min-weighting to prior preservation loss when snr_gamma is set#13797
Dev-X25874 wants to merge 2 commits into
huggingface:mainfrom
Dev-X25874:fix/dreambooth-sdxl-snr-prior-loss

Conversation

@Dev-X25874
Copy link
Copy Markdown
Contributor

What does this PR do?

In examples/dreambooth/train_dreambooth_lora_sdxl.py, when training with
both --snr_gamma and --with_prior_preservation, the Min-SNR loss
weighting (Section 3.4 of https://huggingface.co/papers/2303.09556) is
applied to the instance loss but silently skipped for the prior preservation
loss. The prior loss is computed before the snr_gamma branch runs, so
mse_loss_weights is never multiplied into it — the prior loss always
falls back to plain MSE regardless of --snr_gamma.

This PR fixes the snr_gamma branch to also apply mse_loss_weights to
prior_loss, mirroring the same pattern already used for the instance loss.
A new test covering the --snr_gamma + --with_prior_preservation combination
is also added to test_dreambooth_lora.py.

Fixes # (issue)

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • 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?
  • Did you write any new necessary tests?

Who can review?

@sayakpaul

@github-actions github-actions Bot added size/S PR with diff < 50 LOC examples and removed size/S PR with diff < 50 LOC labels May 23, 2026
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.

1 participant