From df3624663f890b659c0e6203a6dfa3887ba4bc27 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Tue, 6 Jan 2026 17:04:48 +0100 Subject: [PATCH 01/13] add batch dimension to RoPE-Nd --- .gitignore | 1 + src/diffulab/networks/denoisers/ddt.py | 70 ++++++++++++++---------- src/diffulab/networks/denoisers/mmdit.py | 26 ++++++--- src/diffulab/networks/utils/nn.py | 36 ++++++------ 4 files changed, 77 insertions(+), 56 deletions(-) diff --git a/.gitignore b/.gitignore index cd275a9..4a86014 100644 --- a/.gitignore +++ b/.gitignore @@ -104,6 +104,7 @@ venv/ ENV/ env.bak/ venv.bak/ +.vscode/ # Spyder project settings .spyderproject diff --git a/src/diffulab/networks/denoisers/ddt.py b/src/diffulab/networks/denoisers/ddt.py index 41c89bb..8b1ce89 100644 --- a/src/diffulab/networks/denoisers/ddt.py +++ b/src/diffulab/networks/denoisers/ddt.py @@ -313,7 +313,7 @@ def encode_mmddt( context = self.context_embed(context) attn_mask = context_output.get("attn_mask", None) - # pos_ids: [S, n_axes] positional IDs along each axis for rope + # pos_ids: [B, S, n_axes] positional IDs along each axis for rope # in mmdit attention we concat with context first. Context have 0,0 for h w # text: (t>0, 0, 0) text_pos_ids = torch.stack( @@ -336,7 +336,7 @@ def encode_mmddt( dim=-1, ).view(-1, 3) - pos_ids = torch.cat([text_pos_ids, img_pos_ids], dim=0) + pos_ids = torch.cat([text_pos_ids, img_pos_ids], dim=0).unsqueeze(0).repeat(x.size(0), 1, 1) cos_sin_rope = get_cos_sin_ndim_grid(pos_ids, base=self.rope_base, axes_dim=self.rope_axes_dim) features: list[Tensor] | None = [] if intermediate_features else None @@ -379,14 +379,22 @@ def encode_ddt( if self.label_embed is not None: emb = emb + self.label_embed(y, p) - # pos_ids: [S, n_axes] positional IDs along each axis for rope - pos_ids = torch.stack( - torch.meshgrid( - [torch.arange(self.grid_size[0], device=x.device), torch.arange(self.grid_size[1], device=x.device)], - indexing="ij", - ), - dim=-1, - ).view(-1, 2) + # pos_ids: [B, S, n_axes] positional IDs along each axis for rope + pos_ids = ( + torch.stack( + torch.meshgrid( + [ + torch.arange(self.grid_size[0], device=x.device), + torch.arange(self.grid_size[1], device=x.device), + ], + indexing="ij", + ), + dim=-1, + ) + .view(-1, 2) + .unsqueeze(0) + .repeat(x.size(0), 1, 1) + ) cos_sin_rope = get_cos_sin_ndim_grid(pos_ids, base=self.rope_base, axes_dim=self.rope_axes_dim) features: list[Tensor] | None = [] if intermediate_features else None @@ -421,28 +429,32 @@ def decode( emb = self.time_embed(timestep_embedding(timesteps, self.frequency_embedding))[:, None, :] encoder_output = nn.functional.silu(encoder_output + emb) - # pos_ids: [S, n_axes] positional IDs along each axis for rope + # pos_ids: [B, S, n_axes] positional IDs along each axis for rope pos_ids = ( - torch.stack( - torch.meshgrid( - [ + ( + torch.stack( + torch.meshgrid( + [ + torch.arange(self.grid_size[0], device=x.device), + torch.arange(self.grid_size[1], device=x.device), + ], + indexing="ij", + ), + dim=-1, + ).view(-1, 2) + if self.simple_ddt + else torch.stack( + torch.meshgrid( + torch.zeros(1, device=x.device, dtype=torch.long), torch.arange(self.grid_size[0], device=x.device), torch.arange(self.grid_size[1], device=x.device), - ], - indexing="ij", - ), - dim=-1, - ).view(-1, 2) - if self.simple_ddt - else torch.stack( - torch.meshgrid( - torch.zeros(1, device=x.device, dtype=torch.long), - torch.arange(self.grid_size[0], device=x.device), - torch.arange(self.grid_size[1], device=x.device), - indexing="ij", - ), - dim=-1, - ).view(-1, 3) + indexing="ij", + ), + dim=-1, + ).view(-1, 3) + ) + .unsqueeze(0) + .repeat(x.size(0), 1, 1) ) cos_sin_rope = get_cos_sin_ndim_grid(pos_ids, base=self.rope_base, axes_dim=self.rope_axes_dim) diff --git a/src/diffulab/networks/denoisers/mmdit.py b/src/diffulab/networks/denoisers/mmdit.py index cad43f4..ff961a4 100644 --- a/src/diffulab/networks/denoisers/mmdit.py +++ b/src/diffulab/networks/denoisers/mmdit.py @@ -751,7 +751,7 @@ def mmdit_forward( dim=-1, ).view(-1, 3) - pos_ids = torch.cat([text_pos_ids, img_pos_ids], dim=0) + pos_ids = torch.cat([text_pos_ids, img_pos_ids], dim=0).unsqueeze(0).repeat(x.size(0), 1, 1) cos_sin_rope = get_cos_sin_ndim_grid(pos_ids, base=self.rope_base, axes_dim=self.rope_axes_dim) features: list[Tensor] | None = [] if intermediate_features else None @@ -786,14 +786,22 @@ def simple_dit_forward( if self.label_embed is not None: emb = emb + self.label_embed(y, p) - # pos_ids: [S, n_axes] positional IDs along each axis for rope - pos_ids = torch.stack( - torch.meshgrid( - [torch.arange(self.grid_size[0], device=x.device), torch.arange(self.grid_size[1], device=x.device)], - indexing="ij", - ), - dim=-1, - ).view(-1, 2) + # pos_ids: [B, S, n_axes] positional IDs along each axis for rope + pos_ids = ( + torch.stack( + torch.meshgrid( + [ + torch.arange(self.grid_size[0], device=x.device), + torch.arange(self.grid_size[1], device=x.device), + ], + indexing="ij", + ), + dim=-1, + ) + .view(-1, 2) + .unsqueeze(0) + .repeat(x.size(0), 1, 1) + ) cos_sin_rope = get_cos_sin_ndim_grid(pos_ids, base=self.rope_base, axes_dim=self.rope_axes_dim) # Pass through each layer sequentially diff --git a/src/diffulab/networks/utils/nn.py b/src/diffulab/networks/utils/nn.py index 5b1c6cf..607f201 100644 --- a/src/diffulab/networks/utils/nn.py +++ b/src/diffulab/networks/utils/nn.py @@ -260,26 +260,26 @@ def forward( def get_cos_sin_ndim_grid( - pos_id: Int[Tensor, "seq_len n_axes"], base: float, axes_dim: list[int] + pos_id: Int[Tensor, "B seq_len n_axes"], base: float, axes_dim: list[int] ) -> tuple[Tensor, Tensor]: """ Get cos/sin for N-D grid positions. Args: - pos_id: [S, n_axes] positional IDs along each axis. + pos_id: [B, S, n_axes] positional IDs along each axis. base: base frequency for RoPE. axes_dim: list of rotary dimensions per axis. Returns: - cos: [S, dim/2] - sin: [S, dim/2] + cos: [B, S, dim/2] + sin: [B, S, dim/2] """ - assert len(axes_dim) == pos_id.shape[1], "axes_dim length must match pos_id n_axes" + assert len(axes_dim) == pos_id.shape[-1], "axes_dim length must match pos_id n_axes" cos_chunks: list[Tensor] = [] sin_chunks: list[Tensor] = [] for axis_idx, axis_dim in enumerate(axes_dim): # pos along this axis: [S] - pos_i = pos_id[:, axis_idx].to(dtype=torch.float64) + pos_i = pos_id[..., axis_idx].to(dtype=torch.float64) freqs = 1.0 / ( base @@ -295,11 +295,11 @@ def get_cos_sin_ndim_grid( ) ) # [D_i/2] - # angles: [S, D_i/2] - angles_i = torch.outer(pos_i, freqs) + # angles: [B, S, D_i/2] + angles_i = torch.einsum("...s,d->...sd", pos_i, freqs) - cos_i = angles_i.cos().float() # [S, D_i/2] - sin_i = angles_i.sin().float() # [S, D_i/2] + cos_i = angles_i.cos().float() # [B, S, D_i/2] + sin_i = angles_i.sin().float() # [B, S, D_i/2] cos_chunks.append(cos_i) sin_chunks.append(sin_i) @@ -330,17 +330,17 @@ def __init__(self, axes_dim: list[int]) -> None: @staticmethod def _apply_rotary( x: Float[Tensor, "batch_size num_head seq_len dim_rot"], # [B, H, S, D_rot] - cos: Float[Tensor, "seq_len dim_rot_half"], # [S, D_rot/2] - sin: Float[Tensor, "seq_len dim_rot_half"], # [S, D_rot/2] + cos: Float[Tensor, "batch_size seq_len dim_rot_half"], # [B, S, D_rot/2] + sin: Float[Tensor, "batch_size seq_len dim_rot_half"], # [B, S, D_rot/2] ) -> Float[Tensor, "batch_size num_head seq_len dim_rot"]: """ Apply rotary positional embedding to the input tensor, assuming: - x[..., 0::2] and x[..., 1::2] form the complex pairs - cos/sin are per-pair, shape [S, D_rot/2] """ - # [1, 1, S, D_rot/2] - cos = cos[None, None, :, :] - sin = sin[None, None, :, :] + # [B, 1, S, D_rot/2] + cos = cos[:, None, :, :] + sin = sin[:, None, :, :] x_even = x[..., 0::2] # [B, H, S, D_rot/2] x_odd = x[..., 1::2] # [B, H, S, D_rot/2] @@ -361,7 +361,7 @@ def forward( q: Float[Tensor, "batch_size seq_len n_heads head_dim"], k: Float[Tensor, "batch_size seq_len n_heads head_dim"], v: Float[Tensor, "batch_size seq_len n_heads head_dim"], - cos_sin: tuple[Float[Tensor, "seq_len dim/2"], Float[Tensor, "seq_len dim/2"]], + cos_sin: tuple[Float[Tensor, "batch_size seq_len dim/2"], Float[Tensor, "batch_size seq_len dim/2"]], # pos_id: Int[Tensor, "seq_len n_axes"], ) -> tuple[Tensor, Tensor, Tensor]: """ @@ -375,8 +375,8 @@ def forward( (q_rot, k_rot, v): q,k rotated on the first self.dim channels; v unchanged. """ cos, sin = cos_sin # precomputed - cos = cos.to(device=q.device, dtype=q.dtype) # [S, dim/2] - sin = sin.to(device=q.device, dtype=q.dtype) # [S, dim/2] + cos = cos.to(device=q.device, dtype=q.dtype) # [B, S, dim/2] + sin = sin.to(device=q.device, dtype=q.dtype) # [B, S, dim/2] # [B, S, H, D] -> [B, H, S, D] q = q.transpose(1, 2) From bfe26bb107fd41032a6364d4853f10446111ab3c Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Tue, 6 Jan 2026 17:11:44 +0100 Subject: [PATCH 02/13] add sprint model --- src/diffulab/networks/denoisers/sprint.py | 619 ++++++++++++++++++++++ 1 file changed, 619 insertions(+) create mode 100644 src/diffulab/networks/denoisers/sprint.py diff --git a/src/diffulab/networks/denoisers/sprint.py b/src/diffulab/networks/denoisers/sprint.py new file mode 100644 index 0000000..280e8a1 --- /dev/null +++ b/src/diffulab/networks/denoisers/sprint.py @@ -0,0 +1,619 @@ +from typing import Any + +import torch +import torch.nn as nn +from einops import rearrange +from jaxtyping import Float, Int +from torch import Tensor + +from diffulab.networks.denoisers.common import Denoiser, ModelOutput +from diffulab.networks.denoisers.mmdit import DiTBlock, MMDiTBlock, ModulatedLastLayer +from diffulab.networks.embedders.common import ContextEmbedder, ContextEmbedderOutput +from diffulab.networks.utils.nn import ( + LabelEmbed, + Modulation, + get_cos_sin_ndim_grid, + timestep_embedding, +) +from diffulab.networks.utils.utils import zero_module + + +class SprintDiT(Denoiser): + """ + (mm)DiT with sprint integration (https://arxiv.org/pdf/2510.21986). + + Args: + simple_dit (bool): If True, use DiT blocks with class-label conditioning only (no context + or cross-attention). If False, use MMDiT blocks with cross-attention to contextual tokens + produced by `context_embedder`. Default: False. + input_channels (int): Number of channels of the main input x. Default: 3. + output_channels (int | None): Number of channels to predict. If None, equals `input_channels`. + Default: None. + input_dim (int): Token/patch embedding width for the image stream. Also the hidden size used + before the final per-patch projection. Default: 4096. + hidden_dim (int): Inner attention dimension for attention projections. Default: 4096. + embedding_dim (int): Conditioning embedding width (for timestep/labels/pooled context) used + by modulation layers and the last prediction layer. Default: 4096. + num_heads (int): Number of attention heads in each block. Default: 16. + mlp_ratio (int): Expansion ratio for the MLP in each block. Default: 4. + patch_size (int): Side length P of square patches. Images are projected with stride P. Default: 16. + context_dim (int): Model width for contextual tokens after `context_embed` when + `simple_dit=False`. Ignored when `simple_dit=True`. Default: 4096. + encoder_depth (int): Number of transformer blocks in the encoder. Default: 2. + deep_layers_depth (int): Number of transformer blocks in the deep layers path. Default: 8. + decoder_depth (int): Number of transformer blocks in the decoder. Default: 2. + rope_base (int): Base frequency for RoPE. Default: 10000. + partial_rotary_factor (float): Fraction of each head dimension using RoPE. + 1.0 means full rotary. Default: 1.0. + rope_axes_dim (list[int] | None): List of dimensions for rotary positional embeddings. + When `simple_dit=True`, should contain 2 integers for H and W axes. When `simple_dit=False`, + should contain 3 integers for L, H, W axes. If None, defaults are used based on + partial_rotary_factor and the heads_dim. Default: None + frequency_embedding (int): Size of the Fourier timestep embedding before the time MLP. + Default: 256. + n_classes (int | None): Number of classes for label conditioning in `simple_dit` mode. + Required to use classifier-free guidance with labels. Must be None when using + a `context_embedder`. Default: None. + classifier_free (bool): Enables classifier-free guidance. In `simple_dit`, it applies to + dropped labels; in MMDiT mode, it is forwarded to the context embedder which may drop + context. Additionally for Sprint, drops the deep layers path (path free guidance). Default: False. + context_embedder (ContextEmbedder | None): When `simple_dit=False`, a module returning + `ContextEmbedderOutput`. Must be provided for text/image conditioning and must be None + when `simple_dit=True`. If the embedder returns pooled and token embeddings + (`n_output == 2`), pooled features are fused into the timestep embedding via an MLP. + Default: None. + use_checkpoint (bool): Enable torch.utils.checkpoint in blocks to trade compute for memory. + Default: False. + drop_rate (float): Fraction of image tokens to drop in the deep layers path during training. + Default: 0.75. + """ + + def __init__( + self, + simple_dit: bool = False, + input_channels: int = 3, + output_channels: int | None = None, + input_dim: int = 768, + hidden_dim: int = 768, + num_heads: int = 12, + mlp_ratio: int = 4, + patch_size: int = 16, + context_dim: int = 768, + encoder_depth: int = 2, + deep_layers_depth: int = 8, + decoder_depth: int = 2, + rope_base: int = 10_000, + partial_rotary_factor: float = 1, + rope_axes_dim: list[int] | None = None, + frequency_embedding: int = 256, + n_classes: int | None = None, + classifier_free: bool = False, + context_embedder: ContextEmbedder | None = None, + use_checkpoint: bool = False, + drop_rate: float = 0.75, + ): + super().__init__() + assert not (n_classes is not None and context_embedder is not None), ( + "n_classes and context_embedder cannot both be specified" + ) + self.simple_dit = simple_dit + self.patch_size = patch_size + self.input_channels = input_channels + if not output_channels: + output_channels = input_channels + self.output_channels = output_channels + self.context_embedder = context_embedder + self.frequency_embedding = frequency_embedding + self.rope_base = rope_base + + self.n_classes = n_classes + self.classifier_free = classifier_free + + self.mask_token = nn.Parameter(torch.zeros(1, 1, input_dim)) + self.drop_rate = drop_rate + + heads_dim = hidden_dim // num_heads + if not self.simple_dit: + assert self.context_embedder is not None, "for dit with text context embedder must be provided" + assert isinstance(self.context_embedder.output_size, tuple) and all( + isinstance(i, int) for i in self.context_embedder.output_size + ), "context_embedder.output_size must be a tuple of integers" + + self.pooled_embedding = False + self.mlp_pooled_context = None + if self.context_embedder.n_output == 2: + self.pooled_embedding = True + self.mlp_pooled_context = nn.Sequential( + nn.Linear(self.context_embedder.output_size[0], input_dim), + nn.SiLU(), + nn.Linear(input_dim, input_dim), + ) + self.context_embed = nn.Linear(self.context_embedder.output_size[1], context_dim) + else: + assert self.context_embedder.n_output == 1 + self.context_embed = nn.Linear(self.context_embedder.output_size[0], context_dim) + if rope_axes_dim is None: + rope_axes_dim = [ + int((partial_rotary_factor * heads_dim) // 3), # L for text, set to 0 for image tokens + int((partial_rotary_factor * heads_dim) // 3), # H set to 0 for text + int((partial_rotary_factor * heads_dim) // 3), # W set to 0 for text + ] + else: + self.label_embed = ( + LabelEmbed(self.n_classes, input_dim, self.classifier_free) if self.n_classes is not None else None + ) + if rope_axes_dim is None: + rope_axes_dim = [ + int((partial_rotary_factor * heads_dim) // 2), # H + int((partial_rotary_factor * heads_dim) // 2), # W + ] + self.rope_axes_dim = rope_axes_dim + + self.time_embed = nn.Sequential( + nn.Linear(self.frequency_embedding, input_dim), + nn.SiLU(), + nn.Linear(input_dim, input_dim), + ) + + self.conv_proj = nn.Conv2d(self.input_channels, input_dim, kernel_size=self.patch_size, stride=self.patch_size) + + self.fuse = nn.Linear(input_dim * 2, input_dim) + if not self.simple_dit: + self.fuse_context = nn.Linear(2 * context_dim, context_dim) + + self.last_layer = ModulatedLastLayer( + embedding_dim=input_dim, + hidden_size=input_dim, + patch_size=self.patch_size, + out_channels=self.output_channels, + ) + + # -------------- + # encoder layers + # -------------- + self.layers = nn.ModuleList( # name compatibility for RePA + [ + MMDiTBlock( + context_dim=context_dim, + input_dim=input_dim, + hidden_dim=hidden_dim, + embedding_dim=input_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + use_checkpoint=use_checkpoint, + ) + if not self.simple_dit + else DiTBlock( + input_dim=input_dim, + hidden_dim=hidden_dim, + embedding_dim=input_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + use_checkpoint=use_checkpoint, + ) + for _ in range(encoder_depth) + ] + ) + + # -------------- + # deep layers + # -------------- + self.deep_layers = nn.ModuleList( + [ + MMDiTBlock( + context_dim=context_dim, + input_dim=input_dim, + hidden_dim=hidden_dim, + embedding_dim=input_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + use_checkpoint=use_checkpoint, + ) + if not self.simple_dit + else DiTBlock( + input_dim=input_dim, + hidden_dim=hidden_dim, + embedding_dim=input_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + use_checkpoint=use_checkpoint, + ) + for _ in range(deep_layers_depth) + ] + ) + + # -------------- + # decoder layers + # -------------- + self.decoder_layers = nn.ModuleList( + [ + MMDiTBlock( + context_dim=context_dim, + input_dim=input_dim, + hidden_dim=hidden_dim, + embedding_dim=input_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + use_checkpoint=use_checkpoint, + ) + if not self.simple_dit + else DiTBlock( + input_dim=input_dim, + hidden_dim=hidden_dim, + embedding_dim=input_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + use_checkpoint=use_checkpoint, + ) + for _ in range(decoder_depth) + ] + ) + + self.apply(self._init_weights) + + def _init_weights(self, module: nn.Module) -> None: + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: # type: ignore + nn.init.constant_(module.bias, 0) + if isinstance(module, Modulation): + zero_module(module) + if isinstance(module, ModulatedLastLayer): + zero_module(module.adaLN_modulation) + + def patchify( + self, x: Float[Tensor, "batch_size channels height width"] + ) -> Float[Tensor, "batch_size num_patches patch_dim"]: + """ + Convert image tensor into patches using convolutional projection. + Args: + x (Tensor): Input tensor of shape (B, C, H, W) + Returns: + Tensor: Patchified tensor of shape (B, num_patches, patch_dim) + """ + _, _, H, W = x.shape + self.original_size = (H, W) + + x = self.conv_proj(x) + _, _, Hp, Wp = x.shape + self.grid_size = (Hp, Wp) + x = rearrange(x, "b c h w -> b (h w) c") + + return x + + def unpatchify( + self, x: Float[Tensor, "batch_size seq_len patch_dim"] + ) -> Float[Tensor, "batch_size channels height width"]: + """ + Convert patches back into the original image tensor. + Args: + x (Tensor): Patchified tensor of shape (B, num_patches, patch_dim) + Returns: + Tensor: Reconstructed image tensor of shape (B, C, H, W) + """ + # Reshape the tensor to the original image dimensions + x = rearrange( + x, + "b (h w) (p1 p2 c) -> b c (h p1) (w p2)", + h=self.grid_size[0], + w=self.grid_size[1], + p1=self.patch_size, + p2=self.patch_size, + c=self.output_channels, + ) + return x + + def drop_tokens( + self, + x: Float[Tensor, "batch_size seq_len patch_dim"], + cos_sin_rope: tuple[Float[Tensor, "batch_size seq_len rope_dim"], ...], + ) -> tuple[ + Float[Tensor, "batch_size kept_seq_len patch_dim"], + Int[Tensor, "batch_size kept_seq_len"], + tuple[Float[Tensor, "batch_size kept_seq_len rope_dim"], ...], + ]: + """ + Drop a fraction of tokens during training for Sprint. + Args: + x (Tensor): Input tensor of shape (B, seq_len, patch_dim) + cos_sin_rope (tuple): Tuple of RoPE tensors, each of shape (B, seq_len, rope_dim) + Returns: + x_dropped (Tensor): Tensor with dropped tokens of shape (B, kept_seq_len, patch_dim) + kept_indices (Tensor): Indices of kept tokens of shape (B, kept_seq_len) + cos_sin_rope_dropped (tuple): Tuple of RoPE tensors for kept tokens, + each of shape (B, kept_seq_len, rope_dim) + """ + B, S, D = x.shape + + if not self.training: + return x, torch.arange(S, device=x.device).expand(B, S), cos_sin_rope + + k = max(1, int(S * (1.0 - float(self.drop_rate)))) + scores = torch.rand((B, S), device=x.device, dtype=torch.float32) + kept_indices = torch.topk(scores, k=k, dim=1, largest=True, sorted=False).indices.to(torch.long) + order = torch.argsort(kept_indices, dim=1) + kept_indices = torch.gather(kept_indices, dim=1, index=order) + + x_dropped = torch.gather(x, dim=1, index=kept_indices.unsqueeze(-1).expand(B, k, D)) + cos_sin_rope_dropped = tuple( + torch.gather(rope, dim=1, index=kept_indices.unsqueeze(-1).expand(B, k, rope.shape[-1])) + for rope in cos_sin_rope + ) + + return x_dropped, kept_indices, cos_sin_rope_dropped + + def restore_tokens( + self, + x_dropped: Float[Tensor, "batch_size kept_seq_len patch_dim"], + kept_indices: Int[Tensor, "batch_size kept_seq_len"], + path_drop_p: float = 0.0, + ) -> Float[Tensor, "batch_size seq_len patch_dim"]: + """ + Restore dropped tokens to their original positions using a mask token. + Args: + x_dropped (Tensor): Tensor with dropped tokens of shape (B, kept_seq_len, patch_dim) + kept_indices (Tensor): Indices of kept tokens of shape (B, kept_seq_len) + path_drop_p (float): Probability of dropping the dense path during training. Default: 0.0 + Returns: + x_full (Tensor): Tensor with restored tokens of shape (B, seq_len, patch_dim) + """ + H, W = self.grid_size + B, _, D = x_dropped.shape + S = H * W + + # Create full tensor filled with mask token + mask_token = self.mask_token.to(dtype=x_dropped.dtype) + x_full = mask_token.expand(B, S, D).clone() + + # Scatter kept tokens back to their original positions + kept_indices_expanded = kept_indices.unsqueeze(-1).expand(-1, -1, D) # (B, num_kept, D) + x_full.scatter_(dim=1, index=kept_indices_expanded, src=x_dropped) + + if path_drop_p > 0: + drop_mask = torch.rand(B, device=x_full.device) < path_drop_p + x_full = torch.where(drop_mask[:, None, None], mask_token.expand_as(x_full), x_full) + + return x_full + + def _forward_mmdit( + self, + x: Float[Tensor, "batch_size seq_len patch_dim"], + timesteps: Float[Tensor, "batch_size"], + initial_context: Any | None = None, + p: float = 0.0, + intermediate_features: bool = False, + ) -> ModelOutput: + """ + Forward pass through the Sprint mmDiT model. + Args: + x (Tensor): Input tensor of shape (B, seq_len, patch_dim) + timesteps (Tensor): Timestep tensor of shape (B,) + initial_context (Any, optional): Initial context for the context embedder. Defaults to None. + p (float, optional): Probability for classifier/path-free guidance. Defaults to 0.0. + intermediate_features (bool, optional): Whether to return intermediate features. Defaults to False. + Returns: + ModelOutput: Output dictionary containing the final output and optional intermediate features. + + """ + assert self.context_embedder is not None, "for MMDiT context embedder must be provided" + emb = self.time_embed(timestep_embedding(timesteps, self.frequency_embedding)) + context_output: ContextEmbedderOutput = self.context_embedder(initial_context, p) + if self.pooled_embedding: + assert self.mlp_pooled_context is not None, ( + "for MMDiT with pooled context, mlp_pooled_context must be defined" + ) + assert "pooled_embeddings" in context_output, "pooled embeddings must be in context_output" + context_pooled = context_output["pooled_embeddings"] + emb = self.mlp_pooled_context(context_pooled) + emb + + context = context_output["embeddings"] + context = self.context_embed(context) + attn_mask = context_output.get("attn_mask", None) + + # pos_ids: [S, n_axes] positional IDs along each axis for rope + # in mmdit attention we concat with context first. Context have 0,0 for h w + # text: (t>0, 0, 0) + text_pos_ids = torch.stack( + [ + torch.arange(1, context.shape[1] + 1, device=x.device), + torch.zeros(context.shape[1], device=x.device, dtype=torch.long), + torch.zeros(context.shape[1], device=x.device, dtype=torch.long), + ], + dim=-1, + ) + + # image: (0, h, w) + img_pos_ids = torch.stack( + torch.meshgrid( + torch.zeros(1, device=x.device, dtype=torch.long), + torch.arange(self.grid_size[0], device=x.device), + torch.arange(self.grid_size[1], device=x.device), + indexing="ij", + ), + dim=-1, + ).view(-1, 3) + + pos_ids = torch.cat([text_pos_ids, img_pos_ids], dim=0).unsqueeze(0).repeat(x.size(0), 1, 1) + cos_sin_rope = get_cos_sin_ndim_grid(pos_ids, base=self.rope_base, axes_dim=self.rope_axes_dim) + + features: list[Tensor] | None = [] if intermediate_features else None + + # Pass through each encoder layer sequentially + for layer in self.layers: + x, context = layer(x, emb, context, cos_sin_rope=cos_sin_rope, attn_mask=attn_mask) + if features is not None: + features.append(x) + encoder_context = context.clone() + + # drop tokens and forward through deep layers + cos_sin_rope_img = tuple(rope[:, text_pos_ids.shape[0] :] for rope in cos_sin_rope) + x_dropped, kept_indices, cos_sin_rope_img_dropped = self.drop_tokens(x, cos_sin_rope=cos_sin_rope_img) + cos_sin_rope_dropped = tuple( + torch.cat([rope[:, : text_pos_ids.shape[0]], cos_sin_rope_img_dropped[i]], dim=1) + for i, rope in enumerate(cos_sin_rope) + ) + if p < 1: + for layer in self.deep_layers: + x_dropped, context = layer( + x_dropped, emb, context, cos_sin_rope=cos_sin_rope_dropped, attn_mask=attn_mask + ) # type: ignore + if features is not None: + features.append(x_dropped) + x_restored = self.restore_tokens(x_dropped, kept_indices, p) + else: + x_restored = self.mask_token.expand_as(x).clone() + + # fuse deep layers output with residual connection from encoder + x_fused = self.fuse(torch.cat([x_restored, x], dim=-1)) + context_fused = self.fuse_context(torch.cat([context, encoder_context], dim=-1)) # type: ignore + + for layer in self.decoder_layers: + x_fused, context_fused = layer( + x_fused, + emb, + context_fused, + cos_sin_rope=cos_sin_rope, + attn_mask=attn_mask, # type: ignore + ) + if features is not None: + features.append(x_fused) + + x_fused = self.last_layer(x_fused, emb) + if features is not None: + features.append(x_fused) + + x_output = self.unpatchify(x_fused) + + output: ModelOutput = {"x": x_output} + if features is not None: + output["features"] = features + + return output + + def _forward_dit( + self, + x: Float[Tensor, "batch_size seq_len patch_dim"], + timestep: Float[Tensor, "batch_size"], + p: float = 0.0, + y: Int[Tensor, "batch_size"] | None = None, + intermediate_features: bool = False, + ) -> ModelOutput: + """ + Forward pass through the Sprint DiT model. + Args: + x (Tensor): Input tensor of shape (B, seq_len, patch_dim) + timestep (Tensor): Timestep tensor of shape (B,) + p (float, optional): Probability for classifier/path-free guidance. Defaults to 0.0 + y (Tensor, optional): Class labels. Defaults to None. + intermediate_features (bool, optional): Whether to return intermediate features. Defaults to False. + Returns: + ModelOutput: Output dictionary containing the final output and optional intermediate features. + """ + if p > 0: + assert self.n_classes, ( + "probability of dropping for classifier free guidance is only available if a number of classes is set" + ) + + emb = self.time_embed(timestep_embedding(timestep, self.frequency_embedding)) + if self.label_embed is not None: + emb = emb + self.label_embed(y, p) + + # pos_ids: [S, n_axes] positional IDs along each axis for rope + pos_ids = ( + torch.stack( + torch.meshgrid( + [ + torch.arange(self.grid_size[0], device=x.device), + torch.arange(self.grid_size[1], device=x.device), + ], + indexing="ij", + ), + dim=-1, + ) + .view(-1, 2) + .unsqueeze(0) + .repeat(x.size(0), 1, 1) + ) + cos_sin_rope = get_cos_sin_ndim_grid(pos_ids, base=self.rope_base, axes_dim=self.rope_axes_dim) + + features: list[Tensor] | None = [] if intermediate_features else None + # Pass through each layer sequentially + for layer in self.layers: + x = layer(x, emb, cos_sin_rope=cos_sin_rope) + if features is not None: + features.append(x) + + x_dropped, kept_indices, cos_sin_rope_dropped = self.drop_tokens(x, cos_sin_rope=cos_sin_rope) + + if p < 1: + for layer in self.deep_layers: + x_dropped = layer(x_dropped, emb, cos_sin_rope=cos_sin_rope_dropped) + if features is not None: + features.append(x_dropped) + x_restored = self.restore_tokens(x_dropped, kept_indices, p) + else: + x_restored = self.mask_token.expand_as(x).clone() + + x_fused = self.fuse(torch.cat([x_restored, x], dim=-1)) + + for layer in self.decoder_layers: + x_fused = layer(x_fused, emb, cos_sin_rope=cos_sin_rope) + if features is not None: + features.append(x_fused) + + x_fused = self.last_layer(x_fused, emb) + if features is not None: + features.append(x_fused) + + x_output = self.unpatchify(x_fused) + + output: ModelOutput = {"x": x_output} + if features is not None: + output["features"] = features + + return output + + def forward( + self, + x: Float[Tensor, "batch_size channels height width"], + timesteps: Float[Tensor, "batch_size"], + initial_context: Any | None = None, + p: float = 0.0, + y: Int[Tensor, "batch_size"] | None = None, + x_context: Tensor | None = None, + intermediate_features: bool = False, + ) -> ModelOutput: + """ + Forward pass through the Sprint model. + Args: + x (Tensor): Input tensor of shape (B, C, H, W) + timesteps (Tensor): Timestep tensor of shape (B,) + initial_context (Any, optional): Initial context for the context embedder. Defaults to None. + p (float, optional): Probability for classifier/path-free guidance. Defaults to 0.0 + y (Tensor, optional): Class labels. Defaults to None. + x_context (Tensor, optional): Additional context tensor to concatenate with input. Defaults to None. + intermediate_features (bool, optional): Whether to return intermediate features. Defaults to False. + path_drop_p (float, optional): Probability of dropping the dense path. Defaults to 0.0. + Returns: + ModelOutput: Output dictionary containing the final output and optional intermediate features. + """ + assert not (initial_context is not None and y is not None), "initial_context and y cannot both be specified" + if p > 0: + assert self.classifier_free, ( + "probability of dropping for classifier free guidance is only available if model is set up to be classifier free" + ) + if x_context is not None: + x = torch.cat([x, x_context], dim=1) + + encoder_input = self.patchify(x) + + if self.simple_dit: + return self._forward_dit(encoder_input, timesteps, p, y, intermediate_features) + + return self._forward_mmdit(encoder_input, timesteps, initial_context, p, intermediate_features) From f05242ff715a642d9bfda27d5c859a2de917fa65 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Tue, 6 Jan 2026 17:12:32 +0100 Subject: [PATCH 03/13] add PrecomputedEmbedder class --- .../networks/embedders/precomputed.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 src/diffulab/networks/embedders/precomputed.py diff --git a/src/diffulab/networks/embedders/precomputed.py b/src/diffulab/networks/embedders/precomputed.py new file mode 100644 index 0000000..26fc23c --- /dev/null +++ b/src/diffulab/networks/embedders/precomputed.py @@ -0,0 +1,43 @@ +from pathlib import Path + +import torch + +from diffulab.networks.embedders.common import ContextEmbedder, ContextEmbedderOutput + + +class PrecomputedEmbedder(ContextEmbedder): + def __init__(self, path_null_embedding: Path | str, null_embedding_seq_len: int) -> None: + super().__init__() + self.null_embedding = torch.load(path_null_embedding).squeeze() + self.null_embedding_mask = torch.cat( + [ + torch.ones(null_embedding_seq_len, dtype=torch.bool), + torch.zeros(self.null_embedding.shape[0] - null_embedding_seq_len, dtype=torch.bool), + ], + dim=0, + ) + self._output_size = (self.null_embedding.shape[-1],) # type: ignore + self._n_output = 1 + + def drop_conditions(self, context: ContextEmbedderOutput, p: float) -> ContextEmbedderOutput: + batch_size = context["embeddings"].shape[0] + device, dtype = context["embeddings"].device, context["embeddings"].dtype + + drop_mask = torch.rand(batch_size, device=device) < p + null_emb = self.null_embedding.to(device=device, dtype=dtype) + null_mask = self.null_embedding_mask.to(device=device) + + embeddings = torch.where( + drop_mask[:, None, None], null_emb.unsqueeze(0).expand(batch_size, -1, -1), context["embeddings"] + ) + attn_mask = torch.where( + drop_mask[:, None], + null_mask.unsqueeze(0).expand(batch_size, -1), + context["attn_mask"], # type: ignore + ) + + return {"embeddings": embeddings, "attn_mask": attn_mask} + + def forward(self, context: ContextEmbedderOutput, p: float = 0) -> ContextEmbedderOutput: + dropped_context = self.drop_conditions(context, p) + return dropped_context From 9df74c7ea4ab6f4c359e6da9a921a4ff500010c8 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Tue, 6 Jan 2026 17:15:11 +0100 Subject: [PATCH 04/13] add SmolVLMTextEmbedder class --- src/diffulab/networks/embedders/qwen.py | 5 +- src/diffulab/networks/embedders/smolVLM.py | 91 ++++++++++++++++++++++ 2 files changed, 93 insertions(+), 3 deletions(-) create mode 100644 src/diffulab/networks/embedders/smolVLM.py diff --git a/src/diffulab/networks/embedders/qwen.py b/src/diffulab/networks/embedders/qwen.py index 00eead3..b149e06 100644 --- a/src/diffulab/networks/embedders/qwen.py +++ b/src/diffulab/networks/embedders/qwen.py @@ -45,8 +45,7 @@ def drop_conditions(self, context: list[str], p: float) -> list[str]: """ return ["" if random.random() < p else c for c in context] - @torch._dynamo.disable # type: ignore[reportUnknownMemberType] - def forward(self, context: list[str], p: float = 0) -> ContextEmbedderOutput: + def forward(self, context: list[str], p: float = 0, force_padding: bool = False) -> ContextEmbedderOutput: """Compute Qwen text embeddings. Args: @@ -65,7 +64,7 @@ def forward(self, context: list[str], p: float = 0) -> ContextEmbedderOutput: tokens = self.tokenizer( prompt_texts, max_length=self.max_length + self.prompt_template_encoder_start_idx, - padding=True, + padding=True if not force_padding else "max_length", truncation=True, return_tensors="pt", ).to(self.model.device) diff --git a/src/diffulab/networks/embedders/smolVLM.py b/src/diffulab/networks/embedders/smolVLM.py new file mode 100644 index 0000000..0b53f5c --- /dev/null +++ b/src/diffulab/networks/embedders/smolVLM.py @@ -0,0 +1,91 @@ +import random +from typing import cast + +import torch +from torch import Tensor +from transformers import GPT2TokenizerFast, Idefics3ForConditionalGeneration + +from diffulab.networks.embedders.common import ContextEmbedder, ContextEmbedderOutput + + +class SmolVLMTextEmbedder(ContextEmbedder): + def __init__( + self, + model_id: str = "HuggingFaceTB/SmolVLM-256M-Instruct", + device: str | torch.device = "cuda", + max_length: int = 1024, + ) -> None: + super().__init__() + self.model: Idefics3ForConditionalGeneration = Idefics3ForConditionalGeneration.from_pretrained( # type: ignore + model_id, device_map=device, dtype="auto" + ) + self.model.eval() + self.model.requires_grad_(False) + self.tokenizer: GPT2TokenizerFast = GPT2TokenizerFast.from_pretrained(model_id, device_map=device) # type: ignore + self.prompt_template: str = self._get_prompt_template() + self.prompt_template_encoder_start_idx = 33 + self.max_length = max_length + + self._output_size: tuple[int] = (self.model.config.text_config.hidden_size,) # type: ignore + self._n_output = 1 + + def _get_prompt_template(self) -> str: + return ( + "<|im_start|>System: Describe the image by detailing the color, shape, size, texture, quantity, text, " + "spatial relationships of the objects and background.\nUser: {}\n" + "Assistant: " + ) + + def drop_conditions(self, context: list[str], p: float) -> list[str]: + """Randomly drop prompts for classifier-free guidance style training. + + Args: + context (list[str]): Original list of prompt strings. + p (float): Drop probability. + + Returns: + list[str]: context with some entries replaced by empty strings. + """ + return ["" if random.random() < p else c for c in context] + + @torch._dynamo.disable # type: ignore[reportUnknownMemberType] + def forward(self, context: list[str], p: float = 0, force_padding: bool = False) -> ContextEmbedderOutput: + """Compute smolVLM text embeddings. + + Args: + context (list[str]): List of input prompt strings. + p (float): Probability of dropping the condition (replacing with empty string). + + Returns: + ContextEmbedderOutput: + Sequence embeddings of shape ``[B, N_ctx, D]`` where ``N_ctx`` is the number of tokens after prompt template + and ``D`` is the model hidden size. + Attention mask of shape ``[B, N_ctx]`` indicating non-padding tokens. + + """ + context_dropped = self.drop_conditions(context, p) + prompt_texts = [self.prompt_template.format(c) for c in context_dropped] + tokens = self.tokenizer( + prompt_texts, + max_length=self.max_length + self.prompt_template_encoder_start_idx, + padding=True if not force_padding else "max_length", + truncation=True, + return_tensors="pt", + ).to(self.model.device) + + embeddings = cast( + Tensor, + self.model( + input_ids=tokens.input_ids, # type: ignore + attention_mask=tokens.attention_mask, # type: ignore + output_hidden_states=True, + ).hidden_states[-1], + ) + + embeddings = embeddings[:, self.prompt_template_encoder_start_idx :, :] + attn_mask = tokens.attention_mask[:, self.prompt_template_encoder_start_idx :] # type: ignore + + return ContextEmbedderOutput( + embeddings=embeddings, + attn_mask=attn_mask, # type: ignore + ) From 3d0c4ef122eae01428a035a4d02cba3af1a99285 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Tue, 6 Jan 2026 17:16:01 +0100 Subject: [PATCH 05/13] fix repa --- src/diffulab/training/losses/repa.py | 87 ++++++++++------------------ 1 file changed, 31 insertions(+), 56 deletions(-) diff --git a/src/diffulab/training/losses/repa.py b/src/diffulab/training/losses/repa.py index 940c75b..9609e3d 100644 --- a/src/diffulab/training/losses/repa.py +++ b/src/diffulab/training/losses/repa.py @@ -1,5 +1,4 @@ -from typing import Any, Callable, TypedDict -from weakref import WeakKeyDictionary +from typing import Any, TypedDict import torch from jaxtyping import Float @@ -12,13 +11,6 @@ from diffulab.networks.repa.perceiver_resampler import PerceiverResampler from diffulab.training.losses.common import LossFunction -try: - from torch._dynamo import disable as _dynamo_disable # type: ignore -except Exception: - - def _dynamo_disable(fn: Any) -> Any: - return fn - class ResamplerParams(TypedDict): dim: int @@ -82,9 +74,9 @@ def __init__( self, repa_encoder: str = "dinov2", encoder_args: dict[str, Any] = {}, - alignment_layer: int = 8, + alignment_layer: int = 8, # 0 for last layer denoiser_dimension: int = 256, - hidden_dim: int = 256, + hidden_dim: int = 1024, load_dino: bool = True, # whether to load the DINO model weights, if precomputed features are used no need to load it embedding_dim: int = 768, # dimension of the DINO features use_resampler: bool = False, # whether to use the perceiver resampler @@ -118,74 +110,59 @@ def __init__( **resampler_params, ) self.alignment_layer = alignment_layer - self._handles: "WeakKeyDictionary[nn.Module, RemovableHandle]" = WeakKeyDictionary() - self._features: "WeakKeyDictionary[nn.Module, torch.Tensor]" = WeakKeyDictionary() - self._active_model: nn.Module | None = None + self._handles: dict[int, RemovableHandle] = {} + self._captured_features: dict[int, torch.Tensor] = {} + self._active_model_id: int | None = None self._hook_layer_idx = self.alignment_layer - 1 self.coeff = coeff - @_dynamo_disable - def _make_forward_hook(self, key_model: MMDiT) -> Callable[[nn.Module, tuple[Any, ...], torch.Tensor], None]: - def _hook(_mod: nn.Module, _inp: tuple[Any, ...], out: torch.Tensor): - self._features[key_model] = out + def _make_hook(self, model_id: int): + """Create a hook that stores features keyed by model id.""" + + def _hook(_mod: nn.Module, _inp: tuple[Any, ...], out: torch.Tensor) -> None: + self._captured_features[model_id] = out return _hook - def _attach_once(self, model: MMDiT) -> None: - if model in self._handles: + def _attach_hook(self, model: MMDiT) -> None: + """Attach the forward hook to the target layer of the model (once per model).""" + model_id = id(model) + if model_id in self._handles: return + layer = model.layers[self._hook_layer_idx] - handle = layer.register_forward_hook(self._make_forward_hook(model)) # type: ignore - self._handles[model] = handle + handle = layer.register_forward_hook(self._make_hook(model_id)) # type: ignore + self._handles[model_id] = handle def set_model(self, model: MMDiT) -> None: # type: ignore """Register the model to capture features from a specific layer. This attaches a forward hook to the specified ``alignment_layer`` of the - provided model (only once). A forward pass on ``model`` must be executed - after calling this method so that features are captured before computing - the loss. + provided model (only once per model). A forward pass on ``model`` must be + executed after calling this method so that features are captured before + computing the loss. Args: model (MMDiT): The model whose intermediate features will be aligned to the encoder features. """ - self._attach_once(model) - self._active_model = model + self._attach_hook(model) + self._active_model_id = id(model) def _unregister_all(self) -> None: - for h in list(self._handles.values()): - h.remove() + for handle in self._handles.values(): + handle.remove() self._handles.clear() - self._features.clear() - self._active_model = None + self._captured_features.clear() + self._active_model_id = None def forward( self, x0: Float[Tensor, "batch 3 H W"] | None = None, dst_features: Float[Tensor, "batch seq_len n_dim"] | None = None, ) -> Tensor: - """Compute the REPA cosine-similarity loss. - - Either provide input images via ``x0`` to compute destination features - with the encoder, or pass precomputed ``dst_features`` directly. - - Args: - x0 (Tensor): Input images of shape ``[B, 3, H, W]`` used to compute encoder - features when an encoder is available. - dst_features (Tensor): Precomputed encoder features of shape ``[B, S, D]``. - If provided, ``x0`` is ignored. - - Returns: - Tensor: A scalar tensor containing the REPA loss. - - Raises: - RuntimeError: If no captured features are available for the active - model. Ensure ``set_model(...)`` was called and a forward pass - on the model was executed first. - AssertionError: If neither ``x0`` nor ``dst_features`` is provided. - """ - if self._active_model is None or self._active_model not in self._features: + """Compute the REPA cosine-similarity loss.""" + if self._active_model_id is None or self._active_model_id not in self._captured_features: raise RuntimeError( "REPA: no captured features for the active model. Did you call set_model(...) and run a forward pass?" ) @@ -193,12 +170,10 @@ def forward( if dst_features is None: assert self.repa_encoder is not None, "REPA encoder must be initialized to compute features." with torch.no_grad(): - dst_features = self.repa_encoder( - x0 - ) # batch size seqlen embedding_dim # SEE HOW TO HANDLE THE PRE COMPUTING OF FEATURES + dst_features = self.repa_encoder(x0) assert dst_features is not None, "Destination features must be provided or computed." - src_features = self._features[self._active_model] + src_features = self._captured_features[self._active_model_id] if isinstance(src_features, tuple): src_features = src_features[0] projected_src_features: Tensor = self.proj(src_features) From 7499e5ebae6e1a1a536178f36e87eb94afe02733 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Tue, 6 Jan 2026 17:16:34 +0100 Subject: [PATCH 06/13] update init scripts --- src/diffulab/__init__.py | 4 ++++ src/diffulab/networks/__init__.py | 7 +++++-- src/diffulab/networks/denoisers/__init__.py | 3 ++- src/diffulab/networks/embedders/__init__.py | 4 +++- 4 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/diffulab/__init__.py b/src/diffulab/__init__.py index 2397db4..ab9365f 100644 --- a/src/diffulab/__init__.py +++ b/src/diffulab/__init__.py @@ -11,6 +11,8 @@ PerceiverResampler, QwenTextEmbedder, SD3TextEmbedder, + SmolVLMTextEmbedder, + SprintDiT, UNetModel, VisionTower, ) @@ -32,9 +34,11 @@ "Flux2VAE", "MMDiT", "DDT", + "SprintDiT", "PerceiverResampler", "QwenTextEmbedder", "SD3TextEmbedder", + "SmolVLMTextEmbedder", "UNetModel", "VisionTower", "LossFunction", diff --git a/src/diffulab/networks/__init__.py b/src/diffulab/networks/__init__.py index 274ffbb..95398bf 100644 --- a/src/diffulab/networks/__init__.py +++ b/src/diffulab/networks/__init__.py @@ -1,5 +1,5 @@ -from .denoisers import DDT, Denoiser, MMDiT, UNetModel -from .embedders import QwenTextEmbedder, SD3TextEmbedder +from .denoisers import DDT, Denoiser, MMDiT, SprintDiT, UNetModel +from .embedders import PrecomputedEmbedder, QwenTextEmbedder, SD3TextEmbedder, SmolVLMTextEmbedder from .repa import REPA, DinoV2, PerceiverResampler from .vision_towers import DCAE, Flux2VAE, VisionTower @@ -8,8 +8,11 @@ "UNetModel", "MMDiT", "DDT", + "SprintDiT", "SD3TextEmbedder", "QwenTextEmbedder", + "SmolVLMTextEmbedder", + "PrecomputedEmbedder", "REPA", "DinoV2", "PerceiverResampler", diff --git a/src/diffulab/networks/denoisers/__init__.py b/src/diffulab/networks/denoisers/__init__.py index bfd7b36..5136159 100644 --- a/src/diffulab/networks/denoisers/__init__.py +++ b/src/diffulab/networks/denoisers/__init__.py @@ -1,6 +1,7 @@ from .common import Denoiser from .ddt import DDT from .mmdit import MMDiT +from .sprint import SprintDiT from .unet import UNetModel -__all__ = ["Denoiser", "UNetModel", "MMDiT", "DDT"] +__all__ = ["Denoiser", "UNetModel", "MMDiT", "DDT", "SprintDiT"] diff --git a/src/diffulab/networks/embedders/__init__.py b/src/diffulab/networks/embedders/__init__.py index 5a397c8..51c8c36 100644 --- a/src/diffulab/networks/embedders/__init__.py +++ b/src/diffulab/networks/embedders/__init__.py @@ -1,4 +1,6 @@ +from .precomputed import PrecomputedEmbedder from .qwen import QwenTextEmbedder from .sd3 import SD3TextEmbedder +from .smolVLM import SmolVLMTextEmbedder -__all__ = ["SD3TextEmbedder", "QwenTextEmbedder"] +__all__ = ["SD3TextEmbedder", "QwenTextEmbedder", "SmolVLMTextEmbedder", "PrecomputedEmbedder"] From 85614b69f8d44c24b4c18053da1a2ee333fe420a Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Tue, 6 Jan 2026 22:21:54 +0100 Subject: [PATCH 07/13] update readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 50b1413..82e3215 100644 --- a/README.md +++ b/README.md @@ -67,4 +67,4 @@ Here is a To-Do list, feel welcome to help to any point along this list. The alr - [ ] Train our models on toy datasets for different tasks (conditional generation, Image to Image ...) - [ ] Add possibility to train LORA/DORA - [x] add different sampler -- [ ] Add SPRINT (https://arxiv.org/pdf/2510.21986) +- [x] Add SPRINT (https://arxiv.org/pdf/2510.21986) From 96f91975d5c2a6c0565813154c11111c805605b7 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Tue, 6 Jan 2026 22:23:25 +0100 Subject: [PATCH 08/13] update configs --- configs/model/sprint.yaml | 17 +++++ .../train_imagenet_flow_matching_repa.yaml | 2 +- configs/train_imagenet_repa_txt_to_img.yaml | 2 +- ...train_imagenet_repa_txt_to_img_sprint.yaml | 62 +++++++++++++++++++ 4 files changed, 81 insertions(+), 2 deletions(-) create mode 100644 configs/model/sprint.yaml create mode 100644 configs/train_imagenet_repa_txt_to_img_sprint.yaml diff --git a/configs/model/sprint.yaml b/configs/model/sprint.yaml new file mode 100644 index 0000000..40c7514 --- /dev/null +++ b/configs/model/sprint.yaml @@ -0,0 +1,17 @@ +_target_: diffulab.networks.SprintDiT +simple_dit: true +input_channels: 3 +output_channels: 3 +input_dim: 512 +hidden_dim: 512 +context_dim: 512 +num_heads: 8 +mlp_ratio: 4 +patch_size: 2 +encoder_depth: 2 +deep_layers_depth: 8 +decoder_depth: 2 +n_classes: 10 +classifier_free: false +use_checkpoint: false +drop_rate: 0.75 diff --git a/configs/train_imagenet_flow_matching_repa.yaml b/configs/train_imagenet_flow_matching_repa.yaml index 23d422c..bd484c2 100644 --- a/configs/train_imagenet_flow_matching_repa.yaml +++ b/configs/train_imagenet_flow_matching_repa.yaml @@ -15,7 +15,7 @@ trainer: project_name: imagenet_repa_flow_matching n_epoch: 150 precision_type: "bf16" - p_classifier_free_guidance: 0.2 + p_classifier_free_guidance: 0.1 ema_update_after_step: 0 ema_update_every: 10 diff --git a/configs/train_imagenet_repa_txt_to_img.yaml b/configs/train_imagenet_repa_txt_to_img.yaml index 62eecbd..95b63b0 100644 --- a/configs/train_imagenet_repa_txt_to_img.yaml +++ b/configs/train_imagenet_repa_txt_to_img.yaml @@ -16,7 +16,7 @@ trainer: project_name: txt_to_img_imagenet_repa_flux2 n_epoch: 50 precision_type: "bf16" - p_classifier_free_guidance: 0.2 + p_classifier_free_guidance: 0.1 gradient_accumulation_step: 16 val_step_shift: 6.93 compile: false diff --git a/configs/train_imagenet_repa_txt_to_img_sprint.yaml b/configs/train_imagenet_repa_txt_to_img_sprint.yaml new file mode 100644 index 0000000..ae12917 --- /dev/null +++ b/configs/train_imagenet_repa_txt_to_img_sprint.yaml @@ -0,0 +1,62 @@ +# configs/train_cifar10_flow_matching.yaml +# @package _global_ +defaults: + - model: sprint + - diffuser: rectified_flow + - trainer: default + - dataset: imagenet_repa + - dataloader: default + - optimizer: adamw + - vision_tower: flux2 + - embedder: precomputed + - _self_ + +# Override specific settings +trainer: + project_name: txt_to_img_imagenet_multiar_repa_flux2_sprint + n_epoch: 50 + precision_type: "bf16" + p_classifier_free_guidance: 0.1 + gradient_accumulation_step: 8 + val_step_shift: 6.93 + compile: true + ema_update_after_step: 0 + ema_update_every: 1 + ema_rate: 0.9999 + +model: + input_channels: 128 + output_channels: 128 + input_dim: 640 + hidden_dim: 640 + context_dim: 640 + num_heads: 16 + mlp_ratio: 3 + patch_size: 1 + encoder_depth: 3 + deep_layers_depth: 12 + decoder_depth: 3 + classifier_free: true + rope_base: 1000 + rope_axes_dim: [12, 14, 14] + n_classes: null + simple_dit: false + +diffuser: + extra_args: + logits_normal: true + shift: 4.63 + +dataloader: + batch_size: 32 + +# Hydra configuration +hydra: + run: + dir: outputs/${now:%Y-%m-%d}/${now:%H-%M-%S} + sweep: + dir: multirun/${now:%Y-%m-%d}/${now:%H-%M-%S} + subdir: ${hydra.job.num} + +perceiver_resampler: + use_resampler: false From 12f5f4a33ea22633e038e269a204608774b8b323 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Fri, 9 Jan 2026 15:47:28 +0100 Subject: [PATCH 09/13] simplify dim handling in DiT based denoisers + add Single Stream block --- configs/embedder/precomputed.yaml | 3 + configs/model/ddt.yaml | 4 +- configs/model/dit.yaml | 3 +- configs/model/sprint.yaml | 5 +- .../train_imagenet_flow_matching_repa.yaml | 3 +- configs/train_imagenet_repa_txt_to_img.yaml | 4 +- ...train_imagenet_repa_txt_to_img_sprint.yaml | 6 +- src/diffulab/networks/denoisers/ddt.py | 94 ++++---- src/diffulab/networks/denoisers/mmdit.py | 218 ++++++++++++------ src/diffulab/networks/denoisers/sprint.py | 94 ++++---- 10 files changed, 245 insertions(+), 189 deletions(-) create mode 100644 configs/embedder/precomputed.yaml diff --git a/configs/embedder/precomputed.yaml b/configs/embedder/precomputed.yaml new file mode 100644 index 0000000..d16574e --- /dev/null +++ b/configs/embedder/precomputed.yaml @@ -0,0 +1,3 @@ +_target_: diffulab.networks.PrecomputedEmbedder +path_null_embedding: "/path/to/null_embedding.pt" +null_embedding_seq_len: 7 diff --git a/configs/model/ddt.yaml b/configs/model/ddt.yaml index c55a445..d478609 100644 --- a/configs/model/ddt.yaml +++ b/configs/model/ddt.yaml @@ -3,9 +3,7 @@ _target_: diffulab.networks.DDT simple_ddt: true input_channels: 3 output_channels: 3 -input_dim: 512 -hidden_dim: 512 -context_dim: 512 +inner_dim: 512 num_heads: 8 mlp_ratio: 4 patch_size: 2 diff --git a/configs/model/dit.yaml b/configs/model/dit.yaml index 4e63555..244c42a 100644 --- a/configs/model/dit.yaml +++ b/configs/model/dit.yaml @@ -3,8 +3,7 @@ _target_: diffulab.networks.MMDiT simple_dit: true input_channels: 3 output_channels: 3 -input_dim: 512 -hidden_dim: 512 +inner_dim: 512 embedding_dim: 512 num_heads: 8 mlp_ratio: 4 diff --git a/configs/model/sprint.yaml b/configs/model/sprint.yaml index 40c7514..25150b8 100644 --- a/configs/model/sprint.yaml +++ b/configs/model/sprint.yaml @@ -2,9 +2,8 @@ _target_: diffulab.networks.SprintDiT simple_dit: true input_channels: 3 output_channels: 3 -input_dim: 512 -hidden_dim: 512 -context_dim: 512 +inner_dim: 512 +embedding_dim: 512 num_heads: 8 mlp_ratio: 4 patch_size: 2 diff --git a/configs/train_imagenet_flow_matching_repa.yaml b/configs/train_imagenet_flow_matching_repa.yaml index bd484c2..676d567 100644 --- a/configs/train_imagenet_flow_matching_repa.yaml +++ b/configs/train_imagenet_flow_matching_repa.yaml @@ -22,8 +22,7 @@ trainer: model: input_channels: 32 output_channels: 32 - input_dim: 768 - hidden_dim: 768 + inner_dim: 768 embedding_dim: 256 num_heads: 12 mlp_ratio: 4 diff --git a/configs/train_imagenet_repa_txt_to_img.yaml b/configs/train_imagenet_repa_txt_to_img.yaml index 95b63b0..b08d2e3 100644 --- a/configs/train_imagenet_repa_txt_to_img.yaml +++ b/configs/train_imagenet_repa_txt_to_img.yaml @@ -28,9 +28,7 @@ trainer: model: input_channels: 128 output_channels: 128 - input_dim: 640 - hidden_dim: 640 - context_dim: 640 + inner_dim: 640 num_heads: 10 mlp_ratio: 4 patch_size: 1 diff --git a/configs/train_imagenet_repa_txt_to_img_sprint.yaml b/configs/train_imagenet_repa_txt_to_img_sprint.yaml index ae12917..2d3779d 100644 --- a/configs/train_imagenet_repa_txt_to_img_sprint.yaml +++ b/configs/train_imagenet_repa_txt_to_img_sprint.yaml @@ -27,14 +27,14 @@ trainer: model: input_channels: 128 output_channels: 128 - input_dim: 640 - hidden_dim: 640 - context_dim: 640 + inner_dim: 640 + embedding_dim: 640 num_heads: 16 mlp_ratio: 3 patch_size: 1 encoder_depth: 3 deep_layers_depth: 12 + n_single_stream_blocks: 8 decoder_depth: 3 classifier_free: true rope_base: 1000 diff --git a/src/diffulab/networks/denoisers/ddt.py b/src/diffulab/networks/denoisers/ddt.py index 8b1ce89..847f171 100644 --- a/src/diffulab/networks/denoisers/ddt.py +++ b/src/diffulab/networks/denoisers/ddt.py @@ -2,6 +2,7 @@ # Some changes were made to the architecture to reuse DiT and MMDiT components # e.g double stream +import logging from typing import Any import torch @@ -11,34 +12,17 @@ from torch import Tensor from diffulab.networks.denoisers.common import Denoiser, ModelOutput -from diffulab.networks.denoisers.mmdit import DiTBlock, MMDiTBlock +from diffulab.networks.denoisers.mmdit import DiTBlock, MMDiTBlock, MMDiTSingleStreamBlock, ModulatedLastLayer from diffulab.networks.embedders.common import ContextEmbedder, ContextEmbedderOutput from diffulab.networks.utils.nn import ( LabelEmbed, Modulation, get_cos_sin_ndim_grid, - modulate, timestep_embedding, ) from diffulab.networks.utils.utils import zero_module -class ModulatedLastLayerDDT(nn.Module): - def __init__(self, embedding_dim: int, hidden_size: int, patch_size: int, out_channels: int): - super().__init__() # type: ignore - self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) - self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(embedding_dim, 2 * hidden_size, bias=True)) - - def forward( - self, x: Float[Tensor, "batch_size seq_len dim"], vec: Float[Tensor, "batch_size seq_len dim"] - ) -> Tensor: - alpha, beta = self.adaLN_modulation(vec).chunk(2, dim=-1) - x = modulate(self.norm_final(x), scale=alpha, shift=beta) - x = self.linear(x) - return x - - class DDT(Denoiser): """ Decoder-Encoder architecture following https://arxiv.org/pdf/2504.05741 @@ -53,16 +37,12 @@ class DDT(Denoiser): input_channels (int): Number of channels of the main input x. Default: 3. output_channels (int | None): Number of channels to predict. If None, equals `input_channels`. Default: None. - input_dim (int): Token/patch embedding width for both encoder and decoder streams. + inner_dim (int): Token/patch embedding width for both encoder and decoder streams. Also used as the hidden size of the last prediction layer. Default: 768. - hidden_dim (int): Inner attention dimension used by DiT/MMDiT attention blocks. - Default: 768. num_heads (int): Number of attention heads in each block. Default: 12. mlp_ratio (int): Expansion ratio for the MLP in each block. Default: 4. patch_size (int): Side length P of square patches. Images are projected with stride P in both encoder and decoder streams. Default: 16. - context_dim (int): Model width of contextual tokens after `context_embed` when - `simple_ddt=False`. Ignored when `simple_ddt=True`. Default: 1024. encoder_depth (int): Number of DiT/MMDiT blocks in the encoder. Default: 8. decoder_depth (int): Number of DiT blocks in the decoder. Default: 4. partial_rotary_factor (float): Fraction of each head dimension using RoPE. @@ -88,13 +68,12 @@ def __init__( simple_ddt: bool = False, input_channels: int = 3, output_channels: int | None = None, - input_dim: int = 768, - hidden_dim: int = 768, + inner_dim: int = 768, num_heads: int = 12, mlp_ratio: int = 4, patch_size: int = 16, - context_dim: int = 1024, encoder_depth: int = 8, + n_single_stream_blocks: int = 0, decoder_depth: int = 4, rope_base: int = 10_000, partial_rotary_factor: float = 1, @@ -109,6 +88,7 @@ def __init__( assert not (n_classes is not None and context_embedder is not None), ( "n_classes and context_embedder cannot both be specified" ) + assert n_single_stream_blocks < encoder_depth, "n_single_stream_blocks must be less than encoder_depth" self.simple_ddt = simple_ddt self.patch_size = patch_size self.input_channels = input_channels @@ -122,7 +102,7 @@ def __init__( self.n_classes = n_classes self.classifier_free = classifier_free - heads_dim = hidden_dim // num_heads + heads_dim = inner_dim // num_heads if not self.simple_ddt: assert self.context_embedder is not None, "for ddt with text context embedder must be provided" assert isinstance(self.context_embedder.output_size, tuple) and all( @@ -134,14 +114,14 @@ def __init__( if self.context_embedder.n_output == 2: self.pooled_embedding = True self.mlp_pooled_context = nn.Sequential( - nn.Linear(self.context_embedder.output_size[0], input_dim), + nn.Linear(self.context_embedder.output_size[0], inner_dim * 2), nn.SiLU(), - nn.Linear(input_dim, input_dim), + nn.Linear(inner_dim * 2, inner_dim), ) - self.context_embed = nn.Linear(self.context_embedder.output_size[1], context_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim) else: assert self.context_embedder.n_output == 1 - self.context_embed = nn.Linear(self.context_embedder.output_size[0], context_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim) if rope_axes_dim is None: rope_axes_dim = [ int((partial_rotary_factor * heads_dim) // 3), # L for text, set to 0 for image tokens @@ -150,32 +130,37 @@ def __init__( ] else: self.label_embed = ( - LabelEmbed(self.n_classes, input_dim, self.classifier_free) if self.n_classes is not None else None + LabelEmbed(self.n_classes, inner_dim, self.classifier_free) if self.n_classes is not None else None ) if rope_axes_dim is None: rope_axes_dim = [ int((partial_rotary_factor * heads_dim) // 2), # H int((partial_rotary_factor * heads_dim) // 2), # W ] + if n_single_stream_blocks > 0: + logging.warning( + "n_single_stream_blocks is ignored when simple_ddt=True. All blocks are single-stream DiT blocks." + ) + n_single_stream_blocks = encoder_depth self.rope_axes_dim = rope_axes_dim - self.last_layer = ModulatedLastLayerDDT( - embedding_dim=input_dim, - hidden_size=input_dim, + self.last_layer = ModulatedLastLayer( + embedding_dim=inner_dim, + hidden_size=inner_dim, patch_size=self.patch_size, out_channels=self.output_channels, ) self.time_embed = nn.Sequential( - nn.Linear(self.frequency_embedding, input_dim), + nn.Linear(self.frequency_embedding, inner_dim), nn.SiLU(), - nn.Linear(input_dim, input_dim), + nn.Linear(inner_dim, inner_dim), ) self.conv_proj_encoder = nn.Conv2d( - self.input_channels, input_dim, kernel_size=self.patch_size, stride=self.patch_size + self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size ) self.conv_proj_decoder = nn.Conv2d( - self.input_channels, input_dim, kernel_size=self.patch_size, stride=self.patch_size + self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size ) # -------------- @@ -184,10 +169,8 @@ def __init__( self.layers = nn.ModuleList( [ MMDiTBlock( - context_dim=context_dim, - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=inner_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, @@ -195,15 +178,25 @@ def __init__( ) if not self.simple_ddt else DiTBlock( - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=inner_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + use_checkpoint=use_checkpoint, + ) + for _ in range(encoder_depth - n_single_stream_blocks) + ] + + [ + MMDiTSingleStreamBlock( + inner_dim=inner_dim, + embedding_dim=inner_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, use_checkpoint=use_checkpoint, ) - for _ in range(encoder_depth) + for _ in range(n_single_stream_blocks) ] ) @@ -213,9 +206,8 @@ def __init__( self.decoder_layers = nn.ModuleList( [ DiTBlock( - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=inner_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, @@ -234,7 +226,7 @@ def _init_weights(self, module: nn.Module) -> None: nn.init.constant_(module.bias, 0) if isinstance(module, Modulation): zero_module(module) - if isinstance(module, ModulatedLastLayerDDT): + if isinstance(module, ModulatedLastLayer): zero_module(module.adaLN_modulation) def patchify( diff --git a/src/diffulab/networks/denoisers/mmdit.py b/src/diffulab/networks/denoisers/mmdit.py index ff961a4..01ab7bf 100644 --- a/src/diffulab/networks/denoisers/mmdit.py +++ b/src/diffulab/networks/denoisers/mmdit.py @@ -1,5 +1,6 @@ # Recoded from scratch from https://arxiv.org/pdf/2403.03206, if you see any error please report it to the author of the repository +import logging from typing import Any import torch @@ -30,8 +31,7 @@ class DiTAttention(nn.Module): DiTAttention is a multi-head self attention mechanism with rotary positional embeddings. Args: - input_dim (int): Dimension of the input. - dim (int): Inner Attention dimension. + inner_dim (int): Dimension of the input. num_heads (int): Number of attention heads. rope_axes_dim (list[int]): List of dimensions for rotary positional embeddings. @@ -54,23 +54,23 @@ class DiTAttention(nn.Module): Returns: Tensor: output tensor. Example: - >>> dit_attention = DiTAttention(input_dim=512, dim=512, num_heads=8) + >>> dit_attention = DiTAttention(inner_dim=512, num_heads=8) >>> input_tensor = torch.randn(10, 25, 512) >>> cos_sin_rope = (torch.randn(25, 256), torch.randn(25, 256)) >>> output = dit_attention(input_tensor, cos_sin_rope) """ - def __init__(self, input_dim: int, dim: int, num_heads: int, rope_axes_dim: list[int]) -> None: + def __init__(self, inner_dim: int, num_heads: int, rope_axes_dim: list[int]) -> None: super().__init__() # type: ignore self.num_heads = num_heads - self.head_dim = dim // num_heads + self.head_dim = inner_dim // num_heads self.scale = self.head_dim**-0.5 - self.qkv = nn.Linear(input_dim, 3 * dim) - self.qk_norm = QKNorm(dim) + self.qkv = nn.Linear(inner_dim, 3 * inner_dim) + self.qk_norm = QKNorm(inner_dim) self.rope = RotaryPositionalEmbeddingNDim(axes_dim=rope_axes_dim) - self.proj_out = nn.Linear(dim, input_dim) + self.proj_out = nn.Linear(inner_dim, inner_dim) def forward( self, @@ -107,9 +107,7 @@ class MMDiTAttention(nn.Module): MMDiTAttention is a multi-head attention mechanism with rotary positional embeddings. Args: - context_dim (int): Dimension of the context input. - input_dim (int): Dimension of the input. - dim (int): Inner Attention dimension. + inner_dim (int): Dimension of the input. num_heads (int): Number of attention heads. rope_axes_dim (list[int]): List of dimensions for rotary positional embeddings. @@ -149,31 +147,29 @@ class MMDiTAttention(nn.Module): def __init__( self, - context_dim: int, - input_dim: int, - dim: int, + inner_dim: int, num_heads: int, rope_axes_dim: list[int], ): super().__init__() # type: ignore self.num_heads = num_heads - self.head_dim = dim // num_heads + self.head_dim = inner_dim // num_heads self.scale = self.head_dim**-0.5 - self.qkv_input = nn.Linear(input_dim, 3 * dim) - self.qkv_context = nn.Linear(context_dim, 3 * dim) - self.qk_norm_input = QKNorm(dim) - self.qk_norm_context = QKNorm(dim) + self.qkv_input = nn.Linear(inner_dim, 3 * inner_dim) + self.qkv_context = nn.Linear(inner_dim, 3 * inner_dim) + self.qk_norm_input = QKNorm(inner_dim) + self.qk_norm_context = QKNorm(inner_dim) self.rope = RotaryPositionalEmbeddingNDim(axes_dim=rope_axes_dim) - self.input_proj_out = nn.Linear(dim, input_dim) - self.context_proj_out = nn.Linear(dim, context_dim) + self.input_proj_out = nn.Linear(inner_dim, inner_dim) + self.context_proj_out = nn.Linear(inner_dim, inner_dim) def forward( self, - input: Float[Tensor, "batch_size seq_len_input input_dim"], - context: Float[Tensor, "batch_size seq_len_context context_dim"], + input: Float[Tensor, "batch_size seq_len_input inner_dim"], + context: Float[Tensor, "batch_size seq_len_context inner_dim"], cos_sin_rope: tuple[Float[Tensor, "seq_len dim/2"], Float[Tensor, "seq_len dim/2"]], attn_mask: Bool[Tensor, "batch_size seq_len_context"] | Int[Tensor, "batch_size seq_len_context"] | None = None, ) -> tuple[Float[Tensor, "batch_size seq_len input_dim"], Float[Tensor, "batch_size seq_len context_dim"]]: @@ -247,8 +243,7 @@ class DiTBlock(nn.Module): def __init__( self, - input_dim: int, - hidden_dim: int, + inner_dim: int, embedding_dim: int, num_heads: int, mlp_ratio: int, @@ -256,23 +251,23 @@ def __init__( use_checkpoint: bool = False, ): super().__init__() # type: ignore - self.modulation = Modulation(embedding_dim, input_dim) - self.norm_1 = nn.LayerNorm(input_dim) - self.attention = DiTAttention(input_dim, hidden_dim, num_heads, rope_axes_dim=rope_axes_dim) - self.norm_2 = nn.LayerNorm(input_dim) + self.modulation = Modulation(embedding_dim, inner_dim) + self.norm_1 = nn.LayerNorm(inner_dim) + self.attention = DiTAttention(inner_dim, num_heads, rope_axes_dim=rope_axes_dim) + self.norm_2 = nn.LayerNorm(inner_dim) self.mlp_input = nn.Sequential( - nn.Linear(input_dim, mlp_ratio * input_dim * 2), + nn.Linear(inner_dim, mlp_ratio * inner_dim * 2), PackedSwiGLU(), - nn.Linear(mlp_ratio * input_dim, input_dim), + nn.Linear(mlp_ratio * inner_dim, inner_dim), ) self.use_checkpoint = use_checkpoint def forward( self, - input: Float[Tensor, "batch_size seq_len embedding_dim"], + input: Float[Tensor, "batch_size seq_len inner_dim"], y: Float[Tensor, "batch_size embedding_dim"], cos_sin_rope: tuple[Float[Tensor, "seq_len dim/2"], Float[Tensor, "seq_len dim/2"]], - ) -> Float[Tensor, "batch_size seq_len input_dim"]: + ) -> Float[Tensor, "batch_size seq_len inner_dim"]: """ Forward pass of the DiTBlock applying modulation and attention mechanisms. Args: @@ -318,9 +313,7 @@ class MMDiTBlock(nn.Module): attention, and multi-layer perceptron (MLP) operations on input and context tensors. Args: - context_dim (int): Dimension of the context tensor. - input_dim (int): Dimension of the input tensor. - hidden_dim (int): Dimension of the hidden layer in the attention mechanism. + inner_dim (int): Dimension of the input tensor. embedding_dim (int): Dimension of the embedding used in modulation. num_heads (int): Number of attention heads. mlp_ratio (int): Ratio used to determine the size of the MLP layers. @@ -340,7 +333,7 @@ class MMDiTBlock(nn.Module): Tuple[Tensor, Tensor]: The modulated input and context tensors. Example: - >>> mmdit_block = MMDiTBlock(context_dim=512, input_dim=512, hidden_dim=512, embedding_dim=512, num_heads=8, mlp_ratio=4) + >>> mmdit_block = MMDiTBlock(inner_dim=512, embedding_dim=512, num_heads=8, mlp_ratio=4, rope_axes_dim=[256]) >>> input_tensor = torch.randn(10, 25, 512) >>> y = torch.randn(10, 512) >>> context_tensor = torch.randn(10, 32, 512) @@ -352,9 +345,7 @@ class MMDiTBlock(nn.Module): def __init__( self, - context_dim: int, - input_dim: int, - hidden_dim: int, + inner_dim: int, embedding_dim: int, num_heads: int, mlp_ratio: int, @@ -362,34 +353,34 @@ def __init__( use_checkpoint: bool = False, ): super().__init__() # type: ignore - self.modulation_context = Modulation(embedding_dim, context_dim) - self.modulation_input = Modulation(embedding_dim, input_dim) + self.modulation_context = Modulation(embedding_dim, inner_dim) + self.modulation_input = Modulation(embedding_dim, inner_dim) - self.context_norm_1 = nn.LayerNorm(context_dim) - self.input_norm_1 = nn.LayerNorm(input_dim) + self.context_norm_1 = nn.LayerNorm(inner_dim) + self.input_norm_1 = nn.LayerNorm(inner_dim) - self.attention = MMDiTAttention(context_dim, input_dim, hidden_dim, num_heads, rope_axes_dim=rope_axes_dim) + self.attention = MMDiTAttention(inner_dim, num_heads, rope_axes_dim=rope_axes_dim) - self.context_norm_2 = nn.LayerNorm(context_dim) - self.input_norm_2 = nn.LayerNorm(input_dim) + self.context_norm_2 = nn.LayerNorm(inner_dim) + self.input_norm_2 = nn.LayerNorm(inner_dim) self.mlp_context = nn.Sequential( - nn.Linear(context_dim, mlp_ratio * context_dim * 2), + nn.Linear(inner_dim, mlp_ratio * inner_dim * 2), PackedSwiGLU(), - nn.Linear(mlp_ratio * context_dim, context_dim), + nn.Linear(mlp_ratio * inner_dim, inner_dim), ) self.mlp_input = nn.Sequential( - nn.Linear(input_dim, mlp_ratio * input_dim * 2), + nn.Linear(inner_dim, mlp_ratio * inner_dim * 2), PackedSwiGLU(), - nn.Linear(mlp_ratio * input_dim, input_dim), + nn.Linear(mlp_ratio * inner_dim, inner_dim), ) self.use_checkpoint = use_checkpoint def forward( self, - input: Float[Tensor, "batch_size seq_len_input embedding_dim"], + input: Float[Tensor, "batch_size seq_len_input inner_dim"], y: Float[Tensor, "batch_size embedding_dim"], - context: Float[Tensor, "batch_size seq_len_context context_dim"], + context: Float[Tensor, "batch_size seq_len_context inner_dim"], cos_sin_rope: tuple[Float[Tensor, "seq_len dim/2"], Float[Tensor, "seq_len dim/2"]], attn_mask: Bool[Tensor, "batch_size seq_len_context"] | Int[Tensor, "batch_size seq_len_context"] | None = None, ) -> tuple[ @@ -466,16 +457,81 @@ def _forward( return modulated_input, modulated_context +class MMDiTSingleStreamBlock(nn.Module): + def __init__( + self, + inner_dim: int, + embedding_dim: int, + num_heads: int, + mlp_ratio: int, + rope_axes_dim: list[int], + use_checkpoint: bool = False, + ): + super().__init__() # type: ignore + self.mlp = nn.Sequential( + nn.Linear(inner_dim, mlp_ratio * inner_dim * 2), + PackedSwiGLU(), + nn.Linear(mlp_ratio * inner_dim, inner_dim), + ) + self.attention = DiTAttention(inner_dim, num_heads, rope_axes_dim=rope_axes_dim) + self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(embedding_dim, 3 * inner_dim)) + self.norm = nn.LayerNorm(inner_dim) + self.use_checkpoint = use_checkpoint + + def forward( + self, + input: Float[Tensor, "batch_size seq_len inner_dim"], + y: Float[Tensor, "batch_size embedding_dim"], + context: Float[Tensor, "batch_size seq_len inner_dim"], + cos_sin_rope: tuple[Float[Tensor, "seq_len dim/2"], Float[Tensor, "seq_len dim/2"]], + attn_mask: Bool[Tensor, "batch_size seq_len_context"] | Int[Tensor, "batch_size seq_len_context"] | None = None, + ) -> tuple[ + Float[Tensor, "batch_size seq_len_input inner_dim"], Float[Tensor, "batch_size seq_len_context inner_dim"] + ]: + return ( + checkpoint(self._forward, *(input, y, context, cos_sin_rope, attn_mask), use_reentrant=False) + if self.use_checkpoint + else self._forward(input, y, context, cos_sin_rope, attn_mask) + ) # type: ignore + + def _forward( + self, + input: Float[Tensor, "batch_size seq_len inner_dim"], + y: Float[Tensor, "batch_size embedding_dim"], + context: Float[Tensor, "batch_size seq_len inner_dim"], + cos_sin_rope: tuple[Float[Tensor, "seq_len dim/2"], Float[Tensor, "seq_len dim/2"]], + attn_mask: Bool[Tensor, "batch_size seq_len_context"] | Int[Tensor, "batch_size seq_len_context"] | None = None, + ): + latents = torch.cat([context, input], dim=1) + modulation = self.modulation(y) + if modulation.dim() == 2: + modulation = modulation[:, None, :] + alpha, beta, gamma = modulation.chunk(3, dim=-1) + modulated_latents = modulate(self.norm(latents), scale=alpha, shift=beta) + latents = ( + latents + + ( + self.attention(modulated_latents, cos_sin_rope=cos_sin_rope, attn_mask=attn_mask) + + self.mlp(modulated_latents) + ) + * gamma + ) + return latents[:, : context.size(1), :], latents[:, context.size(1) :, :] + + class ModulatedLastLayer(nn.Module): def __init__(self, embedding_dim: int, hidden_size: int, patch_size: int, out_channels: int): super().__init__() # type: ignore self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) - self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(embedding_dim, 2 * hidden_size, bias=True)) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(embedding_dim, 2 * hidden_size)) def forward(self, x: Float[Tensor, "batch_size seq_len dim"], vec: Float[Tensor, "batch_size dim"]) -> Tensor: - alpha, beta = self.adaLN_modulation(vec).chunk(2, dim=1) - x = modulate(self.norm_final(x), scale=alpha[:, None, :], shift=beta[:, None, :]) + modulation = self.adaLN_modulation(vec) + if modulation.dim() == 2: + modulation = modulation[:, None, :] + alpha, beta = modulation.chunk(2, dim=-1) + x = modulate(self.norm_final(x), scale=alpha, shift=beta) x = self.linear(x) return x @@ -499,9 +555,7 @@ class MMDiT(Denoiser): input_channels (int): Number of channels of the main input x. Default: 3. output_channels (int | None): Number of channels to predict. If None, equals `input_channels`. Default: None. - input_dim (int): Token/patch embedding width for the image stream. Also the hidden size used - before the final per-patch projection. Default: 4096. - hidden_dim (int): Inner attention dimension for attention projections. Default: 4096. + inner_dim (int): Token/patch embedding width for the stream. Default: 4096. embedding_dim (int): Conditioning embedding width (for timestep/labels/pooled context) used by modulation layers and the last prediction layer. Default: 4096. num_heads (int): Number of attention heads in each block. Default: 16. @@ -539,14 +593,13 @@ def __init__( simple_dit: bool = False, input_channels: int = 3, output_channels: int | None = None, - input_dim: int = 4096, - hidden_dim: int = 4096, + inner_dim: int = 4096, embedding_dim: int = 4096, num_heads: int = 16, mlp_ratio: int = 4, patch_size: int = 16, depth: int = 38, - context_dim: int = 4096, + n_single_stream_blocks: int = 0, rope_base: int = 10_000, partial_rotary_factor: float = 1, rope_axes_dim: list[int] | None = None, @@ -573,7 +626,7 @@ def __init__( self.n_classes = n_classes self.classifier_free = classifier_free - heads_dim = hidden_dim // num_heads + heads_dim = inner_dim // num_heads if not self.simple_dit: assert self.context_embedder is not None, "for MMDiT context embedder must be provided" assert isinstance(self.context_embedder.output_size, tuple) and all( @@ -585,14 +638,14 @@ def __init__( if self.context_embedder.n_output == 2: self.pooled_embedding = True self.mlp_pooled_context = nn.Sequential( - nn.Linear(self.context_embedder.output_size[0], embedding_dim), + nn.Linear(self.context_embedder.output_size[0], embedding_dim * 2), nn.SiLU(), - nn.Linear(embedding_dim, embedding_dim), + nn.Linear(embedding_dim * 2, embedding_dim), ) - self.context_embed = nn.Linear(self.context_embedder.output_size[1], context_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim) else: assert self.context_embedder.n_output == 1 - self.context_embed = nn.Linear(self.context_embedder.output_size[0], context_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim) if rope_axes_dim is None: rope_axes_dim = [ int((partial_rotary_factor * heads_dim) // 3), # L for text, set to 0 for image tokens @@ -609,11 +662,16 @@ def __init__( int((partial_rotary_factor * heads_dim) // 2), # H int((partial_rotary_factor * heads_dim) // 2), # W ] + if n_single_stream_blocks > 0: + logging.warning( + "n_single_stream_blocks is ignored when simple_dit=True. All blocks are single-stream DiT blocks." + ) + n_single_stream_blocks = depth self.rope_axes_dim = rope_axes_dim self.last_layer = ModulatedLastLayer( embedding_dim=embedding_dim, - hidden_size=input_dim, + hidden_size=inner_dim, patch_size=self.patch_size, out_channels=self.output_channels, ) @@ -623,14 +681,12 @@ def __init__( nn.Linear(embedding_dim, embedding_dim), ) - self.conv_proj = nn.Conv2d(self.input_channels, input_dim, kernel_size=self.patch_size, stride=self.patch_size) + self.conv_proj = nn.Conv2d(self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size) self.layers = nn.ModuleList( [ MMDiTBlock( - context_dim=context_dim, - input_dim=input_dim, - hidden_dim=hidden_dim, + inner_dim=inner_dim, embedding_dim=embedding_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, @@ -639,15 +695,25 @@ def __init__( ) if not self.simple_dit else DiTBlock( - input_dim=input_dim, - hidden_dim=hidden_dim, + inner_dim=inner_dim, + embedding_dim=embedding_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + use_checkpoint=use_checkpoint, + ) + for _ in range(depth - n_single_stream_blocks) + ] + + [ + MMDiTSingleStreamBlock( + inner_dim=inner_dim, embedding_dim=embedding_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, use_checkpoint=use_checkpoint, ) - for _ in range(depth) + for _ in range(n_single_stream_blocks) ] ) diff --git a/src/diffulab/networks/denoisers/sprint.py b/src/diffulab/networks/denoisers/sprint.py index 280e8a1..f63d23e 100644 --- a/src/diffulab/networks/denoisers/sprint.py +++ b/src/diffulab/networks/denoisers/sprint.py @@ -1,3 +1,4 @@ +import logging from typing import Any import torch @@ -7,7 +8,7 @@ from torch import Tensor from diffulab.networks.denoisers.common import Denoiser, ModelOutput -from diffulab.networks.denoisers.mmdit import DiTBlock, MMDiTBlock, ModulatedLastLayer +from diffulab.networks.denoisers.mmdit import DiTBlock, MMDiTBlock, MMDiTSingleStreamBlock, ModulatedLastLayer from diffulab.networks.embedders.common import ContextEmbedder, ContextEmbedderOutput from diffulab.networks.utils.nn import ( LabelEmbed, @@ -29,11 +30,7 @@ class SprintDiT(Denoiser): input_channels (int): Number of channels of the main input x. Default: 3. output_channels (int | None): Number of channels to predict. If None, equals `input_channels`. Default: None. - input_dim (int): Token/patch embedding width for the image stream. Also the hidden size used - before the final per-patch projection. Default: 4096. - hidden_dim (int): Inner attention dimension for attention projections. Default: 4096. - embedding_dim (int): Conditioning embedding width (for timestep/labels/pooled context) used - by modulation layers and the last prediction layer. Default: 4096. + inner_dim (int): Token/patch embedding width for the stream. Default: 4096. num_heads (int): Number of attention heads in each block. Default: 16. mlp_ratio (int): Expansion ratio for the MLP in each block. Default: 4. patch_size (int): Side length P of square patches. Images are projected with stride P. Default: 16. @@ -73,14 +70,14 @@ def __init__( simple_dit: bool = False, input_channels: int = 3, output_channels: int | None = None, - input_dim: int = 768, - hidden_dim: int = 768, + inner_dim: int = 768, + embedding_dim: int = 768, num_heads: int = 12, mlp_ratio: int = 4, patch_size: int = 16, - context_dim: int = 768, encoder_depth: int = 2, deep_layers_depth: int = 8, + n_single_stream_blocks: int = 0, decoder_depth: int = 2, rope_base: int = 10_000, partial_rotary_factor: float = 1, @@ -109,10 +106,10 @@ def __init__( self.n_classes = n_classes self.classifier_free = classifier_free - self.mask_token = nn.Parameter(torch.zeros(1, 1, input_dim)) + self.mask_token = nn.Parameter(torch.zeros(1, 1, inner_dim)) self.drop_rate = drop_rate - heads_dim = hidden_dim // num_heads + heads_dim = inner_dim // num_heads if not self.simple_dit: assert self.context_embedder is not None, "for dit with text context embedder must be provided" assert isinstance(self.context_embedder.output_size, tuple) and all( @@ -124,14 +121,14 @@ def __init__( if self.context_embedder.n_output == 2: self.pooled_embedding = True self.mlp_pooled_context = nn.Sequential( - nn.Linear(self.context_embedder.output_size[0], input_dim), + nn.Linear(self.context_embedder.output_size[0], embedding_dim * 2), nn.SiLU(), - nn.Linear(input_dim, input_dim), + nn.Linear(embedding_dim * 2, embedding_dim), ) - self.context_embed = nn.Linear(self.context_embedder.output_size[1], context_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim) else: assert self.context_embedder.n_output == 1 - self.context_embed = nn.Linear(self.context_embedder.output_size[0], context_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim) if rope_axes_dim is None: rope_axes_dim = [ int((partial_rotary_factor * heads_dim) // 3), # L for text, set to 0 for image tokens @@ -140,30 +137,35 @@ def __init__( ] else: self.label_embed = ( - LabelEmbed(self.n_classes, input_dim, self.classifier_free) if self.n_classes is not None else None + LabelEmbed(self.n_classes, embedding_dim, self.classifier_free) if self.n_classes is not None else None ) if rope_axes_dim is None: rope_axes_dim = [ int((partial_rotary_factor * heads_dim) // 2), # H int((partial_rotary_factor * heads_dim) // 2), # W ] + if n_single_stream_blocks > 0: + logging.warning( + "n_single_stream_blocks is ignored when simple_dit=True. All blocks are single-stream DiT blocks." + ) + n_single_stream_blocks = encoder_depth + deep_layers_depth + decoder_depth + self.rope_axes_dim = rope_axes_dim self.time_embed = nn.Sequential( - nn.Linear(self.frequency_embedding, input_dim), + nn.Linear(self.frequency_embedding, embedding_dim), nn.SiLU(), - nn.Linear(input_dim, input_dim), + nn.Linear(embedding_dim, embedding_dim), ) - self.conv_proj = nn.Conv2d(self.input_channels, input_dim, kernel_size=self.patch_size, stride=self.patch_size) + self.conv_proj = nn.Conv2d(self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size) - self.fuse = nn.Linear(input_dim * 2, input_dim) + self.fuse = nn.Linear(inner_dim * 2, inner_dim) if not self.simple_dit: - self.fuse_context = nn.Linear(2 * context_dim, context_dim) - + self.fuse_context = nn.Linear(2 * inner_dim, inner_dim) self.last_layer = ModulatedLastLayer( - embedding_dim=input_dim, - hidden_size=input_dim, + embedding_dim=embedding_dim, + hidden_size=inner_dim, patch_size=self.patch_size, out_channels=self.output_channels, ) @@ -174,10 +176,8 @@ def __init__( self.layers = nn.ModuleList( # name compatibility for RePA [ MMDiTBlock( - context_dim=context_dim, - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=embedding_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, @@ -185,9 +185,8 @@ def __init__( ) if not self.simple_dit else DiTBlock( - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=embedding_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, @@ -203,10 +202,8 @@ def __init__( self.deep_layers = nn.ModuleList( [ MMDiTBlock( - context_dim=context_dim, - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=embedding_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, @@ -214,15 +211,23 @@ def __init__( ) if not self.simple_dit else DiTBlock( - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=embedding_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, use_checkpoint=use_checkpoint, ) - for _ in range(deep_layers_depth) + for _ in range(deep_layers_depth - n_single_stream_blocks) + ] + + [ + MMDiTSingleStreamBlock( + inner_dim=inner_dim, + embedding_dim=embedding_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + rope_axes_dim=self.rope_axes_dim, + ) ] ) @@ -232,10 +237,8 @@ def __init__( self.decoder_layers = nn.ModuleList( [ MMDiTBlock( - context_dim=context_dim, - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=embedding_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, @@ -243,9 +246,8 @@ def __init__( ) if not self.simple_dit else DiTBlock( - input_dim=input_dim, - hidden_dim=hidden_dim, - embedding_dim=input_dim, + inner_dim=inner_dim, + embedding_dim=embedding_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, From 9a98e049003a7a51b4fd7f35ad58e0e1e3035905 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Fri, 9 Jan 2026 16:10:18 +0100 Subject: [PATCH 10/13] fix single stream block attention --- src/diffulab/networks/denoisers/mmdit.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/diffulab/networks/denoisers/mmdit.py b/src/diffulab/networks/denoisers/mmdit.py index 01ab7bf..91f96d6 100644 --- a/src/diffulab/networks/denoisers/mmdit.py +++ b/src/diffulab/networks/denoisers/mmdit.py @@ -76,6 +76,7 @@ def forward( self, input: Float[Tensor, "batch_size seq_len dim"], cos_sin_rope: tuple[Float[Tensor, "seq_len dim/2"], Float[Tensor, "seq_len dim/2"]], + attn_mask: Bool[Tensor, "batch_size seq_len"] | Int[Tensor, "batch_size seq_len"] | None = None, ) -> Float[Tensor, "batch_size seq_len dim"]: input_q, input_k, input_v = self.qkv(input).chunk(3, dim=-1) input_q, input_k = self.qk_norm(input_q, input_k, input_v) @@ -93,6 +94,7 @@ def forward( key=k, value=v, scale=self.scale, + attn_mask=attn_mask.bool() if attn_mask is not None else None, ) attn_output = rearrange(attn_output, "b h n d -> b n (h d)") @@ -503,11 +505,22 @@ def _forward( attn_mask: Bool[Tensor, "batch_size seq_len_context"] | Int[Tensor, "batch_size seq_len_context"] | None = None, ): latents = torch.cat([context, input], dim=1) + if attn_mask is not None: + attn_mask = torch.cat( + [ + attn_mask.bool(), + torch.ones(attn_mask.size(0), input.size(1), device=attn_mask.device).bool(), + ], + dim=1, + ) + attn_mask = attn_mask[:, None, None, :] + modulation = self.modulation(y) if modulation.dim() == 2: modulation = modulation[:, None, :] alpha, beta, gamma = modulation.chunk(3, dim=-1) modulated_latents = modulate(self.norm(latents), scale=alpha, shift=beta) + latents = ( latents + ( @@ -516,7 +529,7 @@ def _forward( ) * gamma ) - return latents[:, : context.size(1), :], latents[:, context.size(1) :, :] + return latents[:, context.size(1) :, :], latents[:, : context.size(1), :] class ModulatedLastLayer(nn.Module): From 459ac69e08235422c67bd76d941b5aaa6fb6c9ae Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Fri, 9 Jan 2026 22:38:49 +0100 Subject: [PATCH 11/13] remove bias when not needed --- src/diffulab/networks/denoisers/ddt.py | 8 ++--- src/diffulab/networks/denoisers/mmdit.py | 36 ++++++++++++----------- src/diffulab/networks/denoisers/sprint.py | 12 ++++---- 3 files changed, 30 insertions(+), 26 deletions(-) diff --git a/src/diffulab/networks/denoisers/ddt.py b/src/diffulab/networks/denoisers/ddt.py index 847f171..c4ce3d7 100644 --- a/src/diffulab/networks/denoisers/ddt.py +++ b/src/diffulab/networks/denoisers/ddt.py @@ -118,10 +118,10 @@ def __init__( nn.SiLU(), nn.Linear(inner_dim * 2, inner_dim), ) - self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim, bias=False) else: assert self.context_embedder.n_output == 1 - self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim, bias=False) if rope_axes_dim is None: rope_axes_dim = [ int((partial_rotary_factor * heads_dim) // 3), # L for text, set to 0 for image tokens @@ -157,10 +157,10 @@ def __init__( ) self.conv_proj_encoder = nn.Conv2d( - self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size + self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=False ) self.conv_proj_decoder = nn.Conv2d( - self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size + self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=False ) # -------------- diff --git a/src/diffulab/networks/denoisers/mmdit.py b/src/diffulab/networks/denoisers/mmdit.py index 91f96d6..244434e 100644 --- a/src/diffulab/networks/denoisers/mmdit.py +++ b/src/diffulab/networks/denoisers/mmdit.py @@ -67,10 +67,10 @@ def __init__(self, inner_dim: int, num_heads: int, rope_axes_dim: list[int]) -> self.head_dim = inner_dim // num_heads self.scale = self.head_dim**-0.5 - self.qkv = nn.Linear(inner_dim, 3 * inner_dim) + self.qkv = nn.Linear(inner_dim, 3 * inner_dim, bias=False) self.qk_norm = QKNorm(inner_dim) self.rope = RotaryPositionalEmbeddingNDim(axes_dim=rope_axes_dim) - self.proj_out = nn.Linear(inner_dim, inner_dim) + self.proj_out = nn.Linear(inner_dim, inner_dim, bias=False) def forward( self, @@ -158,15 +158,15 @@ def __init__( self.head_dim = inner_dim // num_heads self.scale = self.head_dim**-0.5 - self.qkv_input = nn.Linear(inner_dim, 3 * inner_dim) - self.qkv_context = nn.Linear(inner_dim, 3 * inner_dim) + self.qkv_input = nn.Linear(inner_dim, 3 * inner_dim, bias=False) + self.qkv_context = nn.Linear(inner_dim, 3 * inner_dim, bias=False) self.qk_norm_input = QKNorm(inner_dim) self.qk_norm_context = QKNorm(inner_dim) self.rope = RotaryPositionalEmbeddingNDim(axes_dim=rope_axes_dim) - self.input_proj_out = nn.Linear(inner_dim, inner_dim) - self.context_proj_out = nn.Linear(inner_dim, inner_dim) + self.input_proj_out = nn.Linear(inner_dim, inner_dim, bias=False) + self.context_proj_out = nn.Linear(inner_dim, inner_dim, bias=False) def forward( self, @@ -258,9 +258,9 @@ def __init__( self.attention = DiTAttention(inner_dim, num_heads, rope_axes_dim=rope_axes_dim) self.norm_2 = nn.LayerNorm(inner_dim) self.mlp_input = nn.Sequential( - nn.Linear(inner_dim, mlp_ratio * inner_dim * 2), + nn.Linear(inner_dim, mlp_ratio * inner_dim * 2, bias=False), PackedSwiGLU(), - nn.Linear(mlp_ratio * inner_dim, inner_dim), + nn.Linear(mlp_ratio * inner_dim, inner_dim, bias=False), ) self.use_checkpoint = use_checkpoint @@ -367,14 +367,14 @@ def __init__( self.input_norm_2 = nn.LayerNorm(inner_dim) self.mlp_context = nn.Sequential( - nn.Linear(inner_dim, mlp_ratio * inner_dim * 2), + nn.Linear(inner_dim, mlp_ratio * inner_dim * 2, bias=False), PackedSwiGLU(), - nn.Linear(mlp_ratio * inner_dim, inner_dim), + nn.Linear(mlp_ratio * inner_dim, inner_dim, bias=False), ) self.mlp_input = nn.Sequential( - nn.Linear(inner_dim, mlp_ratio * inner_dim * 2), + nn.Linear(inner_dim, mlp_ratio * inner_dim * 2, bias=False), PackedSwiGLU(), - nn.Linear(mlp_ratio * inner_dim, inner_dim), + nn.Linear(mlp_ratio * inner_dim, inner_dim, bias=False), ) self.use_checkpoint = use_checkpoint @@ -471,9 +471,9 @@ def __init__( ): super().__init__() # type: ignore self.mlp = nn.Sequential( - nn.Linear(inner_dim, mlp_ratio * inner_dim * 2), + nn.Linear(inner_dim, mlp_ratio * inner_dim * 2, bias=False), PackedSwiGLU(), - nn.Linear(mlp_ratio * inner_dim, inner_dim), + nn.Linear(mlp_ratio * inner_dim, inner_dim, bias=False), ) self.attention = DiTAttention(inner_dim, num_heads, rope_axes_dim=rope_axes_dim) self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(embedding_dim, 3 * inner_dim)) @@ -655,10 +655,10 @@ def __init__( nn.SiLU(), nn.Linear(embedding_dim * 2, embedding_dim), ) - self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim, bias=False) else: assert self.context_embedder.n_output == 1 - self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim, bias=False) if rope_axes_dim is None: rope_axes_dim = [ int((partial_rotary_factor * heads_dim) // 3), # L for text, set to 0 for image tokens @@ -694,7 +694,9 @@ def __init__( nn.Linear(embedding_dim, embedding_dim), ) - self.conv_proj = nn.Conv2d(self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size) + self.conv_proj = nn.Conv2d( + self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=False + ) self.layers = nn.ModuleList( [ diff --git a/src/diffulab/networks/denoisers/sprint.py b/src/diffulab/networks/denoisers/sprint.py index f63d23e..cdd9379 100644 --- a/src/diffulab/networks/denoisers/sprint.py +++ b/src/diffulab/networks/denoisers/sprint.py @@ -125,10 +125,10 @@ def __init__( nn.SiLU(), nn.Linear(embedding_dim * 2, embedding_dim), ) - self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[1], inner_dim, bias=False) else: assert self.context_embedder.n_output == 1 - self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim) + self.context_embed = nn.Linear(self.context_embedder.output_size[0], inner_dim, bias=False) if rope_axes_dim is None: rope_axes_dim = [ int((partial_rotary_factor * heads_dim) // 3), # L for text, set to 0 for image tokens @@ -158,11 +158,13 @@ def __init__( nn.Linear(embedding_dim, embedding_dim), ) - self.conv_proj = nn.Conv2d(self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size) + self.conv_proj = nn.Conv2d( + self.input_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=False + ) - self.fuse = nn.Linear(inner_dim * 2, inner_dim) + self.fuse = nn.Linear(inner_dim * 2, inner_dim, bias=False) if not self.simple_dit: - self.fuse_context = nn.Linear(2 * inner_dim, inner_dim) + self.fuse_context = nn.Linear(2 * inner_dim, inner_dim, bias=False) self.last_layer = ModulatedLastLayer( embedding_dim=embedding_dim, hidden_size=inner_dim, From 2022370d4280c7143b7d5e9a1c61a5260da373b7 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Sat, 10 Jan 2026 18:25:10 +0100 Subject: [PATCH 12/13] fix loop for single stream blocks in SprintDiT --- src/diffulab/networks/denoisers/sprint.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/diffulab/networks/denoisers/sprint.py b/src/diffulab/networks/denoisers/sprint.py index cdd9379..c16a2f2 100644 --- a/src/diffulab/networks/denoisers/sprint.py +++ b/src/diffulab/networks/denoisers/sprint.py @@ -230,6 +230,7 @@ def __init__( mlp_ratio=mlp_ratio, rope_axes_dim=self.rope_axes_dim, ) + for _ in range(n_single_stream_blocks) ] ) From 111566821057451a9bcf81d1b18ed154db837dd6 Mon Sep 17 00:00:00 2001 From: LouisRouss Date: Sat, 10 Jan 2026 18:34:15 +0100 Subject: [PATCH 13/13] update config sprint imagenet --- ...train_imagenet_repa_txt_to_img_sprint.yaml | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/configs/train_imagenet_repa_txt_to_img_sprint.yaml b/configs/train_imagenet_repa_txt_to_img_sprint.yaml index 2d3779d..04fcb93 100644 --- a/configs/train_imagenet_repa_txt_to_img_sprint.yaml +++ b/configs/train_imagenet_repa_txt_to_img_sprint.yaml @@ -27,18 +27,18 @@ trainer: model: input_channels: 128 output_channels: 128 - inner_dim: 640 - embedding_dim: 640 - num_heads: 16 - mlp_ratio: 3 + inner_dim: 768 + embedding_dim: 768 + num_heads: 12 + mlp_ratio: 4 patch_size: 1 - encoder_depth: 3 - deep_layers_depth: 12 + encoder_depth: 2 + deep_layers_depth: 8 n_single_stream_blocks: 8 - decoder_depth: 3 + decoder_depth: 2 classifier_free: true - rope_base: 1000 - rope_axes_dim: [12, 14, 14] + rope_base: 2000 + rope_axes_dim: [16, 24, 24] n_classes: null simple_dit: false @@ -48,7 +48,7 @@ diffuser: shift: 4.63 dataloader: - batch_size: 32 + batch_size: 64 # Hydra configuration hydra: