Implementation:AUTOMATIC1111 Stable diffusion webui MMDiT Architecture

Knowledge Sources	AUTOMATIC1111_Stable_diffusion_webui
Domains	Diffusion Models, Transformer Architecture, Stable Diffusion 3
Last Updated	2025-05-15 00:00 GMT

Overview

Implements the MM-DiT (Multi-Modal Diffusion Transformer) architecture, which is the core denoising backbone of Stable Diffusion 3, using joint attention blocks that process both image latent patches and text conditioning tokens together.

Description

The module defines the complete MM-DiT transformer architecture through the following components:

• PatchEmbed: Converts 2D image latents into patch embeddings using a convolutional projection layer, similar to Vision Transformer (ViT) patch embedding.

• TimestepEmbedder and VectorEmbedder: Embed scalar timesteps and flat conditioning vectors into high-dimensional representations using sinusoidal encodings and MLPs.

• SelfAttention: Multi-head self-attention with support for QK normalization (RMSNorm or LayerNorm) and multiple attention backends (xformers, torch, math).

• DismantledBlock: A DiT block with gated adaptive layer norm (adaLN) conditioning that separates pre-attention and post-attention phases, enabling the joint attention pattern.

• JointBlock: The core building block that applies block_mixing to jointly attend over context (text) and x (image) streams by concatenating their QKV projections before computing shared attention, then splitting the results back.

• FinalLayer: Projects the transformer output back to patch space with adaptive layer norm modulation.

• MMDiT: The top-level model that orchestrates patchification, 2D sinusoidal positional embedding (with cropping support for variable resolutions), timestep/vector embedding, a stack of JointBlocks, and unpatchification.

The hidden size is computed as 64 * depth, giving a head dimension of 64 across all configurations.

Usage

Use this module as the core denoising backbone for Stable Diffusion 3 inference. The MMDiT is instantiated by the BaseModel wrapper in sd3_impls.py and processes noisy latent patches with text conditioning through the joint attention transformer stack during each denoising step.

Code Reference

Source Location

• Repository: AUTOMATIC1111_Stable_diffusion_webui
• File: modules/models/sd3/mmdit.py
• Lines: 1-622

Signature

class MMDiT(nn.Module):
    def __init__(self, input_size=32, patch_size=2, in_channels=4,
                 depth=28, mlp_ratio=4.0, learn_sigma=False,
                 adm_in_channels=None, context_embedder_config=None,
                 register_length=0, attn_mode="torch", rmsnorm=False,
                 scale_mod_only=False, swiglu=False, out_channels=None,
                 pos_embed_scaling_factor=None, pos_embed_offset=None,
                 pos_embed_max_size=None, num_patches=None,
                 qk_norm=None, qkv_bias=True, dtype=None, device=None):

    def forward(self, x: torch.Tensor, t: torch.Tensor,
                y: Optional[torch.Tensor] = None,
                context: Optional[torch.Tensor] = None) -> torch.Tensor:

Import

from modules.models.sd3.mmdit import MMDiT

I/O Contract

Inputs

Name	Type	Required	Description
x	torch.Tensor	Yes	Spatial input tensor (N, C, H, W) of images or latent representations
t	torch.Tensor	Yes	Diffusion timestep tensor (N,)
y	torch.Tensor	No	Class/vector conditioning tensor (N, adm_in_channels)
context	torch.Tensor	No	Text conditioning context tensor (N, L, D) from text encoders

Outputs

Name	Type	Description
output	torch.Tensor	Denoised prediction tensor (N, out_channels, H, W) after unpatchification

Usage Examples

from modules.models.sd3.mmdit import MMDiT

# Create an MMDiT model (typically done by BaseModel from state_dict shapes)
model = MMDiT(
    input_size=None,
    patch_size=2,
    in_channels=16,
    depth=24,
    pos_embed_max_size=192,
    num_patches=36864,
    adm_in_channels=2048,
    context_embedder_config={
        "target": "torch.nn.Linear",
        "params": {"in_features": 4096, "out_features": 1536}
    },
    dtype=torch.float16,
    device="cuda"
)

# Forward pass
x = torch.randn(1, 16, 64, 64, dtype=torch.float16, device="cuda")
t = torch.tensor([500.0], device="cuda")
y = torch.randn(1, 2048, dtype=torch.float16, device="cuda")
context = torch.randn(1, 154, 4096, dtype=torch.float16, device="cuda")

output = model(x, t, y=y, context=context)
# output shape: (1, 16, 64, 64)

Related Pages

• Principle:AUTOMATIC1111_Stable_diffusion_webui_Transformer_Diffusion_Architecture