Fix incorrect batch temporal IDs for cond_model_input in dreambooth flux2 img2img

[Jefsky]

Problem

In train_dreambooth_lora_flux2_img2img.py, Flux2Pipeline._prepare_image_ids is designed for multiple reference images within a single sample, assigning different temporal embeddings (T=10, T=20, T=30...) to distinguish them.

However, cond_model_input has shape (B, C, H, W) where each batch element is an independent training sample with a single conditional image. The old code split the batch into individual items:

cond_model_input_list = [cond_model_input[i].unsqueeze(0) for i in range(cond_model_input.shape[0])]
cond_model_input_ids = Flux2Pipeline._prepare_image_ids(cond_model_input_list)

This produced incorrect cross-sample temporal offsets:

• sample 0 → T=10, sample 1 → T=20, sample 2 → T=30, ...

Fix

Generate temporal IDs for a single sample and expand across the batch dimension:

cond_model_input_ids = Flux2Pipeline._prepare_image_ids([cond_model_input[0:1]])
cond_model_input_ids = cond_model_input_ids.expand(cond_model_input.shape[0], -1, -1)

All batch elements now use the same temporal ID (T=10), which is the correct behavior since each sample only has one conditional image.

Related PRs

• fix(dreambooth): correct batch handling in _prepare_image_ids for flux2 img2img #13821 (previous fix for the same file, different issue)

Fixes #13811