chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+"""StableDiffusion3 Transformer model implementation.
+
+NOTE: This initial implementation uses diffusers' JointTransformerBlock directly.
+A native SGLang attention implementation is needed for FlashAttention, TP/SP,
+quantization, and LoRA support.
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import JointTransformerBlock
+from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
+from diffusers.models.normalization import AdaLayerNormContinuous
+
+from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import (
+    StableDiffusion3TransformerConfig,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SD3Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["JointTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+    layer_names = ["transformer_blocks"]
+
+    def __init__(
+        self,
+        config: StableDiffusion3TransformerConfig,
+        hf_config: dict[str, Any] | None = None,
+        quant_config=None,
+    ):
+        super().__init__(config=config, hf_config=hf_config)
+        self.config = config
+        arch_config = config.arch_config
+        sample_size = arch_config.sample_size
+        patch_size = arch_config.patch_size
+        in_channels = arch_config.in_channels
+        num_layers = arch_config.num_layers
+        attention_head_dim = arch_config.attention_head_dim
+        num_attention_heads = arch_config.num_attention_heads
+        joint_attention_dim = arch_config.joint_attention_dim
+        caption_projection_dim = arch_config.caption_projection_dim
+        pooled_projection_dim = arch_config.pooled_projection_dim
+        out_channels = arch_config.out_channels
+        pos_embed_max_size = arch_config.pos_embed_max_size
+        dual_attention_layers = arch_config.dual_attention_layers
+        qk_norm = arch_config.qk_norm
+
+        self.out_channels = out_channels if out_channels is not None else in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.patch_size = patch_size
+
+        self.pos_embed = PatchEmbed(
+            height=sample_size,
+            width=sample_size,
+            patch_size=patch_size,
+            in_channels=in_channels,
+            embed_dim=self.inner_dim,
+            pos_embed_max_size=pos_embed_max_size,
+        )
+        self.time_text_embed = CombinedTimestepTextProjEmbeddings(
+            embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
+        )
+        self.context_embedder = nn.Linear(joint_attention_dim, caption_projection_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                JointTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    context_pre_only=i == num_layers - 1,
+                    qk_norm=qk_norm,
+                    use_dual_attention=i in dual_attention_layers,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
+        )
+        self.proj_out = nn.Linear(
+            self.inner_dim, patch_size * patch_size * self.out_channels, bias=True
+        )
+
+        self.gradient_checkpointing = False
+
+    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 | None = None,
+        guidance: torch.Tensor | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        skip_layers: list[int] | None = None,
+    ) -> torch.Tensor:
+        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.")
+
+        encoder_embeddings = encoder_hidden_states
+
+        height, width = hidden_states.shape[-2:]
+
+        hidden_states = self.pos_embed(hidden_states)
+        temb = self.time_text_embed(timestep, pooled_projections)
+        encoder_embeddings = self.context_embedder(encoder_embeddings)
+
+        skip_layer_set = set(skip_layers) if skip_layers else set()
+
+        if block_controlnet_hidden_states is not None:
+            interval_control = len(self.transformer_blocks) / len(
+                block_controlnet_hidden_states
+            )
+        else:
+            interval_control = 0
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if index_block not in skip_layer_set:
+                encoder_embeddings, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_embeddings,
+                    temb=temb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if (
+                block_controlnet_hidden_states is not None
+                and block.context_pre_only is False
+            ):
+                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)
+
+        # unpatchify
+        patch_size = self.patch_size
+        height = height // patch_size
+        width = width // patch_size
+
+        hidden_states = hidden_states.reshape(
+            shape=(
+                hidden_states.shape[0],
+                height,
+                width,
+                patch_size,
+                patch_size,
+                self.out_channels,
+            )
+        )
+        hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4)
+        output = hidden_states.reshape(
+            shape=(
+                hidden_states.shape[0],
+                self.out_channels,
+                height * patch_size,
+                width * patch_size,
+            )
+        )
+
+        return output
+
+
+# Entry class for registry
+EntryClass = SD3Transformer2DModel