huggingface / huggingface/diffusers

Wrong default layers for Flux 2 Klein

Open
#13,445 0 comments 0 reactions 0 assignees View on GitHub
bug lora needs-code-example pipelines
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### Describe the bug

Thanks a lot for taking the time to review this issue 🤗

There seems to be a mismatch between the default values exposed via CLI arguments and the layers that are actually used internally in the script:

Script:
[pipeline_flux2_klein.py](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py)
[train_dreambooth_lora_flux2_klein.py](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_flux2_klein.py)

The argument definition specifies:

```
parser.add_argument(
"--text_encoder_out_layers",
type=int,
nargs="+",
default=[10, 20, 30],
help="Text encoder hidden layers to compute the final text embeddings.",
)
```
However, inside the implementation the following layers are actually used:
`hidden_states_layers: list[int] = (9, 18, 27)`
**Expected behavior**
The default CLI argument values should match the layers that are actually used internally, or the internal implementation should respect the provided CLI values.
**Actual behavior**
There is an off-by-one inconsistency between:
CLI defaults: [10, 20, 30]
Internal usage: (9, 18, 27)
This can lead to confusion and potentially incorrect assumptions when tuning or debugging training behavior.
**Possible cause**
This might be a copy-paste artifact from:
[train_dreambooth_lora_flux2.py](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_flux2.py)
which uses a different setup (e.g. Mistral-based text encoder in a dev version), where layer indexing may differ.
**Additional context**
This discrepancy cost me ~72 GPU hours before I realized what was going on, so I figured it’s worth documenting 😅
**Suggested fix**
Either align defaults with (9, 18, 27)
Or make sure --text_encoder_out_layers is actually used consistently throughout the script

### Reproduction

No special setup required — this is directly visible from reading the script.

### Logs

```shell

```

### System Info

N/A

### Who can help?

_No response_

Contributor guide

Open the contributing guide

Research direction

Start with src/diffusers/pipelines/flux2/pipeline_flux2_klein.py and examples/dreambooth/train_dreambooth_lora_flux2_klein.py, comparing the CLI default for --text_encoder_out_layers with the internal hidden_states_layers value. Trace how the argument reaches the text-embedding computation and verify the intended layer indexing. Done means the defaults and actual layers are consistent, and explicitly provided CLI values are respected if that is the chosen behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.