I am studying the code related to Flux. Could you please explain the underlying logic behind this core part of the code?

[deadsmither5]

I’m curious about why you chose the dimensions below. Looking forward to your reply. Thank you!

(1,24,4608,4608) -> (1,24,4096,512)#?why do this?

    attention_probs = attention_probs[:,:,text_length:,:text_length].cpu()

[wooyeolbaek]

@deadsmither5
SD3 flattens the text hidden state and image hidden state and concatenates them for computation. Therefore, it was used to extract only the cross-attention map from the entire attention map.

[luchaoqi]

@wooyeolbaek
I'm wondering if the current version only stores the attention map of cross-attention and no self-attention in the FLUX model. Is there any way to extract self-attention from the FLUX model as well?

[wooyeolbaek]

@luchaoqi Sure, you can get the self-attention for image sequences by indexing with attention_probs = attention_prob[:, :, text_length:, text_length:]

[luchaoqi]

@luchaoqi Sure, you can get the self-attention for image sequences by indexing with attention_probs = attention_prob[:, :, text_length:, text_length:]

what about the self-attention part from the FluxSingleTransformerBlock? I guess currently it only extracts 19 blocks regarding the cross-attention for dualTransformerBlock but not attention from the 38 FluxSingleTransformerBlock?

I believe the current code detects processors for all transformer blocks in register_cross_attention_hook but only records 19 attentions with something like transformer_blocks.0~18.attn. Did I miss anything here?