flux2_single_block

FLUX.2 single-stream block (parallel)

Description

Parameterless LayerNorm with external modulation, then the fused parallel attention+MLP. Operates on the pre-concatenated [text; image] sequence (the reference model concatenates once before the stack). Reference: Flux2SingleTransformerBlock.

Usage

flux2_single_block(dim, num_attention_heads, attention_head_dim, mlp_ratio = 3,
                   eps = 1e-06, bias = FALSE)

Arguments

  • dim: Integer. Model dimension.
  • num_attention_heads: Integer. Attention heads.
  • attention_head_dim: Integer. Per-head dimension.
  • mlp_ratio: Numeric. MLP multiplier (FLUX.2: 3.0).
  • eps: Numeric. Norm epsilon.
  • bias: Logical.

Value

Module whose forward(hidden_states, temb_mod, image_rotary_emb) returns the joint hidden states.