flux_transformer

FLUX transformer model

Description

19 double-stream (MMDiT) blocks followed by 38 single-stream blocks over the joint [text; image] sequence, with adaLN-Zero conditioning on timestep + pooled CLIP text. Rotary embeddings are precomputed by the caller with flux_pos_embed (they are static across denoise steps). Defaults are the FLUX.1-schnell configuration.

Usage

flux_transformer(in_channels = 64L, num_layers = 19L, num_single_layers = 38L,
                 attention_head_dim = 128L, num_attention_heads = 24L,
                 joint_attention_dim = 4096L, pooled_projection_dim = 768L,
                 axes_dims_rope = c(16L, 56L, 56L), out_channels = NULL)

Arguments

in_channels: Integer. Packed latent channels (64).
num_layers: Integer. Double-stream block count.
num_single_layers: Integer. Single-stream block count.
attention_head_dim: Integer. Per-head dimension.
num_attention_heads: Integer. Attention heads.
joint_attention_dim: Integer. T5 embedding dim (4096).
pooled_projection_dim: Integer. CLIP pooled dim (768).
axes_dims_rope: Integer vector. Per-axis rotary dims.
out_channels: Integer or NULL. Output channels (defaults to in_channels).

Value

Module whose forward(hidden_states, encoder_hidden_states, pooled_projections, timestep, image_rotary_emb) returns the predicted velocity for the image tokens [B, S_img, out_channels]. timestep is in sigma space (0-1); it is scaled by 1000 internally, matching the reference.

Packages