dew.nn.backbones

Model backbones.

Each class registers itself with dew.registry.models where it is defined, so models["simple_dit"], models.SimpleDiT and the class are one object and training and inference build the same model from a logged name and fields. The classes are exported for direct use in notebooks and tests.

Name	Summary
CausalTransformer	Decoder-only transformer over token ids: [B, S] int32 -> [B, S, vocab] fp32. Documented in dew.nn.backbones.causal_transformer.
FluxTransformer	Diffusers 0.34.0’s FluxTransformer2DModel over Dew’s interface. Documented in dew.nn.backbones.flux.
HierarchicalMMDiT	U-shaped MM-DiT: dual-stream blocks per stage with patch merging on the way down and expansion + skip fusion on the way up.
HybridSSMAttentionDiT	DiT that interleaves SSM blocks with attention blocks.
QwenImageTransformer	QwenImage21Transformer2DModel over Dew’s interface.
SD3Transformer	Diffusers 0.34.0’s SD3Transformer2DModel over Dew’s interface. Documented in dew.nn.backbones.sd3.
SimpleDiT	Standard DiT: a plain stack of adaLN-Zero attention blocks.
SimpleMMDiT	SD3-style MM-DiT: a plain stack of dual-stream blocks.
SimpleUDiT	A U-shaped DiT: SimpleDiT’s adaLN-Zero blocks with the first half’s outputs skipping into the second half through a dense layer over the concatenation.
UNet2DCondition	Denoise NHWC latents on text, pooled/size and optional inpaint conditions.
UNet3D	Video UNet over (B, T, H, W, C): the 2D Unet body per frame, with a TemporalBlock at every resolution level.
UNetStage	Describes one resolution level: its width, heads and transformer depth.
UViT	Denoise patches as U-ViT does (Bao et al.
Unet	A convolutional UNet with residual blocks and cross-attention stages.
VideoDiT	Factorized spatial-temporal DiT over (B, T, H, W, C) inputs.

HierarchicalMMDiT

Flax module source

class HierarchicalMMDiT(
    output_channels: int = 3,
    base_patch_size: int = 8,
    emb_features: Sequence[int] = (512, 768, 1024),
    num_layers: Sequence[int] = (4, 4, 14),
    num_heads: Sequence[int] = (8, 12, 16),
    mlp_ratio: int = 4,
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
    remat: RematChoice = False,
)

U-shaped MM-DiT: dual-stream blocks per stage with patch merging on the way down and expansion + skip fusion on the way up.

Raster order only: the merge/expand grid reshapes assume row-major token order, so a hilbert scan would scramble the neighborhoods being merged.

HierarchicalMMDiT.stage_blocks

def stage_blocks(stage: int, prefix: str) -> list

Build one stage’s MMDiT blocks, at that stage’s width and heads.

HierarchicalMMDiT.encoder_path

def encoder_path(num_stages: int)

Build the encoder, fine to coarse: each stage’s blocks and its merger.

Sets encoder_blocks, one list per stage, and patch_mergers, one between each pair of stages.

HierarchicalMMDiT.decoder_path

def decoder_path(num_stages: int)

Build the decoder, coarse to fine, for stages N-2 down to 0.

Sets patch_expanders, the fusion_layers that take an expanded stage beside its skip, and decoder_blocks. All three are indexed by decoder step, not by stage.

HierarchicalMMDiT.setup

def setup()

Build the patch embedding, the two paths of stages, and the head.

cond_projs and txt_embeds carry the conditioning and the text into each stage’s own width. encoder_path and decoder_path set the blocks and the resampling layers between them.

HybridSSMAttentionDiT

Flax module source

class HybridSSMAttentionDiT(
    output_channels: int = 3,
    patch_size: int = 16,
    emb_features: int = 768,
    num_layers: int = 12,
    num_heads: int = 12,
    mlp_ratio: int = 4,
    ssm_state_dim: int = 64,
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
    remat: RematChoice = False,
    scan_order: Literal['raster', 'hilbert', 'zigzag'] = 'raster',
    block_pattern: Sequence[str] | None = None,
    ssm_attention_ratio: str = DEFAULT_SSM_RATIO,
    bidirectional_ssm: bool = True,
    use_2d_fusion: bool = False,
)

DiT that interleaves SSM blocks with attention blocks.

The mixer of every layer comes from ssm_attention_ratio, a shorthand that reads the same at any depth (“3:1”, “all-ssm”), or from block_pattern, which names each layer. Setting both raises a ValueError at setup.

HybridSSMAttentionDiT.block

def block(index: int, block_type: str) -> ModulatedBlock

Build layer index’s block, an SSM mixer or an attention one.

The two share the width, the MLP ratio and the norms. They differ in the mixer they name, the fields only that mixer reads, the remat policy, and the name a checkpoint stores them under.

HybridSSMAttentionDiT.setup

def setup()

QwenImageTransformer

Flax module source

class QwenImageTransformer(
    in_channels: int = 64,
    out_channels: int = 64,
    num_layers: int = 32,
    heads: int = 32,
    head_dim: int = 128,
    context_in_dim: int = 4096,
    mlp_ratio: int = 3,
    axes_dims_rope: Sequence[int] = (16, 56, 56),
    eps: float = 1e-06,
    causal_condition: bool = True,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
)

QwenImage21Transformer2DModel over Dew’s interface.

__call__ takes NHWC latents, one token per position, the model time the schedule supplies - the sigma times the training count, which is the product the source reaches by dividing its timestep by a thousand and multiplying it back - and a DenoisingCondition whose context is the encoder’s text states after the system prompt, right-padded, with mask marking the real ones. Each row’s real tokens lead its text, the layout the rotary positions and the attention’s key lengths both read. It returns the prediction at the image’s tokens.

SimpleDiT

Flax module source

class SimpleDiT(
    output_channels: int = 3,
    patch_size: int = 16,
    emb_features: int = 768,
    num_layers: int = 12,
    num_heads: int = 12,
    mlp_ratio: int = 4,
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
    remat: RematChoice = False,
    scan_order: Literal['raster', 'hilbert', 'zigzag'] = 'raster',
)

Standard DiT: a plain stack of adaLN-Zero attention blocks.

SimpleDiT.setup

def setup()

SimpleMMDiT

Flax module source

class SimpleMMDiT(
    output_channels: int = 3,
    patch_size: int = 16,
    emb_features: int = 768,
    num_layers: int = 12,
    num_heads: int = 12,
    mlp_ratio: int = 4,
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
    remat: RematChoice = False,
    scan_order: Literal['raster', 'hilbert', 'zigzag'] = 'raster',
)

SD3-style MM-DiT: a plain stack of dual-stream blocks.

SimpleMMDiT.setup

def setup()

SimpleUDiT

Flax module source

class SimpleUDiT(
    output_channels: int = 3,
    patch_size: int = 16,
    emb_features: int = 768,
    num_layers: int = 12,
    num_heads: int = 12,
    mlp_ratio: int = 4,
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    attention_impl: str = 'auto',
    remat: RematChoice = False,
    norm_epsilon: float = 1e-05,
    scan_order: Literal['raster', 'hilbert'] = 'raster',
)

A U-shaped DiT: SimpleDiT’s adaLN-Zero blocks with the first half’s outputs skipping into the second half through a dense layer over the concatenation. Position comes from RoPE over the sequence index, so a hilbert scan carries the rotation of its curve index and no 2D signal.

SimpleUDiT.setup

def setup()

UNet2DCondition

Flax module source

class UNet2DCondition(
    in_channels: int = 4,
    out_channels: int = 4,
    blocks_per_level: int = 2,
    linear_projection: bool = False,
    additional_time_features: int = 0,
    middle_attention: bool = True,
    frequency_shift: float = 0,
    cosine_first: bool = True,
    dropout: float = 0.0,
    dtype: Dtype = jnp.float32,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
    norm_groups: int = 32,
    norm_epsilon: float = 1e-05,
    attention_norm_epsilon: float = 1e-05,
    approximate_gelu: bool = True,
)

Denoise NHWC latents on text, pooled/size and optional inpaint conditions.

Stages specify spatial width, attention heads and transformer depth. The encoder saves each residual output and each downsample; the decoder consumes those skips in reverse order with one extra residual block per level.

UNet3D

class source

class UNet3D(Unet)

Video UNet over (B, T, H, W, C): the 2D Unet body per frame, with a TemporalBlock at every resolution level. Spatial param paths are identical to Unet, so 2D checkpoints inflate directly.

UNetStage

dataclass source

class UNetStage(
    features: int,
    heads: int,
    depth: int = 1,
    cross_attention: bool = True,
    cross_only: bool = False,
)

Describes one resolution level: its width, heads and transformer depth.

cross_attention gives the level a spatial attention block at all; cross_only drops the self-attention inside it.

UViT

Flax module source

class UViT(
    output_channels: int = 3,
    patch_size: int = 16,
    emb_features: int = 768,
    num_layers: int = 12,
    num_heads: int = 12,
    use_projection: bool = False,
    use_self_and_cross: bool = False,
    force_fp32_for_softmax: bool = True,
    attention_impl: str = 'auto',
    activation: Callable = jax.nn.swish,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    add_residualblock_output: bool = False,
    norm_inputs: bool = False,
    explicitly_add_residual: bool = True,
    norm_epsilon: float = 1e-05,
    scan_order: Literal['raster', 'hilbert'] = 'raster',
)

Denoise patches as U-ViT does (Bao et al. 2023).

The time embedding and the text are tokens beside the patches, the blocks are plain transformer blocks, and each first-half output skips into the second half through a dense layer over the concatenation. Position is a learned table over the raster index, sized for a 512 pixel image.

add_residualblock_output refines the unpatchified prediction with two convolutions over it and the input image. use_projection, use_self_and_cross, norm_inputs and explicitly_add_residual are TransformerBlock’s own settings, passed to every block.

UViT.setup

def setup()

Unet

Flax module source

class Unet(
    output_channels: int = 3,
    emb_features: int = 64 * 4,
    feature_depths: Sequence[int] = (64, 128, 256, 512),
    attention_configs: Sequence[Stage | None] = (Stage(heads=8), Stage(heads=8), Stage(heads=8), Stage(heads=8)),
    num_res_blocks: int = 2,
    num_middle_res_blocks: int = 1,
    activation: Callable = jax.nn.swish,
    norm_groups: int = 8,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
)

A convolutional UNet with residual blocks and cross-attention stages.

Without text the attention stages self-attend (TransformerBlock’s context defaults to its input). Cross-attention reads the whole text sequence; the mask the DiT family pools with has no place here.

attention_configs: Sequence[Stage | None]

Attention per resolution stage, one entry per feature depth; None is a stage with no attention.

Unet.setup

def setup()

VideoDiT

Flax module source

class VideoDiT(
    output_channels: int = 3,
    patch_size: int = 16,
    emb_features: int = 768,
    num_layers: int = 12,
    num_heads: int = 12,
    mlp_ratio: int = 4,
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
    remat: RematChoice = False,
    scan_order: Literal['raster', 'hilbert', 'zigzag'] = 'raster',
)

Factorized spatial-temporal DiT over (B, T, H, W, C) inputs.

VideoDiT.setup

def setup()