dew.nn.backbones.flux

Flux’s transformer, as Diffusers 0.34.0’s FluxTransformer2DModel runs it.

The published model reads a packed latent: 2x2 patches already folded into the channel axis by its pipeline, one position per patch. Its stack has two halves. The double-stream blocks carry the image and the text in separate residual streams with their own modulation and feed-forwards, joined only inside attention; the single-stream blocks concatenate the two and run one stream whose attention and feed-forward share a projection. Every attention call rotates its queries and keys with a three-axis rotary table over the text and image ids, in the interleaved-real form Flux uses.

SD3Transformer and this module share the modulation, the feed-forward and the layer norm the MM-DiT family uses; what differs is here.

Name	Summary
FluxBlock	One double-stream block: two modulated residual streams, joined in attention and separate through their feed-forwards.
FluxSingleBlock	One single-stream block: attention and a feed-forward over the joined sequence, whose outputs are concatenated and projected together.
FluxTransformer	Diffusers 0.34.0’s FluxTransformer2DModel over Dew’s interface.
apply_rotary	The rotation apply_rotary_emb applies with use_real_unbind_dim=-1: adjacent channels are one complex pair, so the rotated copy is (-x1, x0) within each pair.
flux_positions	The ids a Flux pipeline lays out: the text at the origin, then one position per packed patch, indexed by its row and column.
pack	_prepare_latents’ packing: one position per 2x2 patch, whose channels run channel-major over the patch’s own two rows and columns.
rotary_table	FluxPosEmbed over ids: one angle per channel pair, per axis.
unpack	_unpack_latents, back to NHWC.

FluxBlock

Flax module source

class FluxBlock(
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
)

One double-stream block: two modulated residual streams, joined in attention and separate through their feed-forwards.

FluxSingleBlock

Flax module source

class FluxSingleBlock(
    mlp_ratio: float = 4.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
)

One single-stream block: attention and a feed-forward over the joined sequence, whose outputs are concatenated and projected together.

FluxTransformer

Flax module source

class FluxTransformer(
    patch_size: int = 1,
    in_channels: int = 64,
    out_channels: int = 64,
    num_layers: int = 19,
    num_single_layers: int = 38,
    heads: int = 24,
    head_dim: int = 128,
    joint_attention_dim: int = 4096,
    pooled_projection_dim: int = 768,
    guidance_embeds: bool = False,
    axes_dims_rope: Sequence[int] = (16, 56, 56),
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
)

Diffusers 0.34.0’s FluxTransformer2DModel over Dew’s interface.

__call__ takes NHWC latents, the model time the schedule supplies - the sigma times the training count, which is the product the source reaches by dividing its timestep and multiplying it back - and a DenoisingCondition whose context is the T5 token states, whose pooled is the CLIP pooled vector and whose guidance is the distilled guidance value a guidance-embedded checkpoint reads. The 2x2 packing its pipeline performs is here, so a caller works in latents and the position ids follow the latent grid.

apply_rotary

function source

def apply_rotary(x: jax.Array, cos: jax.Array, sin: jax.Array) -> jax.Array

The rotation apply_rotary_emb applies with use_real_unbind_dim=-1: adjacent channels are one complex pair, so the rotated copy is (-x1, x0) within each pair.

flux_positions

function source

def flux_positions(rows: int, columns: int, text: int) -> np.ndarray

The ids a Flux pipeline lays out: the text at the origin, then one position per packed patch, indexed by its row and column.

pack

function source

def pack(x: jax.Array) -> jax.Array

_prepare_latents’ packing: one position per 2x2 patch, whose channels run channel-major over the patch’s own two rows and columns.

rotary_table

function source

def rotary_table(
    positions: np.ndarray,
    axes: Sequence[int],
    *,
    theta: float = 10000.0,
) -> tuple[np.ndarray, np.ndarray]

FluxPosEmbed over ids: one angle per channel pair, per axis.

The source computes its frequencies and their cosines in float64 from a static id grid, so this is the same host computation, and the pair of channels that shares an angle is adjacent rather than half a width apart.

unpack

function source

def unpack(x: jax.Array, rows: int, columns: int) -> jax.Array

_unpack_latents, back to NHWC.