huggingface / huggingface/diffusers

Fix incorrect batch handling in _prepare_image_ids usage in train_dreambooth_lora_flux2_img2img.py

Aperta Adatta ai principianti
#13,811 1 commento 0 reazioni 0 assegnatari Vedi su GitHub
bug
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### 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_

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Apri examples/research_projects/flux2/train_dreambooth_lora_flux2_img2img.py e analizza la chiamata a _prepare_image_ids intorno alla riga 1703, quindi confrontala con Flux2Pipeline._prepare_image_ids. Assicurati che gli ID immagine di un campione vengano riutilizzati nell’intero batch senza offset temporali tra i campioni. Il lavoro è completato quando ogni elemento del batch ha lo stesso ID temporale dell’immagine condizionale e la forma risultante rimane compatibile con model_input_ids.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
2/5
Tempo stimato
1-3 ore
Stato di attività
Tranquilla
Chiarezza
Specificata chiaramente
Idoneità per principianti
72/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.