examples/dreambooth: fix missing weighting chunk when using prior preservation in Flux and SD3 LoRA training

[Dev-X25874]

What does this PR do?

When --with_prior_preservation is enabled, the training batch concatenates
instance and class (prior) samples, so every per-sample tensor —
model_pred, target, sigmas, and therefore weighting — has shape
(2 * train_batch_size, ...).

Inside the loss block, model_pred and target are correctly split via
torch.chunk(..., 2, dim=0), but weighting was never chunked. This means:

• weighting (size 2B) is broadcast against model_pred_prior and
  target_prior (size B), producing a loss tensor of the wrong shape and
  applying incorrectly paired timestep weights to the prior loss term.
• The instance loss term also gets weights from the full unsplit weighting
  instead of only the instance-sample half.

The correct pattern already exists in train_dreambooth_lora_flux2.py:

weighting, weighting_prior = torch.chunk(weighting, 2, dim=0)

This PR applies the same fix to train_dreambooth_lora_flux.py and
train_dreambooth_lora_sd3.py, which were both missing it.

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? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@sayakpaul

[Dev-X25874]

Hi @sayakpaul, would you mind taking a look at this when you get a chance?

The bug is present in both train_dreambooth_lora_flux.py and train_dreambooth_lora_sd3.py — when --with_prior_preservation is enabled, weighting is never chunked alongside model_pred and target, causing incorrect timestep weights to be applied to the prior loss term. The fix already exists in train_dreambooth_lora_flux2.py (line 1832), so this PR simply backports it to the two older scripts. Happy to make any changes if needed!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Dev-X25874]

Thank you @sayakpaul for the review and the quick merge! Really appreciate it.
Looking forward to contributing more! 🙏

[Dev-X25874]

Hi @sayakpaul, hope you're doing well! A quick etiquette question — for PRs that fall under your area, should contributors tag you right when opening the PR? And if not, how long should we wait before tagging you if there's been no activity? Also, do you usually go through all PRs on your own or is an explicit tag always helpful to make sure nothing gets missed? Just want to be respectful of your time. Thanks a lot! 🙏

[sayakpaul]

You can tag me whenever you feel comfortable. But as a sign of respect, waiting till 4-5 days is a good idea I think.