dew.nn.backbones.sd3

Stable Diffusion 3’s own MM-DiT, as the published transformer computes it.

SimpleMMDiT is Dew’s own dual-stream model, trained from scratch on Dew’s conventions. This module is the other thing: the arithmetic of Diffusers 0.34.0’s SD3Transformer2DModel, so a published checkpoint’s tensors mean here what they mean there. The differences from the scratch model are not cosmetic - the modulation channel order, the joint attention’s image-then- context concatenation, the position buffer’s centred crop, the separate timestep and pooled-text embedders that are summed, the last block’s context-only continuous norm, and SD3.5’s ninefold modulation with a second self-attention - so the two stay separate rather than one growing flags.

The interface is Dew’s: NHWC noisy latents, a model time, a DenoisingCondition carrying the text tokens and the pooled text vector, and NHWC velocity out. Position embeddings are a persistent sin/cos buffer in the source, not a learned parameter, so they ride in the buffers collection: an optimizer never sees them and the checkpoint’s own stored values are what the model reads and what export writes back.

Name	Summary
SD3Block	One JointTransformerBlock.
SD3Transformer	Diffusers 0.34.0’s SD3Transformer2DModel over Dew’s interface.
sincos_position	get_2d_sincos_pos_embed at the grid the source builds its buffer on.

SD3Block

Flax module source

class SD3Block(
    context_pre_only: bool = False,
    dual_attention: bool = False,
    qk_norm: str | None = None,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
)

One JointTransformerBlock.

The image stream modulates, attends jointly with the context, gates, then modulates and runs its feed-forward. The context stream does the same unless it is the last block, where context_pre_only gives it a continuous scale/shift norm, no output projection and no feed-forward. Under use_dual_attention the image modulation has nine pieces and a second self-attention reads a differently modulated copy of the same normalized input.

SD3Transformer

Flax module source

class SD3Transformer(
    patch_size: int = 2,
    in_channels: int = 16,
    out_channels: int = 16,
    num_layers: int = 18,
    heads: int = 18,
    head_dim: int = 64,
    joint_attention_dim: int = 4096,
    caption_projection_dim: int = 1152,
    pooled_projection_dim: int = 2048,
    sample_size: int = 128,
    pos_embed_max_size: int = 96,
    dual_attention_layers: Sequence[int] = (),
    qk_norm: str | None = None,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
)

Diffusers 0.34.0’s SD3Transformer2DModel over Dew’s interface.

__call__ takes NHWC latents, the model time and a DenoisingCondition whose context is the text token states and whose pooled is the pooled text vector, and returns NHWC velocity. The latent grid may be any even rectangle the position buffer covers; the buffer is cropped centred on it, the way the source crops.

SD3Transformer.position

def position(height: int, width: int)

The stored position buffer cropped centred on this patch grid.

sincos_position

function source

def sincos_position(channels: int, grid: int, *, base_size: int)

get_2d_sincos_pos_embed at the grid the source builds its buffer on.

The source meshes width first and then reads that first mesh into the leading half of the channels, so the leading half carries the column coordinate and the trailing half the row, each as sine then cosine over frequencies 10000^-(2i/half); both axes are divided by grid / base_size.

This is only the buffer’s initializer. A published checkpoint stores the buffer and that stored value is what a load reads; this is here so a model built without one starts where the source starts.