mmdit

An MMDiTBlock follows a specific sequence to process multiple modalities:

1. Residual Application: Applies the residual function (either Residual or ManifoldConstrainedHyperConnections) to the input tokens.
2. Attention Normalization: Applies AdaptiveLayerNorm to each modality's tokens (using time_cond if available).
3. Joint Attention: Passes all modality tokens through a JointAttention layer, allowing modalities to interact.
4. Post-Attention Scaling: If time_cond is present, it applies learned attn_gammas to scale the tokens.
5. Attention Residual: Adds the attention output back to the original tokens via the residual function.
6. Feedforward Normalization: Applies AdaptiveLayerNorm to the tokens again.
7. Feedforward: Passes tokens through modality-specific FeedForward layers.
8. Post-FF Scaling: If time_cond is present, applies learned ff_gammas to scale the tokens.
9. Feedforward Residual: Adds the feedforward output back to the tokens via the residual function.

The MMDiTBlock implements a single layer of the MMDiT architecture proposed in Stable Diffusion 3. It processes multiple modalities (e.g., text and image) along with a conditioning signal (e.g., time embedding).

Parameters

• dim_cond: Dimension of the conditioning signal.
• dim_text: Dimension of the text modality tokens.
• dim_image: Dimension of the image modality tokens.
• qk_rmsnorm: Boolean flag to enable RMSNorm for queries and keys.

Forward Pass Inputs

• time_cond: Conditioning tensor of shape (batch, dim_cond).
• text_tokens: Text modality tokens of shape (batch, text_seq_len, dim_text).
• text_mask: Boolean mask for text tokens of shape (batch, text_seq_len).
• image_tokens: Image modality tokens of shape (batch, image_seq_len, dim_image).

Returns the updated tokens for both modalities.

import torch
from mmdit import MMDiTBlock

block = MMDiTBlock(
    dim_cond = 256,
    dim_text = 768,
    dim_image = 512,
    qk_rmsnorm = True
)

time_cond = torch.randn(2, 256)
text_tokens = torch.randn(2, 512, 768)
text_mask = torch.ones((2, 512)).bool()
image_tokens = torch.randn(2, 1024, 512)

text_tokens_next, image_tokens_next = block(
    time_cond = time_cond,
    text_tokens = text_tokens,
    text_mask = text_mask,
    image_tokens = image_tokens
)

The MMDiT class from mmdit.mmdit_generalized_pytorch allows for an arbitrary number of modalities, making it suitable for complex multi-modal tasks (e.g., text, video, and audio).

Parameters

• depth: Number of transformer layers.
• dim_modalities: A tuple containing the dimension for each modality in order.
• dim_cond: Dimension of the conditioning signal.
• qk_rmsnorm: Boolean flag to enable RMSNorm for queries and keys.

Forward Pass Inputs

• modality_tokens: A tuple of tensors, where each tensor represents a modality's tokens.
• modality_masks: A tuple of masks corresponding to each modality (use None for unmasked modalities).
• time_cond: Conditioning tensor of shape (batch, dim_cond).

Returns a tuple containing the updated tokens for all input modalities.

import torch
from mmdit.mmdit_generalized_pytorch import MMDiT

mmdit = MMDiT(
    depth = 2, 
    dim_modalities = (768, 512, 384), # e.g., Text, Video, Audio
    dim_cond = 256,
    qk_rmsnorm = True
)

time_cond = torch.randn(2, 256)
text_tokens = torch.randn(2, 512, 768)
text_mask = torch.ones((2, 512)).bool()
video_tokens = torch.randn(2, 1024, 512)
audio_tokens = torch.randn(2, 256, 384)

text_tokens, video_tokens, audio_tokens = mmdit(
    modality_tokens = (text_tokens, video_tokens, audio_tokens),
    modality_masks = (text_mask, None, None),
    time_cond = time_cond,
)

The AdaptiveAttention class implements a self-attention mechanism inspired by adaptive convolutions. It allows for multiple sets of weights (Q, K, V and output projections) that can be dynamically selected or combined per token using a gating mechanism.

Key Parameters

• dim: The input feature dimension.
• dim_head: The dimension of each attention head (default: 64).
• heads: The number of attention heads (default: 8).
• num_adaptive_weights: The number of different weight sets to learn. Setting this to 1 results in regular self-attention. Setting it to > 1 enables gating.
• softclamp: If True, applies a softclamp to the attention similarity scores to prevent extreme values.
• softclamp_value: The value used for clamping if softclamp is enabled (default: 50.0).

Gating Mechanism

When num_adaptive_weights > 1, the module learns a to_gates linear layer that produces a softmax distribution over the adaptive weights for every token. This distribution is used to:

1. Weight the Q, K, and V projections.
2. Weight the final output projection.

Input/Output Shape

• Input: (batch, sequence, dim)
• Output: (batch, sequence, dim * num_adaptive_weights) if gating is enabled, or (batch, sequence, heads * dim_head) if num_adaptive_weights=1.

