Rationale behind converting proj_out of FluxSingleTransformerBlock to ConcatLinear

[vinovo]

in this code

We can see only proj_out of FluxSingleTransformerBlock are converted to ConcatLinear with only a single split [module.proj_out.out_features].

Can anyone help explain the reasoning behind this?

• Why we only cares about proj_out of FluxSingleTransformerBlock. Should this operation be performed on all transformer based diffusion models?
• Why do we only care to create one single split? What is the benefit of using ConcatLinear?

Thank you in advance.

[synxlin]

Hi @vinovo ,

FluxSingleTransformerBlock parallelizes Attention and Feed Forward Network (FFN), and thus fuses the to_out linear layer in Attention Block and second linear layer in FFN into one proj_out layer. We split this proj_out layer back to two linear layers, one used in Attention Block and the other used in FC layer, following the Nunchaku implementation.