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
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
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
System Info
N/A
Who can help?
No response
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
train_dreambooth_lora_flux2_klein.py
The argument definition specifies:
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
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
System Info
N/A
Who can help?
No response