Skip to content

fix: remove incorrect text_ids concatenation in Flux DreamBooth LoRA training - #14983

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

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

Conversation

@Subodh-17

Copy link
Copy Markdown
Contributor

Description

Fixes #10722.

The Flux DreamBooth LoRA training scripts still concatenate text_ids with class_text_ids when prior preservation is enabled.

text_ids is now a 2D tensor containing sequence positional IDs. FluxTransformer2DModel broadcasts this tensor internally across the batch, so concatenating instance and class text_ids along dimension 0 incorrectly doubles the sequence dimension.

This causes a sequence-length mismatch later when the transformer constructs the rotary embeddings.

Changes

Remove the incorrect text_ids concatenation from the static prior-preservation path in:

  • examples/dreambooth/train_dreambooth_lora_flux.py
  • examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py

The resulting text_ids shape matches the format expected by FluxTransformer2DModel and the existing dynamic prompt path.

Background

PR #9653 changed Flux text_ids from a batched tensor to a 2D tensor because the transformer handles batch broadcasting internally. The prior-preservation static caching path retained the older concatenation logic.

Testing

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

The repository currently does not provide an isolated prior-preservation regression test for these training scripts, so no new end-to-end test is included in this minimal fix.

Fixes #10722

Signed-off-by: Subodh Bansode <subodh171011@gmail.com>
@github-actions github-actions Bot added 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

examples size/S PR with diff < 50 LOC

Projects

None yet

1 participant