from mmdit.adaptive_attention import AdaptiveAttention
import torch

# Initialize with 4 adaptive weight sets
adaptive_attn = AdaptiveAttention(
    dim = 512,
    num_adaptive_weights = 4
)

# Input shape: (batch, sequence, dim)
tokens = torch.randn(1, 1408, 512)

# Forward pass
out = adaptive_attn(tokens)
# Output shape: (1, 1408, 512 * 4) = (1, 1408, 2048)

The MMDiTBlock implements a single layer of the MMDiT architecture, featuring joint modality attention and modality-specific feedforward networks. It supports optional time conditioning via AdaptiveLayerNorm.

Key parameters:

• dim_modalities: A tuple of integers representing the dimension of each modality.
• dim_cond: Dimension of the conditioning signal (e.g., time embedding). If set, time_cond must be provided during forward.
• dim_head: Dimension of each attention head.
• heads: Number of attention heads.
• qk_rmsnorm: Boolean for QK RMSNorm (not explicitly used in provided snippet).
• flash_attn: Boolean to enable Flash Attention.
• softclamp: Boolean for softclamping.
• softclamp_value: Value for softclamping.
• num_residual_streams: Number of residual streams (determines if Residual or ManifoldConstrainedHyperConnections is used).
• ff_kwargs: Dictionary of arguments for the FeedForward layers.

block = MMDiTBlock(
    dim_modalities = (512, 768),
    dim_cond = 256,
    dim_head = 64,
    heads = 8
)

# modality_tokens: tuple of Tensors, one per modality
# modality_masks: tuple of optional masks, one per modality
# time_cond: Tensor of shape (batch, dim_cond)
output = block(
    modality_tokens = (tokens_1, tokens_2),
    modality_masks = (mask_1, mask_2),
    time_cond = time_embedding
)

The MMDiT class is a transformer composed of multiple MMDiTBlock layers designed for multi-modal processing. It supports multiple residual streams via ManifoldConstrainedHyperConnections if num_residual_streams > 1.

Key parameters:

• depth: Number of blocks in the transformer.
• dim_modalities: A tuple of integers representing the dimension of each modality (e.g., (512, 768) for two modalities).
• final_norms: Boolean (not explicitly used in the provided logic, but part of the signature).
• num_residual_streams: Number of residual streams to use. If 1, it uses standard Residual connections; otherwise, it uses ManifoldConstrainedHyperConnections.
• mhc_kwargs: Dictionary of arguments passed to ManifoldConstrainedHyperConnections.
• **block_kwargs: Any additional keyword arguments are passed directly to each MMDiTBlock.

model = MMDiT(
    depth = 12,
    dim_modalities = (512, 768),
    num_residual_streams = 4,
    dim_head = 64,
    heads = 8
)

The JointAttention module allows multiple input modalities to attend to each other in a single joint operation. It projects each modality's input to QKV, packs them together, performs attention, and then unpacks/projects them back to their original dimensions.

Arguments:

• dim_inputs: A tuple of input dimensions, e.g., (dim_text, dim_image).
• dim_head: Dimension of each head.
• heads: Number of heads.
• qk_rmsnorm: If True, applies MultiHeadRMSNorm to Q and K per modality.
• flash: If True, uses flash attention.
• softclamp: If True, applies softclamping to logits.
• softclamp_value: The value used for softclamping.
• attend_kwargs: Additional keyword arguments passed to the underlying Attend module.

attn = JointAttention(
    dim_inputs = (768, 1024),
    dim_head = 64,
    heads = 8,
    qk_rmsnorm = True,
    flash = True
)

# inputs is a tuple of tensors, masks is a tuple of optional masks
outs = attn(
    inputs = (torch.randn(1, 77, 768), torch.randn(1, 256, 1024)),
    masks = (torch.ones(1, 77, dtype=torch.bool), None)
)
# outs is a tuple of tensors: (text_out, image_out)

The MMDiT class implements a multi-modal diffusion transformer architecture. It manages a sequence of MMDiTBlocks and supports multiple residual streams via HyperConnections. It can also handle optional register tokens for the image modality and a final normalization step.

Key features:

• Multiple Residual Streams: Controlled by num_residual_streams. If $>1$, it uses HyperConnections to expand and reduce streams.
• Register Tokens: If num_register_tokens > 0, these are prepended to the image_tokens.
• Time Conditioning: Supports adaptive layer normalization via time_cond passed during the forward pass.
• Final Norm: An optional RMSNorm applied to image_tokens at the end of the forward pass.

model = MMDiT(
    depth = 12,
    dim_image = 1024,
    num_register_tokens = 4,
    num_residual_streams = 4,
    dim_text = 768,
    dim_cond = 256,
    dim_head = 64,
    heads = 16
)

text_tokens, image_tokens = model(
    text_tokens = torch.randn(1, 77, 768),
    image_tokens = torch.randn(1, 256, 1024),
    time_cond = torch.randn(1, 256)
)