huggingface / huggingface/diffusers

Eliminate code ambiguity in models.transformers.FluxAttention

Open
#12,402 1 comment 0 reactions 0 assignees View on GitHub
stale
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

In the initialization of [FluxAttention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py#L319), we determine whether to create to_out based on the parameter pre_only.
However, during inference, we decide whether to call to_out based on whether [encoder_hidden_states](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py#L129) is provided.

This asymmetry can be somewhat confusing for beginners.

Although this does not actually cause any runtime errors.
Only [FluxTransformerBlock](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py#L410) sets pre_only to False, and likewise, only FluxTransformerBlock passes encoder_hidden_states during inference.

Another issue is that context_pre_only appears to be an unused parameter.

I was wondering if it might be better to:
- Remove context_pre_only
- During inference, rely on self.pre_only to decide whether to_out should be called

Contributor guide

Open the contributing guide

Research direction

Start in src/diffusers/models/transformers/transformer_flux.py at FluxAttention initialization and its inference path, then compare how FluxTransformerBlock sets pre_only and passes encoder_hidden_states. Remove the unused context_pre_only parameter and make the inference decision use self.pre_only, while preserving the behavior of FluxTransformerBlock.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.