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
@Dev-X25874
Copy link
Copy Markdown
Contributor Author

Hey @sayakpaul — would love to get your eyes on this when you have a moment! 🙏

This fixes a silent bug where mse_loss_weights from --snr_gamma was never applied to the prior preservation loss in train_dreambooth_lora_sdxl.py — so prior loss always fell back to plain MSE even when SNR weighting was explicitly requested. The fix mirrors the exact same pattern already used for instance loss, and I've added a test covering the --snr_gamma + --with_prior_preservation combination.

Happy to adjust anything based on your feedback!

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