Skip to content

Fixed an issue with abnormal training loss in Dreambooth Flux2. #12774

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tcaimm wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fixed an issue with abnormal training loss in Dreambooth Flux2. #12774

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ba51279
Fixed an issue with abnormal training loss in Dreambooth Flux2.
Dec 2, 2025
ed47a79
Merge branch 'main' into fix/dreambooth-flux2-img2img
tcaimm Dec 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions examples/dreambooth/train_dreambooth_lora_flux2_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1652,6 +1652,8 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):

# concatenate the model inputs with the cond inputs
packed_noisy_model_input = torch.cat([packed_noisy_model_input, packed_cond_model_input], dim=1)
orig_inp_shape = packed_noisy_model_input.shape
orig_inp_ids_shape = model_input_ids.shape
model_input_ids = torch.cat([model_input_ids, cond_model_input_ids], dim=1)

# handle guidance
Expand All @@ -1668,8 +1670,8 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
img_ids=model_input_ids, # B, image_seq_len, 4
return_dict=False,
)[0]
model_pred = model_pred[:, : packed_noisy_model_input.size(1) :]

model_pred = model_pred[:, : orig_inp_shape[1] :]
model_input_ids = model_input_ids[:, :orig_inp_ids_shape[1] :]
model_pred = Flux2Pipeline._unpack_latents_with_ids(model_pred, model_input_ids)

# these weighting schemes use a uniform timestep sampling
Expand Down