huggingface / huggingface/diffusers
Wrong default layers for Flux 2 Klein
- 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
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