# SPDX-License-Identifier: Apache-2.0 from __future__ import annotations from contextlib import nullcontext from dataclasses import dataclass from typing import Any import torch import torch.nn as nn import torch.nn.functional as F from fastvideo.attention import DistributedAttention from fastvideo.configs.models import DiTConfig from fastvideo.forward_context import get_forward_context, set_forward_context from fastvideo.layers.activation import get_act_fn from fastvideo.layers.layernorm import RMSNorm from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.visual_embedding import Timesteps from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum @dataclass class SD3Transformer2DModelOutput: sample: torch.Tensor def _chunked_feed_forward( ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int, ) -> torch.Tensor: if hidden_states.shape[chunk_dim] % chunk_size != 0: raise ValueError( f"hidden_states.shape[{chunk_dim}] ({hidden_states.shape[chunk_dim]})" f" must be divisible by chunk_size ({chunk_size})" ) num_chunks = hidden_states.shape[chunk_dim] // chunk_size return torch.cat( [ff(h) for h in hidden_states.chunk(num_chunks, dim=chunk_dim)], dim=chunk_dim, ) def _get_1d_sincos_pos_embed_from_grid( embed_dim: int, pos: torch.Tensor, dtype: torch.dtype | None = None, ) -> torch.Tensor: if embed_dim % 2 != 0: raise ValueError("embed_dim must be divisible by 2") if dtype is None: dtype = torch.float32 if pos.device.type == "mps" else torch.float64 omega = torch.arange(embed_dim // 2, device=pos.device, dtype=dtype) omega /= embed_dim / 2.0 omega = 1.0 / (10000**omega) pos = pos.reshape(-1) out = torch.outer(pos, omega) emb_sin = torch.sin(out) emb_cos = torch.cos(out) return torch.cat([emb_sin, emb_cos], dim=1) def _get_2d_sincos_pos_embed_from_grid( embed_dim: int, grid: torch.Tensor, ) -> torch.Tensor: if embed_dim % 2 != 0: raise ValueError("embed_dim must be divisible by 2") emb_h = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) emb_w = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) return torch.cat([emb_h, emb_w], dim=1) def _get_2d_sincos_pos_embed( embed_dim: int, grid_size: int | tuple[int, int], interpolation_scale: float = 1.0, base_size: int = 16, device: torch.device | None = None, ) -> torch.Tensor: if isinstance(grid_size, int): grid_size = (grid_size, grid_size) grid_h = ( torch.arange(grid_size[0], device=device, dtype=torch.float32) / (grid_size[0] / base_size) / interpolation_scale ) grid_w = ( torch.arange(grid_size[1], device=device, dtype=torch.float32) / (grid_size[1] / base_size) / interpolation_scale ) grid = torch.meshgrid(grid_w, grid_h, indexing="xy") grid = torch.stack(grid, dim=0) grid = grid.reshape([2, 1, grid_size[1], grid_size[0]]) return _get_2d_sincos_pos_embed_from_grid(embed_dim, grid) class SD3PatchEmbed(nn.Module): """2D patch embedding with SD3 positional embedding cropping behavior.""" def __init__( self, height: int = 224, width: int = 224, patch_size: int = 16, in_channels: int = 3, embed_dim: int = 768, layer_norm: bool = False, flatten: bool = True, bias: bool = True, interpolation_scale: float = 1.0, pos_embed_type: str | None = "sincos", pos_embed_max_size: int | None = None, ) -> None: super().__init__() num_patches = (height // patch_size) * (width // patch_size) self.flatten = flatten self.layer_norm = layer_norm self.pos_embed_max_size = pos_embed_max_size self.proj = nn.Conv2d( in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias, ) if layer_norm: self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) else: self.norm = None self.patch_size = patch_size self.height = height // patch_size self.width = width // patch_size self.base_size = height // patch_size self.interpolation_scale = interpolation_scale grid_size = pos_embed_max_size or int(num_patches ** 0.5) if pos_embed_type is None: self.pos_embed = None elif pos_embed_type == "sincos": pos_embed = _get_2d_sincos_pos_embed( embed_dim, grid_size, base_size=self.base_size, interpolation_scale=self.interpolation_scale, ) persistent = True if pos_embed_max_size else False self.register_buffer( "pos_embed", pos_embed.float().unsqueeze(0), persistent=persistent, ) else: raise ValueError(f"Unsupported pos_embed_type: {pos_embed_type}") def cropped_pos_embed(self, height: int, width: int) -> torch.Tensor: if self.pos_embed_max_size is None: raise ValueError("`pos_embed_max_size` must be set for cropping") height = height // self.patch_size width = width // self.patch_size if height > self.pos_embed_max_size: raise ValueError( f"Height ({height}) cannot be greater than pos_embed_max_size " f"({self.pos_embed_max_size})" ) if width > self.pos_embed_max_size: raise ValueError( f"Width ({width}) cannot be greater than pos_embed_max_size " f"({self.pos_embed_max_size})" ) top = (self.pos_embed_max_size - height) // 2 left = (self.pos_embed_max_size - width) // 2 spatial = self.pos_embed.reshape( 1, self.pos_embed_max_size, self.pos_embed_max_size, -1, ) spatial = spatial[:, top:top + height, left:left + width, :] return spatial.reshape(1, -1, spatial.shape[-1]) def forward(self, latent: torch.Tensor) -> torch.Tensor: if self.pos_embed_max_size is not None: height, width = latent.shape[-2:] else: height = latent.shape[-2] // self.patch_size width = latent.shape[-1] // self.patch_size latent = self.proj(latent) if self.flatten: latent = latent.flatten(2).transpose(1, 2) if self.layer_norm: latent = self.norm(latent) if self.pos_embed is None: return latent.to(latent.dtype) if self.pos_embed_max_size: pos_embed = self.cropped_pos_embed(height, width) else: if self.height != height or self.width != width: pos_embed = _get_2d_sincos_pos_embed( embed_dim=self.pos_embed.shape[-1], grid_size=(height, width), base_size=self.base_size, interpolation_scale=self.interpolation_scale, device=latent.device, ) pos_embed = pos_embed.float().unsqueeze(0) else: pos_embed = self.pos_embed return (latent + pos_embed).to(latent.dtype) class SD3TimestepEmbedding(nn.Module): def __init__( self, in_channels: int, time_embed_dim: int, act_fn: str = "silu", ) -> None: super().__init__() self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True) self.act = get_act_fn(act_fn) self.linear_2 = ReplicatedLinear(time_embed_dim, time_embed_dim, bias=True) def forward(self, sample: torch.Tensor) -> torch.Tensor: sample, _ = self.linear_1(sample) sample = self.act(sample) sample, _ = self.linear_2(sample) return sample class SD3TextProjection(nn.Module): def __init__( self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "silu", ) -> None: super().__init__() if out_features is None: out_features = hidden_size self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True) self.act_1 = get_act_fn(act_fn) self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True) def forward(self, caption: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.linear_1(caption) hidden_states = self.act_1(hidden_states) hidden_states, _ = self.linear_2(hidden_states) return hidden_states class CombinedTimestepTextProjEmbeddings(nn.Module): def __init__(self, embedding_dim: int, pooled_projection_dim: int): super().__init__() self.time_proj = Timesteps( num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, ) self.timestep_embedder = SD3TimestepEmbedding( in_channels=256, time_embed_dim=embedding_dim, act_fn="silu", ) self.text_embedder = SD3TextProjection( pooled_projection_dim, embedding_dim, act_fn="silu", ) def forward( self, timestep: torch.Tensor, pooled_projection: torch.Tensor, ) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder( timesteps_proj.to(dtype=pooled_projection.dtype)) pooled_projections = self.text_embedder(pooled_projection) return timesteps_emb + pooled_projections class SD35AdaLayerNormZeroX(nn.Module): def __init__( self, embedding_dim: int, norm_type: str = "layer_norm", bias: bool = True, ) -> None: super().__init__() self.silu = nn.SiLU() self.linear = ReplicatedLinear(embedding_dim, 9 * embedding_dim, bias=bias) if norm_type != "layer_norm": raise ValueError( f"Unsupported norm_type ({norm_type}); expected layer_norm" ) self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) def forward( self, hidden_states: torch.Tensor, emb: torch.Tensor, ) -> tuple[torch.Tensor, ...]: emb, _ = self.linear(self.silu(emb)) (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2) = emb.chunk(9, dim=1) norm_hidden_states = self.norm(hidden_states) hidden_states = norm_hidden_states * ( 1 + scale_msa[:, None]) + shift_msa[:, None] norm_hidden_states2 = norm_hidden_states * ( 1 + scale_msa2[:, None]) + shift_msa2[:, None] return ( hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2, ) class SD3AdaLayerNormZero(nn.Module): def __init__( self, embedding_dim: int, norm_type: str = "layer_norm", bias: bool = True, ) -> None: super().__init__() self.silu = nn.SiLU() self.linear = ReplicatedLinear(embedding_dim, 6 * embedding_dim, bias=bias) if norm_type != "layer_norm": raise ValueError( f"Unsupported norm_type ({norm_type}); expected layer_norm" ) self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) def forward( self, x: torch.Tensor, emb: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: emb, _ = self.linear(self.silu(emb)) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( emb.chunk(6, dim=1)) x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] return x, gate_msa, shift_mlp, scale_mlp, gate_mlp class SD3AdaLayerNormContinuous(nn.Module): def __init__( self, embedding_dim: int, conditioning_embedding_dim: int, elementwise_affine: bool = True, eps: float = 1e-5, bias: bool = True, norm_type: str = "layer_norm", ) -> None: super().__init__() self.silu = nn.SiLU() self.linear = ReplicatedLinear( conditioning_embedding_dim, embedding_dim * 2, bias=bias, ) if norm_type != "layer_norm": raise ValueError(f"Unsupported norm_type: {norm_type}") self.norm = nn.LayerNorm( embedding_dim, eps=eps, elementwise_affine=elementwise_affine, ) def forward( self, x: torch.Tensor, conditioning_embedding: torch.Tensor, ) -> torch.Tensor: emb, _ = self.linear(self.silu(conditioning_embedding).to(x.dtype)) scale, shift = torch.chunk(emb, 2, dim=1) x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] return x class SD3GELU(nn.Module): def __init__( self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True, ) -> None: super().__init__() self.proj = ReplicatedLinear(dim_in, dim_out, bias=bias) self.approximate = approximate def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.proj(hidden_states) return F.gelu(hidden_states, approximate=self.approximate) class SD3FeedForward(nn.Module): def __init__( self, dim: int, dim_out: int | None = None, mult: int = 4, dropout: float = 0.0, activation_fn: str = "geglu", final_dropout: bool = False, inner_dim: int | None = None, bias: bool = True, ) -> None: super().__init__() if inner_dim is None: inner_dim = int(dim * mult) if dim_out is None: dim_out = dim if activation_fn == "gelu-approximate": act_fn: nn.Module = SD3GELU( dim, inner_dim, approximate="tanh", bias=bias, ) elif activation_fn == "gelu": act_fn = SD3GELU(dim, inner_dim, approximate="none", bias=bias) else: raise ValueError(f"Unsupported activation_fn: {activation_fn}") self.net = nn.ModuleList() self.net.append(act_fn) self.net.append(nn.Dropout(dropout)) self.net.append(ReplicatedLinear(inner_dim, dim_out, bias=bias)) if final_dropout: self.net.append(nn.Dropout(dropout)) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: for module in self.net: if isinstance(module, ReplicatedLinear): hidden_states, _ = module(hidden_states) else: hidden_states = module(hidden_states) return hidden_states def _build_qk_norm( qk_norm: str | None, dim_head: int, eps: float, ) -> tuple[nn.Module | None, nn.Module | None]: if qk_norm is None: return None, None if qk_norm == "layer_norm": return ( nn.LayerNorm(dim_head, eps=eps, elementwise_affine=True), nn.LayerNorm(dim_head, eps=eps, elementwise_affine=True), ) if qk_norm == "rms_norm": return ( RMSNorm(dim_head, eps=eps, has_weight=True), RMSNorm(dim_head, eps=eps, has_weight=True), ) raise ValueError(f"Unsupported qk_norm: {qk_norm}") class SD3Attention(nn.Module): """Joint self-attention used in SD3 blocks.""" def __init__( self, query_dim: int, heads: int, dim_head: int, out_dim: int, added_kv_proj_dim: int | None = None, context_pre_only: bool | None = None, qk_norm: str | None = None, eps: float = 1e-6, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, ) -> None: super().__init__() self.heads = heads self.head_dim = dim_head self.inner_dim = out_dim self.context_pre_only = context_pre_only self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=True) self.to_k = ReplicatedLinear(query_dim, self.inner_dim, bias=True) self.to_v = ReplicatedLinear(query_dim, self.inner_dim, bias=True) self.norm_q, self.norm_k = _build_qk_norm(qk_norm, dim_head, eps) self.add_q_proj: ReplicatedLinear | None = None self.add_k_proj: ReplicatedLinear | None = None self.add_v_proj: ReplicatedLinear | None = None self.norm_added_q: nn.Module | None = None self.norm_added_k: nn.Module | None = None if added_kv_proj_dim is not None: self.add_q_proj = ReplicatedLinear( added_kv_proj_dim, self.inner_dim, bias=True, ) self.add_k_proj = ReplicatedLinear( added_kv_proj_dim, self.inner_dim, bias=True, ) self.add_v_proj = ReplicatedLinear( added_kv_proj_dim, self.inner_dim, bias=True, ) self.norm_added_q, self.norm_added_k = _build_qk_norm( qk_norm, dim_head, eps, ) self.to_out = nn.ModuleList([ ReplicatedLinear(self.inner_dim, out_dim, bias=True), nn.Dropout(0.0), ]) if context_pre_only is not None and not context_pre_only: self.to_add_out = ReplicatedLinear(self.inner_dim, out_dim, bias=True) else: self.to_add_out = None self.attn = DistributedAttention( num_heads=heads, head_size=dim_head, causal=False, supported_attention_backends=supported_attention_backends, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor: batch_size = hidden_states.shape[0] residual_seq_len = hidden_states.shape[1] query, _ = self.to_q(hidden_states) key, _ = self.to_k(hidden_states) value, _ = self.to_v(hidden_states) query = query.view(batch_size, -1, self.heads, self.head_dim) key = key.view(batch_size, -1, self.heads, self.head_dim) value = value.view(batch_size, -1, self.heads, self.head_dim) if self.norm_q is not None: query = self.norm_q(query) if self.norm_k is not None: key = self.norm_k(key) if encoder_hidden_states is not None: if (self.add_q_proj is None or self.add_k_proj is None or self.add_v_proj is None): raise RuntimeError( "encoder_hidden_states provided but attention does not have" " added projections" ) encoder_query, _ = self.add_q_proj(encoder_hidden_states) encoder_key, _ = self.add_k_proj(encoder_hidden_states) encoder_value, _ = self.add_v_proj(encoder_hidden_states) encoder_query = encoder_query.view( batch_size, -1, self.heads, self.head_dim, ) encoder_key = encoder_key.view( batch_size, -1, self.heads, self.head_dim, ) encoder_value = encoder_value.view( batch_size, -1, self.heads, self.head_dim, ) if self.norm_added_q is not None: encoder_query = self.norm_added_q(encoder_query) if self.norm_added_k is not None: encoder_key = self.norm_added_k(encoder_key) query = torch.cat([query, encoder_query], dim=1) key = torch.cat([key, encoder_key], dim=1) value = torch.cat([value, encoder_value], dim=1) attn_output, _ = self.attn(query, key, value) attn_output = attn_output.reshape(batch_size, -1, self.heads * self.head_dim) if encoder_hidden_states is not None: hidden_out = attn_output[:, :residual_seq_len] context_out = attn_output[:, residual_seq_len:] if self.to_add_out is not None: context_out, _ = self.to_add_out(context_out) else: hidden_out = attn_output context_out = None hidden_out, _ = self.to_out[0](hidden_out) hidden_out = self.to_out[1](hidden_out) if context_out is not None: return hidden_out, context_out return hidden_out class SD3JointTransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, context_pre_only: bool = False, qk_norm: str | None = None, use_dual_attention: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, ) -> None: super().__init__() self.use_dual_attention = use_dual_attention self.context_pre_only = context_pre_only if use_dual_attention: self.norm1 = SD35AdaLayerNormZeroX(dim) else: self.norm1 = SD3AdaLayerNormZero(dim) if context_pre_only: self.norm1_context = SD3AdaLayerNormContinuous( dim, dim, elementwise_affine=False, eps=1e-6, bias=True, norm_type="layer_norm", ) else: self.norm1_context = SD3AdaLayerNormZero(dim) self.attn = SD3Attention( query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, added_kv_proj_dim=dim, out_dim=dim, context_pre_only=context_pre_only, qk_norm=qk_norm, eps=1e-6, supported_attention_backends=supported_attention_backends, ) if use_dual_attention: self.attn2 = SD3Attention( query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, out_dim=dim, added_kv_proj_dim=None, context_pre_only=None, qk_norm=qk_norm, eps=1e-6, supported_attention_backends=supported_attention_backends, ) else: self.attn2 = None self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) self.ff = SD3FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) if not context_pre_only: self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) self.ff_context = SD3FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) else: self.norm2_context = None self.ff_context = None self._chunk_size: int | None = None self._chunk_dim = 0 def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0): self._chunk_size = chunk_size self._chunk_dim = dim def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, joint_attention_kwargs: dict[str, Any] | None = None, ) -> tuple[torch.Tensor | None, torch.Tensor]: del joint_attention_kwargs if self.use_dual_attention: ( norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2, ) = self.norm1(hidden_states, emb=temb) else: ( norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, ) = self.norm1(hidden_states, emb=temb) if self.context_pre_only: norm_encoder_hidden_states = self.norm1_context( encoder_hidden_states, temb, ) else: ( norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp, ) = self.norm1_context(encoder_hidden_states, emb=temb) attn_output, context_attn_output = self.attn( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, ) hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output if self.use_dual_attention: assert self.attn2 is not None attn_output2 = self.attn2(hidden_states=norm_hidden_states2) hidden_states = hidden_states + gate_msa2.unsqueeze(1) * attn_output2 norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * ( 1 + scale_mlp[:, None]) + shift_mlp[:, None] if self._chunk_size is not None: ff_output = _chunked_feed_forward( self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size, ) else: ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output if self.context_pre_only: encoder_hidden_states = None else: assert context_attn_output is not None assert self.norm2_context is not None assert self.ff_context is not None context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output encoder_hidden_states = encoder_hidden_states + context_attn_output norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) norm_encoder_hidden_states = norm_encoder_hidden_states * ( 1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] if self._chunk_size is not None: context_ff_output = _chunked_feed_forward( self.ff_context, norm_encoder_hidden_states, self._chunk_dim, self._chunk_size, ) else: context_ff_output = self.ff_context(norm_encoder_hidden_states) encoder_hidden_states = encoder_hidden_states + ( c_gate_mlp.unsqueeze(1) * context_ff_output) return encoder_hidden_states, hidden_states class SD3Transformer2DModel(BaseDiT): """FastVideo-native SD3 Transformer2DModel.""" _fsdp_shard_conditions = [ lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit(), ] _compile_conditions = _fsdp_shard_conditions param_names_mapping: dict[str, Any] = {} reverse_param_names_mapping: dict[str, Any] = {} lora_param_names_mapping: dict[str, Any] = {} _supported_attention_backends = ( AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, ) def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs): del kwargs super().__init__(config=config, hf_config=hf_config) self.fastvideo_config = config self.hf_config = hf_config arch = config.arch_config self.out_channels = arch.out_channels self.inner_dim = arch.num_attention_heads * arch.attention_head_dim self.patch_size = arch.patch_size self.num_layers = arch.num_layers self.hidden_size = self.inner_dim self.num_attention_heads = arch.num_attention_heads self.num_channels_latents = arch.in_channels self.pos_embed = SD3PatchEmbed( height=arch.sample_size, width=arch.sample_size, patch_size=arch.patch_size, in_channels=arch.in_channels, embed_dim=self.inner_dim, pos_embed_max_size=arch.pos_embed_max_size, ) self.time_text_embed = CombinedTimestepTextProjEmbeddings( embedding_dim=self.inner_dim, pooled_projection_dim=arch.pooled_projection_dim, ) self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, arch.caption_projection_dim) dual_layers = getattr(arch, "dual_attention_layers", ()) dual_layers = tuple(dual_layers) if isinstance(dual_layers, list) else dual_layers self.transformer_blocks = nn.ModuleList([ SD3JointTransformerBlock( dim=self.inner_dim, num_attention_heads=arch.num_attention_heads, attention_head_dim=arch.attention_head_dim, context_pre_only=i == arch.num_layers - 1, qk_norm=arch.qk_norm, use_dual_attention=i in dual_layers, supported_attention_backends=self._supported_attention_backends, ) for i in range(arch.num_layers) ]) self.norm_out = SD3AdaLayerNormContinuous( self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6, bias=True, norm_type="layer_norm", ) self.proj_out = ReplicatedLinear( self.inner_dim, arch.patch_size * arch.patch_size * self.out_channels, bias=True, ) self.gradient_checkpointing = False self.__post_init__() def enable_forward_chunking( self, chunk_size: int | None = None, dim: int = 0, ) -> None: if dim not in [0, 1]: raise ValueError(f"dim must be 0 or 1, got {dim}") chunk_size = chunk_size or 1 def fn_recursive_feed_forward(module: torch.nn.Module, chunk: int, chunk_dim: int): if hasattr(module, "set_chunk_feed_forward"): module.set_chunk_feed_forward(chunk_size=chunk, dim=chunk_dim) for child in module.children(): fn_recursive_feed_forward(child, chunk, chunk_dim) for module in self.children(): fn_recursive_feed_forward(module, chunk_size, dim) def disable_forward_chunking(self): def fn_recursive_feed_forward(module: torch.nn.Module): if hasattr(module, "set_chunk_feed_forward"): module.set_chunk_feed_forward(chunk_size=None, dim=0) for child in module.children(): fn_recursive_feed_forward(child) for module in self.children(): fn_recursive_feed_forward(module) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, pooled_projections: torch.Tensor | None = None, timestep: torch.LongTensor | None = None, block_controlnet_hidden_states: list[torch.Tensor] | None = None, joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, skip_layers: list[int] | None = None, **kwargs, ) -> tuple[torch.Tensor] | SD3Transformer2DModelOutput: del kwargs if encoder_hidden_states is None: raise ValueError("encoder_hidden_states must be provided") if pooled_projections is None: raise ValueError("pooled_projections must be provided") if timestep is None: raise ValueError("timestep must be provided") if timestep.dim() == 0: timestep = timestep[None] if timestep.dim() > 1: timestep = timestep.reshape(-1) if timestep.shape[0] == 1 and hidden_states.shape[0] > 1: timestep = timestep.expand(hidden_states.shape[0]) try: get_forward_context() forward_context = nullcontext() except AssertionError: timestep_val = int(timestep[0].item()) if timestep.numel() > 0 else 0 forward_context = set_forward_context( current_timestep=timestep_val, attn_metadata=None, ) with forward_context: height, width = hidden_states.shape[-2:] hidden_states = self.pos_embed(hidden_states) temb = self.time_text_embed(timestep, pooled_projections) encoder_hidden_states, _ = self.context_embedder( encoder_hidden_states) for index_block, block in enumerate(self.transformer_blocks): is_skip = bool( skip_layers is not None and index_block in skip_layers) if is_skip: continue encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, joint_attention_kwargs=joint_attention_kwargs, ) if (block_controlnet_hidden_states is not None and not block.context_pre_only): interval_control = len(self.transformer_blocks) / len( block_controlnet_hidden_states) hidden_states = hidden_states + block_controlnet_hidden_states[ int(index_block / interval_control)] hidden_states = self.norm_out(hidden_states, temb) hidden_states, _ = self.proj_out(hidden_states) patch_size = self.patch_size height = height // patch_size width = width // patch_size hidden_states = hidden_states.reshape( hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels, ) hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) output = hidden_states.reshape( hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size, ) if not return_dict: return (output, ) return SD3Transformer2DModelOutput(sample=output) EntryClass = SD3Transformer2DModel