huggingface / huggingface/diffusers

Fix incorrect batch handling in _prepare_image_ids usage in train_dreambooth_lora_flux2_img2img.py

Offen Anfängerfreundlich
#13,811 1 Kommentar 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

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

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Öffne examples/research_projects/flux2/train_dreambooth_lora_flux2_img2img.py und untersuche den Aufruf von _prepare_image_ids um Zeile 1703. Vergleiche ihn anschließend mit Flux2Pipeline._prepare_image_ids. Stelle sicher, dass die Bild-IDs eines Samples im gesamten Batch ohne temporale Offsets zwischen den Samples wiederverwendet werden. Als erledigt gilt die Aufgabe, wenn jedes Batch-Element dieselbe temporale ID des konditionierten Bildes hat und die resultierende Form weiterhin mit model_input_ids kompatibel ist.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning
Issue-Typ
Bug
Schwierigkeit
2/5
Geschätzter Aufwand
1-3 Stunden
Aktivitätsstatus
Ruhig
Klarheit
Klar beschrieben
Anfängerfreundlichkeit
72/100

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.