FLUX single-stream transformer block
Description
Parallel attention + MLP over the joint [text; image] sequence with a
shared gate: x + gate * proj_out(cat(attn, gelu(mlp))). The
reference concatenates the streams inside every block and splits after;
here the caller concatenates once before the single-block stack, which
is numerically identical. Reference: diffusers
FluxSingleTransformerBlock.
Usage
flux_single_block(dim, num_attention_heads, attention_head_dim, mlp_ratio = 4)
Arguments
dim: Integer. Model dimension.num_attention_heads: Integer. Attention heads.attention_head_dim: Integer. Per-head dimension.mlp_ratio: Numeric. MLP hidden dim multiplier.
Value
Module whose forward(hidden_states, temb, image_rotary_emb) returns the joint hidden states.