sd3

@dataclass
class SD3Transformer2DModelOutput:
    sample: torch.Tensor

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)</code></pre></details>
fastvideo.models.dits.sd3.SD3PatchEmbed.cropped_pos_embed
cropped_pos_embed(height: int, width: int) -> torch.Tensor
Source
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 &gt; self.pos_embed_max_size:
    raise ValueError(f&quot;Height (&#123;height&#125;) cannot be greater than pos_embed_max_size &quot;
                     f&quot;(&#123;self.pos_embed_max_size&#125;)&quot;)
if width &gt; self.pos_embed_max_size:
    raise ValueError(f&quot;Width (&#123;width&#125;) cannot be greater than pos_embed_max_size &quot;
                     f&quot;(&#123;self.pos_embed_max_size&#125;)&quot;)

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&#91;:, top:top + height, left:left + width, :&#93;
return spatial.reshape(1, -1, spatial.shape&#91;-1&#93;)</code></pre></details>
fastvideo.models.dits.sd3.SD3PatchEmbed.forward
forward(latent: torch.Tensor) -> torch.Tensor
Source
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&#91;-1&#93;,
            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)</code></pre></details>
fastvideo.models.dits.sd3.SD3TimestepEmbedding
SD3TimestepEmbedding(in_channels: int, time_embed_dim: int, act_fn: str = 'silu')
Source
class SD3TimestepEmbedding(nn.Module):
def __init__(
    self,
    in_channels: int,
    time_embed_dim: int,
    act_fn: str = &quot;silu&quot;,
) -&gt; 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) -&gt; torch.Tensor:
    sample, _ = self.linear_1(sample)
    sample = self.act(sample)
    sample, _ = self.linear_2(sample)
    return sample</code></pre></details>
fastvideo.models.dits.sd3.SD3TimestepEmbedding.forward
forward(sample: torch.Tensor) -> torch.Tensor
Source
def forward(self, sample: torch.Tensor) -> torch.Tensor:
    sample, _ = self.linear_1(sample)
    sample = self.act(sample)
    sample, _ = self.linear_2(sample)
    return sample
fastvideo.models.dits.sd3.SD3TextProjection
SD3TextProjection(in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = 'silu')
Source
class SD3TextProjection(nn.Module):
def __init__(
    self,
    in_features: int,
    hidden_size: int,
    out_features: int | None = None,
    act_fn: str = &quot;silu&quot;,
) -&gt; 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) -&gt; 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</code></pre></details>
fastvideo.models.dits.sd3.SD3TextProjection.forward
forward(caption: torch.Tensor) -> torch.Tensor
Source
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
fastvideo.models.dits.sd3.CombinedTimestepTextProjEmbeddings
CombinedTimestepTextProjEmbeddings(embedding_dim: int, pooled_projection_dim: int)
Source
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=&quot;silu&quot;,
    )
    self.text_embedder = SD3TextProjection(
        pooled_projection_dim,
        embedding_dim,
        act_fn=&quot;silu&quot;,
    )

def forward(
    self,
    timestep: torch.Tensor,
    pooled_projection: torch.Tensor,
) -&gt; 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</code></pre></details>
fastvideo.models.dits.sd3.CombinedTimestepTextProjEmbeddings.forward
forward(timestep: torch.Tensor, pooled_projection: torch.Tensor) -> torch.Tensor
Source
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
fastvideo.models.dits.sd3.SD35AdaLayerNormZeroX
SD35AdaLayerNormZeroX(embedding_dim: int, norm_type: str = 'layer_norm', bias: bool = True)
Source
class SD35AdaLayerNormZeroX(nn.Module):
def __init__(
    self,
    embedding_dim: int,
    norm_type: str = &quot;layer_norm&quot;,
    bias: bool = True,
) -&gt; None:
    super().__init__()

    self.silu = nn.SiLU()
    self.linear = ReplicatedLinear(embedding_dim, 9 * embedding_dim, bias=bias)

    if norm_type != &quot;layer_norm&quot;:
        raise ValueError(f&quot;Unsupported norm_type (&#123;norm_type&#125;); expected layer_norm&quot;)
    self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)

def forward(
    self,
    hidden_states: torch.Tensor,
    emb: torch.Tensor,
) -&gt; tuple&#91;torch.Tensor, ...&#93;:
    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&#91;:, None&#93;) + shift_msa&#91;:, None&#93;
    norm_hidden_states2 = norm_hidden_states * (1 + scale_msa2&#91;:, None&#93;) + shift_msa2&#91;:, None&#93;
    return (
        hidden_states,
        gate_msa,
        shift_mlp,
        scale_mlp,
        gate_mlp,
        norm_hidden_states2,
        gate_msa2,
    )</code></pre></details>
fastvideo.models.dits.sd3.SD35AdaLayerNormZeroX.forward
forward(hidden_states: torch.Tensor, emb: torch.Tensor) -> tuple[torch.Tensor, ...]
Source
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&#91;:, None&#93;) + shift_msa&#91;:, None&#93;
norm_hidden_states2 = norm_hidden_states * (1 + scale_msa2&#91;:, None&#93;) + shift_msa2&#91;:, None&#93;
return (
    hidden_states,
    gate_msa,
    shift_mlp,
    scale_mlp,
    gate_mlp,
    norm_hidden_states2,
    gate_msa2,
)</code></pre></details>
fastvideo.models.dits.sd3.SD3AdaLayerNormZero
SD3AdaLayerNormZero(embedding_dim: int, norm_type: str = 'layer_norm', bias: bool = True)
Source
class SD3AdaLayerNormZero(nn.Module):
def __init__(
    self,
    embedding_dim: int,
    norm_type: str = &quot;layer_norm&quot;,
    bias: bool = True,
) -&gt; None:
    super().__init__()

    self.silu = nn.SiLU()
    self.linear = ReplicatedLinear(embedding_dim, 6 * embedding_dim, bias=bias)

    if norm_type != &quot;layer_norm&quot;:
        raise ValueError(f&quot;Unsupported norm_type (&#123;norm_type&#125;); expected layer_norm&quot;)
    self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)

def forward(
    self,
    x: torch.Tensor,
    emb: torch.Tensor,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor&#93;:
    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&#91;:, None&#93;) + shift_msa&#91;:, None&#93;
    return x, gate_msa, shift_mlp, scale_mlp, gate_mlp</code></pre></details>
fastvideo.models.dits.sd3.SD3AdaLayerNormZero.forward
forward(x: torch.Tensor, emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
Source
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
fastvideo.models.dits.sd3.SD3AdaLayerNormContinuous
SD3AdaLayerNormContinuous(embedding_dim: int, conditioning_embedding_dim: int, elementwise_affine: bool = True, eps: float = 1e-05, bias: bool = True, norm_type: str = 'layer_norm')
Source
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 = &quot;layer_norm&quot;,
) -&gt; None:
    super().__init__()

    self.silu = nn.SiLU()
    self.linear = ReplicatedLinear(
        conditioning_embedding_dim,
        embedding_dim * 2,
        bias=bias,
    )

    if norm_type != &quot;layer_norm&quot;:
        raise ValueError(f&quot;Unsupported norm_type: &#123;norm_type&#125;&quot;)

    self.norm = nn.LayerNorm(
        embedding_dim,
        eps=eps,
        elementwise_affine=elementwise_affine,
    )

def forward(
    self,
    x: torch.Tensor,
    conditioning_embedding: torch.Tensor,
) -&gt; 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)&#91;:, None, :&#93; + shift&#91;:, None, :&#93;
    return x</code></pre></details>
fastvideo.models.dits.sd3.SD3AdaLayerNormContinuous.forward
forward(x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor
Source
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
fastvideo.models.dits.sd3.SD3GELU
SD3GELU(dim_in: int, dim_out: int, approximate: str = 'none', bias: bool = True)
Source
class SD3GELU(nn.Module):
def __init__(
    self,
    dim_in: int,
    dim_out: int,
    approximate: str = &quot;none&quot;,
    bias: bool = True,
) -&gt; None:
    super().__init__()
    self.proj = ReplicatedLinear(dim_in, dim_out, bias=bias)
    self.approximate = approximate

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    hidden_states, _ = self.proj(hidden_states)
    return F.gelu(hidden_states, approximate=self.approximate)</code></pre></details>
fastvideo.models.dits.sd3.SD3GELU.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    hidden_states, _ = self.proj(hidden_states)
    return F.gelu(hidden_states, approximate=self.approximate)
fastvideo.models.dits.sd3.SD3FeedForward
SD3FeedForward(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)
Source
class SD3FeedForward(nn.Module):
def __init__(
    self,
    dim: int,
    dim_out: int | None = None,
    mult: int = 4,
    dropout: float = 0.0,
    activation_fn: str = &quot;geglu&quot;,
    final_dropout: bool = False,
    inner_dim: int | None = None,
    bias: bool = True,
) -&gt; None:
    super().__init__()

    if inner_dim is None:
        inner_dim = int(dim * mult)
    if dim_out is None:
        dim_out = dim

    if activation_fn == &quot;gelu-approximate&quot;:
        act_fn: nn.Module = SD3GELU(
            dim,
            inner_dim,
            approximate=&quot;tanh&quot;,
            bias=bias,
        )
    elif activation_fn == &quot;gelu&quot;:
        act_fn = SD3GELU(dim, inner_dim, approximate=&quot;none&quot;, bias=bias)
    else:
        raise ValueError(f&quot;Unsupported activation_fn: &#123;activation_fn&#125;&quot;)

    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) -&gt; 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</code></pre></details>
fastvideo.models.dits.sd3.SD3FeedForward.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
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
fastvideo.models.dits.sd3.SD3Attention
SD3Attention(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-06, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None)
Joint self-attention used in SD3 blocks.
Source
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&#91;AttentionBackendEnum, ...&#93; | None = None,
) -&gt; 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(&#91;
        ReplicatedLinear(self.inner_dim, out_dim, bias=True),
        nn.Dropout(0.0),
    &#93;)

    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,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor | None&#93; | torch.Tensor:
    batch_size = hidden_states.shape&#91;0&#93;
    residual_seq_len = hidden_states.shape&#91;1&#93;

    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(&quot;encoder_hidden_states provided but attention does not have&quot;
                               &quot; added projections&quot;)

        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(&#91;query, encoder_query&#93;, dim=1)
        key = torch.cat(&#91;key, encoder_key&#93;, dim=1)
        value = torch.cat(&#91;value, encoder_value&#93;, 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&#91;:, :residual_seq_len&#93;
        context_out = attn_output&#91;:, residual_seq_len:&#93;
        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&#91;0&#93;(hidden_out)
    hidden_out = self.to_out&#91;1&#93;(hidden_out)

    if context_out is not None:
        return hidden_out, context_out
    return hidden_out</code></pre></details>
fastvideo.models.dits.sd3.SD3Attention.forward
forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor
Source
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(&quot;encoder_hidden_states provided but attention does not have&quot;
                           &quot; added projections&quot;)

    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(&#91;query, encoder_query&#93;, dim=1)
    key = torch.cat(&#91;key, encoder_key&#93;, dim=1)
    value = torch.cat(&#91;value, encoder_value&#93;, 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&#91;:, :residual_seq_len&#93;
    context_out = attn_output&#91;:, residual_seq_len:&#93;
    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&#91;0&#93;(hidden_out)
hidden_out = self.to_out&#91;1&#93;(hidden_out)

if context_out is not None:
    return hidden_out, context_out
return hidden_out</code></pre></details>
fastvideo.models.dits.sd3.SD3JointTransformerBlock
SD3JointTransformerBlock(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)
Source
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&#91;AttentionBackendEnum, ...&#93; | None = None,
) -&gt; 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=&quot;layer_norm&quot;,
        )
    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=&quot;gelu-approximate&quot;,
    )

    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=&quot;gelu-approximate&quot;,
        )
    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&#91;str, Any&#93; | None = None,
) -&gt; tuple&#91;torch.Tensor | None, torch.Tensor&#93;:
    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&#91;:, None&#93;) + shift_mlp&#91;:, None&#93;

    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&#91;:, None&#93;) + c_shift_mlp&#91;:, None&#93;

        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</code></pre></details>
fastvideo.models.dits.sd3.SD3JointTransformerBlock.set_chunk_feed_forward
set_chunk_feed_forward(chunk_size: int | None, dim: int = 0)
Source
def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0):
    self._chunk_size = chunk_size
    self._chunk_dim = dim
fastvideo.models.dits.sd3.SD3JointTransformerBlock.forward
forward(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]
Source
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&#91;:, None&#93;) + shift_mlp&#91;:, None&#93;

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&#91;:, None&#93;) + c_shift_mlp&#91;:, None&#93;

    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</code></pre></details>
fastvideo.models.dits.sd3.SD3Transformer2DModel
SD3Transformer2DModel(config: DiTConfig, hf_config: dict[str, Any], **kwargs)
FastVideo-native SD3 Transformer2DModel.
Source
class SD3Transformer2DModel(BaseDiT):
    """FastVideo-native SD3 Transformer2DModel."""
_fsdp_shard_conditions = &#91;
    lambda n, m: &quot;transformer_blocks&quot; in n and n.split(&quot;.&quot;)&#91;-1&#93;.isdigit(),
&#93;
_compile_conditions = _fsdp_shard_conditions
param_names_mapping: dict&#91;str, Any&#93; = &#123;&#125;
reverse_param_names_mapping: dict&#91;str, Any&#93; = &#123;&#125;
lora_param_names_mapping: dict&#91;str, Any&#93; = &#123;&#125;
_supported_attention_backends = (
    AttentionBackendEnum.FLASH_ATTN,
    AttentionBackendEnum.TORCH_SDPA,
)

def __init__(self, config: DiTConfig, hf_config: dict&#91;str, Any&#93;, **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, &quot;dual_attention_layers&quot;, ())
    dual_layers = tuple(dual_layers) if isinstance(dual_layers, list) else dual_layers

    self.transformer_blocks = nn.ModuleList(&#91;
        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)
    &#93;)

    self.norm_out = SD3AdaLayerNormContinuous(
        self.inner_dim,
        self.inner_dim,
        elementwise_affine=False,
        eps=1e-6,
        bias=True,
        norm_type=&quot;layer_norm&quot;,
    )
    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,
) -&gt; None:
    if dim not in &#91;0, 1&#93;:
        raise ValueError(f&quot;dim must be 0 or 1, got &#123;dim&#125;&quot;)

    chunk_size = chunk_size or 1

    def fn_recursive_feed_forward(module: torch.nn.Module, chunk: int, chunk_dim: int):
        if hasattr(module, &quot;set_chunk_feed_forward&quot;):
            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, &quot;set_chunk_feed_forward&quot;):
            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&#91;torch.Tensor&#93; | None = None,
    joint_attention_kwargs: dict&#91;str, Any&#93; | None = None,
    return_dict: bool = True,
    skip_layers: list&#91;int&#93; | None = None,
    **kwargs,
) -&gt; tuple&#91;torch.Tensor&#93; | SD3Transformer2DModelOutput:
    del kwargs

    if encoder_hidden_states is None:
        raise ValueError(&quot;encoder_hidden_states must be provided&quot;)
    if pooled_projections is None:
        raise ValueError(&quot;pooled_projections must be provided&quot;)
    if timestep is None:
        raise ValueError(&quot;timestep must be provided&quot;)

    if timestep.dim() == 0:
        timestep = timestep&#91;None&#93;
    if timestep.dim() &gt; 1:
        timestep = timestep.reshape(-1)
    if timestep.shape&#91;0&#93; == 1 and hidden_states.shape&#91;0&#93; &gt; 1:
        timestep = timestep.expand(hidden_states.shape&#91;0&#93;)

    try:
        get_forward_context()
        forward_context = nullcontext()
    except AssertionError:
        timestep_val = int(timestep&#91;0&#93;.item()) if timestep.numel() &gt; 0 else 0
        forward_context = set_forward_context(
            current_timestep=timestep_val,
            attn_metadata=None,
        )

    with forward_context:
        height, width = hidden_states.shape&#91;-2:&#93;

        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&#91;int(index_block / interval_control)&#93;

        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&#91;0&#93;,
            height,
            width,
            patch_size,
            patch_size,
            self.out_channels,
        )
        hidden_states = torch.einsum(&quot;nhwpqc-&gt;nchpwq&quot;, hidden_states)
        output = hidden_states.reshape(
            hidden_states.shape&#91;0&#93;,
            self.out_channels,
            height * patch_size,
            width * patch_size,
        )

        if not return_dict:
            return (output, )

        return SD3Transformer2DModelOutput(sample=output)</code></pre></details>
fastvideo.models.dits.sd3.SD3Transformer2DModel.enable_forward_chunking
enable_forward_chunking(chunk_size: int | None = None, dim: int = 0) -> None
Source
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, &quot;set_chunk_feed_forward&quot;):
        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)</code></pre></details>
fastvideo.models.dits.sd3.SD3Transformer2DModel.disable_forward_chunking
disable_forward_chunking()
Source
def disable_forward_chunking(self):
def fn_recursive_feed_forward(module: torch.nn.Module):
    if hasattr(module, &quot;set_chunk_feed_forward&quot;):
        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)</code></pre></details>
fastvideo.models.dits.sd3.SD3Transformer2DModel.forward
forward(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
Source
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(&quot;encoder_hidden_states must be provided&quot;)
if pooled_projections is None:
    raise ValueError(&quot;pooled_projections must be provided&quot;)
if timestep is None:
    raise ValueError(&quot;timestep must be provided&quot;)

if timestep.dim() == 0:
    timestep = timestep&#91;None&#93;
if timestep.dim() &gt; 1:
    timestep = timestep.reshape(-1)
if timestep.shape&#91;0&#93; == 1 and hidden_states.shape&#91;0&#93; &gt; 1:
    timestep = timestep.expand(hidden_states.shape&#91;0&#93;)

try:
    get_forward_context()
    forward_context = nullcontext()
except AssertionError:
    timestep_val = int(timestep&#91;0&#93;.item()) if timestep.numel() &gt; 0 else 0
    forward_context = set_forward_context(
        current_timestep=timestep_val,
        attn_metadata=None,
    )

with forward_context:
    height, width = hidden_states.shape&#91;-2:&#93;

    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&#91;int(index_block / interval_control)&#93;

    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&#91;0&#93;,
        height,
        width,
        patch_size,
        patch_size,
        self.out_channels,
    )
    hidden_states = torch.einsum(&quot;nhwpqc-&gt;nchpwq&quot;, hidden_states)
    output = hidden_states.reshape(
        hidden_states.shape&#91;0&#93;,
        self.out_channels,
        height * patch_size,
        width * patch_size,
    )

    if not return_dict:
        return (output, )

    return SD3Transformer2DModelOutput(sample=output)</code></pre></details>