flux kontext transformer single blocks forward behavior changed

[WingEdge777]

Describe the bug

I observed that this line of code(

diffusers/src/diffusers/models/transformers/transformer_flux.py

Line 88 in 0454fbb

hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)

) was added for First Block Cache.

However, it not only increases the computation, but also thoroughly changes the forward logic; encoder_hidden_states will affect hidden_states because of attention in this and subsequent blocks within the whole denoising loop. Is that right?

I'm also confused by the questions. The FluxSingleTransformerBlock is designed to mix hidden_states with encoder_hidden_states( prompt_embedding, in this case) or not?

Reproduction

no need

Logs

System Info

no need

Who can help?

No response

[DN6]

cc: @a-r-r-o-w

[a-r-r-o-w]

Hi, thanks for the issue. It's a good question.

Actually, the amount of computation remains unchanged. Before the First Block Cache changes, the concatenation was done once, before the calls to single transformer blocks:

diffusers/src/diffusers/models/transformers/transformer_flux.py

Line 510 in cbc8ced

hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)

After FBC changes, the concatenation and split happen within the blocks themselves. This change was made to simplify the implementation of caching techniques, and other algorithms I anticipate will be supported in diffusers. It is a good question whether doing the cat/split at each block impacts performance.

From my tests, the effect on runtime is negligible. This is because cat/split operations are very cheap to perform, and there are many other overheads that are worse than this in eager mode. Even when compiled, the performance difference was negligible.

[WingEdge777]

Hi, thanks for the issue. It's a good question.

Actually, the amount of computation remains unchanged. Before the First Block Cache changes, the concatenation was done once, before the calls to single transformer blocks:

diffusers/src/diffusers/models/transformers/transformer_flux.py

Line 510 in cbc8ced

hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
After FBC changes, the concatenation and split happen within the blocks themselves. This change was made to simplify the implementation of caching techniques, and other algorithms I anticipate will be supported in diffusers. It is a good question whether doing the cat/split at each block impacts performance.

From my tests, the effect on runtime is negligible. This is because cat/split operations are very cheap to perform, and there are many other overheads that are worse than this in eager mode. Even when compiled, the performance difference was negligible.

Thanks for your reply, indeed I missed that line of code before, so it's just been moved to a place(adding num_blocks' cat/slice), the compute logic did not change. Then I understand this makes sense. Looking forward to the next release. First Block Cache's performance looks promising. And I'm closing this issue.