huggingface / huggingface/diffusers
Fix incorrect batch handling in _prepare_image_ids usage in train_dreambooth_lora_flux2_img2img.py
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
### Describe the bug
## Description
In `train_dreambooth_lora_flux2_img2img.py`, there is a bug in how `_prepare_image_ids` is used for conditional image inputs (`cond_model_input`).
### Problem
`_prepare_image_ids` in `Flux2Pipeline` is designed for **multiple reference images within a single sample**, where each image is assigned a different temporal embedding (e.g., T=10, T=20, T=30...) to distinguish multiple reference images of the same instance.
However, in the training script, `cond_model_input` has shape:
```
(B, C, H, W)
```
i.e., each batch element corresponds to an **independent training sample**, and each sample contains only **one conditional image**.
### Buggy Behavior
In the current implementation (around line ~1703), the code:
1. Splits the batch into a list of single-image tensors
2. Calls `_prepare_image_ids` on the list
3. Produces different temporal IDs per batch index
As a result:
* sample 0 → T=10
* sample 1 → T=20
* sample 2 → T=30
* ...
This is **incorrect**, because batch elements are independent samples and should not have inter-sample temporal relationships.
## Expected Behavior
Each sample in the batch should use the same temporal id (e.g., T=10), since each sample only has one conditional image.
There should be **no cross-sample temporal offset**.
### Reproduction
## Incorrect Implementation
The current implementation is:
```python
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).to(
device=cond_model_input.device
)
cond_model_input_ids = cond_model_input_ids.view(
cond_model_input.shape[0], -1, model_input_ids.shape[-1]
)
```
## Suggested Fix
Instead of computing image IDs per batch element, generate IDs for a single sample and then expand across the batch dimension.
### Fix:
```python
cond_model_input_ids = Flux2Pipeline._prepare_image_ids(
[cond_model_input[0:1]]
).to(device=cond_model_input.device)
# Expand across batch dimension
cond_model_input_ids = cond_model_input_ids.expand(
cond_model_input.shape[0], -1, -1
)
```
### Logs
```shell
```
### System Info
- 🤗 Diffusers version: 0.38.0
- Platform: Linux-5.15.0-72-generic-x86_64-with-glibc2.35
- Running on Google Colab?: No
- Python version: 3.11.13
- PyTorch version (GPU?): 2.5.1+cu121 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 1.15.0
- Transformers version: 5.9.0
- Accelerate version: 1.13.0
- PEFT version: 0.19.1
- Bitsandbytes version: not installed
- Safetensors version: 0.8.0-rc.0
- xFormers version: not installed
- Accelerator: NVIDIA A800-SXM4-80GB, 81920 MiB
- Using GPU in script?:
- Using distributed or parallel set-up in script?:
### Who can help?
_No response_
Contributor guide
Research direction
Open examples/research_projects/flux2/train_dreambooth_lora_flux2_img2img.py and inspect the _prepare_image_ids call around line 1703, then compare it with Flux2Pipeline._prepare_image_ids. Ensure one-sample image IDs are reused across the batch without inter-sample temporal offsets. Done means every batch element has the same conditional-image temporal ID and the resulting shape remains compatible with model_input_ids.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100