From 4df60434d7e15284f56f72b881db7c218eac7563 Mon Sep 17 00:00:00 2001 From: Xiaole Guo <119391905+gxlvera@users.noreply.github.com> Date: Mon, 13 Apr 2026 01:07:06 -0700 Subject: [PATCH] [diffusion] model: support stable-diffusion-3-medium-diffusers (#19225) Co-authored-by: zhaochenyang20 Co-authored-by: Kangrui Du Co-authored-by: Xiaole Guo --- .../configs/models/dits/__init__.py | 4 + .../configs/models/dits/stablediffusion3.py | 37 ++++ .../configs/models/vaes/__init__.py | 4 + .../configs/models/vaes/stablediffusion3.py | 72 +++++++ .../configs/pipeline_configs/__init__.py | 4 + .../configs/pipeline_configs/base.py | 27 +++ .../configs/pipeline_configs/flux.py | 15 ++ .../pipeline_configs/stablediffusion3.py | 202 ++++++++++++++++++ .../configs/sample/stablediffusion3.py | 20 ++ python/sglang/multimodal_gen/registry.py | 26 +++ .../component_loaders/text_encoder_loader.py | 47 ++-- .../loader/component_loaders/vae_loader.py | 7 + .../multimodal_gen/runtime/loader/utils.py | 4 +- .../runtime/models/dits/stablediffusion3.py | 180 ++++++++++++++++ .../runtime/models/encoders/clip.py | 44 +++- .../runtime/pipelines/stable_diffusion_3.py | 111 ++++++++++ .../pipelines_core/stages/text_encoding.py | 25 ++- 17 files changed, 800 insertions(+), 29 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/configs/models/vaes/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/configs/sample/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index ba114c18f..63ac893dd 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -5,6 +5,9 @@ from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConf from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig +from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import ( + StableDiffusion3TransformerConfig, +) from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig __all__ = [ @@ -14,4 +17,5 @@ __all__ = [ "Hunyuan3DDiTConfig", "MOVAAudioConfig", "MOVAVideoConfig", + "StableDiffusion3TransformerConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py b/python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py new file mode 100644 index 000000000..cbe3141b7 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 Transformer model configuration""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class StableDiffusion3TransformerArchConfig(DiTArchConfig): + """Architecture configuration for StableDiffusion3 Transformer, applicable to SD3-medium, SD3.5-medium, SD3.5-large.""" + + sample_size: int = 128 + patch_size: int = 2 + in_channels: int = 16 + out_channels: int = 16 + num_layers: int = 18 + attention_head_dim: int = 64 + num_attention_heads: int = 18 + cross_attention_dim: int = 4096 + joint_attention_dim: int = 4096 + caption_projection_dim: int = 1152 + pooled_projection_dim: int = 2048 + pos_embed_max_size: int = 96 + dual_attention_layers: tuple[int, ...] = () + qk_norm: str | None = None + + _class_name: str = "SD3Transformer2DModel" + + +@dataclass +class StableDiffusion3TransformerConfig(DiTConfig): + """Configuration for StableDiffusion3 Transformer model.""" + + arch_config: StableDiffusion3TransformerArchConfig = field( + default_factory=StableDiffusion3TransformerArchConfig + ) diff --git a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py index 8b9dd60fe..3438b1b89 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py @@ -3,11 +3,15 @@ from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig +from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import ( + StableDiffusion3VAEConfig, +) from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig __all__ = [ "DacVAEConfig", "HunyuanVAEConfig", + "StableDiffusion3VAEConfig", "WanVAEConfig", "Hunyuan3DVAEConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/vaes/stablediffusion3.py b/python/sglang/multimodal_gen/configs/models/vaes/stablediffusion3.py new file mode 100644 index 000000000..fe6351b2e --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/vaes/stablediffusion3.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 VAE configuration.""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class StableDiffusion3VAEArchConfig(VAEArchConfig): + """Architecture configuration for StableDiffusion3 VAE.""" + + scaling_factor: float = 1.5305 + shift_factor: float = 0.0609 + + spatial_compression_ratio: int = 8 + temporal_compression_ratio: int = 1 + + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 16 + sample_size: int = 128 + + block_out_channels: tuple[int, ...] = (128, 256, 512, 512) + layers_per_block: int = 2 + act_fn: str = "silu" + norm_num_groups: int = 32 + + down_block_types: tuple[str, ...] = ( + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + ) + up_block_types: tuple[str, ...] = ( + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + ) + + attention_head_dim: int = 8 + mid_block_add_attention: bool = True + use_quant_conv: bool = False + use_post_quant_conv: bool = False + + +@dataclass +class StableDiffusion3VAEConfig(VAEConfig): + """Configuration for StableDiffusion3 VAE.""" + + arch_config: StableDiffusion3VAEArchConfig = field( + default_factory=StableDiffusion3VAEArchConfig + ) + + tile_sample_min_height: int = 512 + tile_sample_min_width: int = 512 + tile_sample_min_num_frames: int = 1 + tile_sample_stride_height: int = 448 + tile_sample_stride_width: int = 448 + tile_sample_stride_num_frames: int = 1 + + use_tiling: bool = True + use_temporal_tiling: bool = False + use_parallel_tiling: bool = True + use_temporal_scaling_frames: bool = False + + def __post_init__(self) -> None: + """Post initialization for SD3 VAE specific setup.""" + super().__post_init__() + self.update_model_arch({"_class_name": "AutoencoderKL"}) + self.blend_num_frames = 0 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index e2e7c36dc..005bc862c 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -30,6 +30,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import ( + StableDiffusion3PipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.wan import ( SelfForcingWanT2V480PConfig, WanI2V480PConfig, @@ -55,6 +58,7 @@ __all__ = [ "SanaPipelineConfig", "SlidingTileAttnConfig", "MOVAPipelineConfig", + "StableDiffusion3PipelineConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 6406af765..46086cba7 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -459,6 +459,33 @@ class PipelineConfig: sharded_tensor = sharded_tensor[:, :, rank_in_sp_group, :, :, :] return sharded_tensor, True + def get_text_encoder_attention_mask( + self, text_inputs: dict, encoder_index: int + ) -> "torch.Tensor | None": + """Return the attention mask for the given text encoder. + + Override to suppress (return None) or modify the mask per model. + """ + return text_inputs.get("attention_mask") + + def get_text_encoder_pooler_output( + self, outputs: "BaseEncoderOutput", encoder_index: int + ) -> "torch.Tensor | None": + """Return the pooler output for the given text encoder, or None to skip. + + Override for models that need pooled embeddings (e.g. FLUX v1, SD3). + """ + return None + + def select_vae_weight_files( + self, + safetensors_list: list[str], + component_model_path: str, + component_name: str, + vae_precision: str, + ) -> list[str]: + return safetensors_list + def get_pos_prompt_embeds(self, batch): return batch.prompt_embeds diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index 76c7ddea6..38829098f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -81,6 +81,13 @@ class FluxPipelineConfig(ImagePipelineConfig): ] ) + def get_text_encoder_attention_mask(self, text_inputs, encoder_index): + # Flux v1 does not use attention masks for text encoders. + return None + + def get_text_encoder_pooler_output(self, outputs, encoder_index): + return outputs.pooler_output + def prepare_sigmas(self, sigmas, num_inference_steps): return self._prepare_sigmas(sigmas, num_inference_steps) @@ -391,6 +398,14 @@ class Flux2PipelineConfig(FluxPipelineConfig): ] ) + def get_text_encoder_attention_mask(self, text_inputs, encoder_index): + # Flux2 uses standard attention masks (unlike Flux v1). + return text_inputs.get("attention_mask") + + def get_text_encoder_pooler_output(self, outputs, encoder_index): + # Flux2 does not use pooler output. + return None + def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict: messages = build_flux2_text_messages(prompts) inputs = tokenizer.apply_chat_template( diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/stablediffusion3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/stablediffusion3.py new file mode 100644 index 000000000..beea21097 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/stablediffusion3.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Diffusion 3 pipeline configuration.""" + +import os +from dataclasses import dataclass, field +from typing import Callable + +import torch + +from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits import StableDiffusion3TransformerConfig +from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput +from sglang.multimodal_gen.configs.models.encoders.base import TextEncoderArchConfig +from sglang.multimodal_gen.configs.models.encoders.clip import ( + CLIPTextArchConfig, + CLIPTextConfig, +) +from sglang.multimodal_gen.configs.models.encoders.t5 import ( + T5ArchConfig, + T5Config, +) +from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import ( + StableDiffusion3VAEConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + SpatialImagePipelineConfig, +) + + +def sd3_clip_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: + """Extract pre-final hidden state for SD3 CLIP encoders.""" + if outputs.hidden_states is None: + raise ValueError( + "SD3 CLIP postprocessing requires hidden_states from encoder output." + ) + return outputs.hidden_states[-2] + + +def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: + return outputs.last_hidden_state + + +def select_sd3_vae_weight_files( + safetensors_list: list[str], + component_model_path: str, + component_name: str, + vae_precision: str, +) -> list[str]: + """Select SD3 VAE checkpoint file candidates with minimal policy.""" + if component_name not in ("vae", "video_vae"): + return safetensors_list + + base_name = "diffusion_pytorch_model" + if vae_precision == "fp16": + fp16_path = os.path.join(component_model_path, f"{base_name}.fp16.safetensors") + if os.path.exists(fp16_path): + return [fp16_path] + + full_path = os.path.join(component_model_path, f"{base_name}.safetensors") + if os.path.exists(full_path): + return [full_path] + return safetensors_list + + +@dataclass +class SD3CLIPTextArchConfig(CLIPTextArchConfig): + def __post_init__(self) -> None: + super().__post_init__() + self.tokenizer_kwargs.update( + { + "max_length": self.text_len, + "padding": "max_length", + } + ) + + +@dataclass +class SD3CLIPTextConfig(CLIPTextConfig): + arch_config: TextEncoderArchConfig = field(default_factory=SD3CLIPTextArchConfig) + + +@dataclass +class SD3T5ArchConfig(T5ArchConfig): + def __post_init__(self) -> None: + super().__post_init__() + self.tokenizer_kwargs.update({"max_length": 256}) + + +@dataclass +class SD3T5Config(T5Config): + arch_config: TextEncoderArchConfig = field(default_factory=SD3T5ArchConfig) + + +@dataclass +class StableDiffusion3PipelineConfig(SpatialImagePipelineConfig): + """Configuration for SD3 image generation pipeline. + + This config intentionally relies on SD3-specific encoder configs to provide + tokenizer kwargs, instead of stage-level tokenizer overrides. + """ + + task_type: ModelTaskType = ModelTaskType.T2I + + dit_config: DiTConfig = field(default_factory=StableDiffusion3TransformerConfig) + vae_config: VAEConfig = field(default_factory=StableDiffusion3VAEConfig) + + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: ( + SD3CLIPTextConfig(), + SD3CLIPTextConfig(), + SD3T5Config(), + ) + ) + + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp16", "fp16", "fp32") + ) + + preprocess_text_funcs: tuple[Callable[[str], str] | None, ...] = field( + default_factory=lambda: ( + None, + None, + None, + ) + ) + + postprocess_text_funcs: tuple[ + Callable[[BaseEncoderOutput, dict], torch.Tensor], ... + ] = field( + default_factory=lambda: ( + sd3_clip_postprocess_text, + sd3_clip_postprocess_text, + t5_postprocess_text, + ) + ) + + should_use_guidance: bool = False + guidance_scale: float = 7.0 + + def __post_init__(self) -> None: + configs = list(self.text_encoder_configs) + configs[0].update_model_arch({"_class_name": "CLIPTextModelWithProjection"}) + configs[1].update_model_arch({"_class_name": "CLIPTextModelWithProjection"}) + configs[2].update_model_arch({"_class_name": "T5EncoderModel"}) + self.text_encoder_configs = tuple(configs) + + def get_text_encoder_pooler_output(self, outputs, encoder_index): + # SD3 uses pooled embeddings only from the two CLIP encoders (indices 0 and 1). + if encoder_index <= 1: + return outputs.pooler_output + return None + + def select_vae_weight_files( + self, + safetensors_list: list[str], + component_model_path: str, + component_name: str, + vae_precision: str, + ) -> list[str]: + return select_sd3_vae_weight_files( + safetensors_list=safetensors_list, + component_model_path=component_model_path, + component_name=component_name, + vae_precision=vae_precision, + ) + + def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict: + text_inputs = tokenizer(prompt, **tok_kwargs) + text_inputs["attention_mask"] = None + return text_inputs + + def get_pos_prompt_embeds(self, batch): + return batch.prompt_embeds[0] + + def get_neg_prompt_embeds(self, batch): + return batch.negative_prompt_embeds[0] + + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): + return { + "pooled_projections": ( + batch.pooled_embeds[0] if batch.pooled_embeds else None + ) + } + + def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): + return { + "pooled_projections": ( + batch.neg_pooled_embeds[0] if batch.neg_pooled_embeds else None + ) + } + + # SD3 image latents are spatial (B, C, H, W), not video-like (B, C, T, H, W). + def prepare_latent_shape(self, batch, batch_size, num_frames): # noqa: ARG002 + spatial_ratio = self.vae_config.arch_config.spatial_compression_ratio + in_channels = self.dit_config.arch_config.in_channels + return ( + batch_size, + in_channels, + batch.height // spatial_ratio, + batch.width // spatial_ratio, + ) diff --git a/python/sglang/multimodal_gen/configs/sample/stablediffusion3.py b/python/sglang/multimodal_gen/configs/sample/stablediffusion3.py new file mode 100644 index 000000000..f586ee073 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/stablediffusion3.py @@ -0,0 +1,20 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 sampling parameters configuration.""" + +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + + +@dataclass +class StableDiffusion3SamplingParams(SamplingParams): + """Sampling parameters for StableDiffusion3.""" + + # A single space ensures tokenizers produce valid (non-empty) input for CFG. + negative_prompt: str = " " + num_frames: int = 1 + num_inference_steps: int = 50 + guidance_scale: float = 7.0 + guidance_rescale: float = 0.0 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 93bbd9cc2..aaa49b08e 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -68,6 +68,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import ( + StableDiffusion3PipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.wan import ( FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, @@ -109,6 +112,9 @@ from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImageSamplingParams, ) from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams +from sglang.multimodal_gen.configs.sample.stablediffusion3 import ( + StableDiffusion3SamplingParams, +) from sglang.multimodal_gen.configs.sample.wan import ( FastWanT2V480PConfig, Turbo_Wan2_2_I2V_A14B_SamplingParam, @@ -849,6 +855,26 @@ def _register_configs(): hf_model_paths=["Qwen/Qwen-Image-Layered"], model_detectors=[lambda hf_id: "qwen-image-layered" in hf_id.lower()], ) + register_configs( + sampling_param_cls=StableDiffusion3SamplingParams, + pipeline_config_cls=StableDiffusion3PipelineConfig, + hf_model_paths=[ + "stabilityai/stable-diffusion-3-medium", + "stabilityai/stable-diffusion-3-medium-diffusers", + "stabilityai/stable-diffusion-3.5-medium", + "stabilityai/stable-diffusion-3.5-medium-diffusers", + "stabilityai/stable-diffusion-3.5-large", + "stabilityai/stable-diffusion-3.5-large-diffusers", + ], + model_detectors=[ + lambda hf_id: "stable-diffusion-3-medium" in hf_id.lower() + or "stable-diffusion-3.5-medium" in hf_id.lower() + or "stable-diffusion-3.5-large" in hf_id.lower() + or "sd3-medium" in hf_id.lower() + or "sd3.5-medium" in hf_id.lower() + or "sd3.5-large" in hf_id.lower() + ], + ) register_configs( sampling_param_cls=GlmImageSamplingParams, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 01be87142..d7f9b0aec 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -1,6 +1,7 @@ import dataclasses import glob import os +import re from collections.abc import Generator, Iterable from typing import cast @@ -218,21 +219,21 @@ class TextEncoderLoader(ComponentLoader): component_path=component_model_path ) - def is_not_first_encoder(module_name): - return "2" in module_name - # TODO(mick): had to throw an exception for different text-encoder arch - if not is_not_first_encoder(component_name): - encoder_config = server_args.pipeline_config.text_encoder_configs[0] - encoder_config.update_model_arch(model_config) + encoder_index = self._extract_encoder_index(component_name) + assert encoder_index < len( + server_args.pipeline_config.text_encoder_configs + ) and encoder_index < len(server_args.pipeline_config.text_encoder_precisions) + + encoder_config = server_args.pipeline_config.text_encoder_configs[encoder_index] + encoder_config.update_model_arch(model_config) + + if encoder_index == 0: for key, value in diffusers_pretrained_config.__dict__.items(): setattr(encoder_config.arch_config, key, value) - encoder_dtype = server_args.pipeline_config.text_encoder_precisions[0] - else: - assert len(server_args.pipeline_config.text_encoder_configs) == 2 - encoder_config = server_args.pipeline_config.text_encoder_configs[1] - encoder_config.update_model_arch(model_config) - encoder_dtype = server_args.pipeline_config.text_encoder_precisions[1] + encoder_dtype = server_args.pipeline_config.text_encoder_precisions[ + encoder_index + ] # TODO(will): add support for other dtypes return self.load_model( component_model_path, @@ -242,6 +243,28 @@ class TextEncoderLoader(ComponentLoader): cpu_offload_flag=cpu_offload_flag, ) + @staticmethod + def _extract_encoder_index(component_name: str) -> int: + """ + Map text encoder component names to zero-based indices. + + Examples: + - text_encoder -> 0 + - text_encoder_2 -> 1 + - text_encoder_3 -> 2 + """ + match = re.search(r"_(\d+)$", component_name) + if match is None: + return 0 + + suffix_num = int(match.group(1)) + if suffix_num <= 0: + raise ValueError( + f"Invalid text encoder component name '{component_name}': " + "numeric suffix must be >= 1." + ) + return suffix_num - 1 + def load_model( self, model_path: str, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 52c998ac8..c49aa4d67 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -129,6 +129,13 @@ class VAELoader(ComponentLoader): vae = vae_cls(vae_config).to(target_device) safetensors_list = _list_safetensors_files(component_model_path) + safetensors_list = server_args.pipeline_config.select_vae_weight_files( + safetensors_list=safetensors_list, + component_model_path=component_model_path, + component_name=component_name, + vae_precision=vae_precision, + ) + assert ( len(safetensors_list) >= 1 ), f"Found no safetensors files in {component_model_path}" diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index d8ed09c2e..edbcf65fb 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -181,9 +181,7 @@ class skip_init_modules: def _normalize_component_type(module_type: str) -> str: """Normalize module types like 'text_encoder_2' -> 'text_encoder'.""" - if module_type.endswith("_2"): - return module_type[:-2] - return module_type + return re.sub(r"_\d+$", "", module_type) def _clean_hf_config_inplace(model_config: dict) -> None: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py b/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py new file mode 100644 index 000000000..1ceeec2d9 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py @@ -0,0 +1,180 @@ +# 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.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class SD3Transformer2DModel(CachableDiT): + _supports_gradient_checkpointing = True + _no_split_modules = ["JointTransformerBlock"] + _skip_layerwise_casting_patterns = ["pos_embed", "norm"] + + 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 diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py index 56b4932e5..cda4f90ce 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py @@ -594,6 +594,48 @@ class CLIPTextModel(TextEncoder): return loaded_params +class CLIPTextModelWithProjection(CLIPTextModel): + """ + CLIP text encoder with projection head for pooled_output. + """ + + def __init__( + self, + config: CLIPTextConfig, + ) -> None: + super().__init__(config) + self.text_projection = nn.Linear( + config.hidden_size, config.projection_dim, bias=False + ) + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + outputs: BaseEncoderOutput = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_hidden_states=output_hidden_states, + ) + + pooled_output = outputs.pooler_output + if pooled_output is not None: + pooled_output = self.text_projection(pooled_output) + + return BaseEncoderOutput( + last_hidden_state=outputs.last_hidden_state, + pooler_output=pooled_output, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + class CLIPVisionTransformer(nn.Module): def __init__( @@ -759,4 +801,4 @@ class BertModel(CLIPTextModel): pass -EntryClass = [CLIPTextModel, CLIPVisionModel] +EntryClass = [CLIPTextModel, CLIPTextModelWithProjection, CLIPVisionModel] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py b/python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py new file mode 100644 index 000000000..f0b8c0b9f --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py @@ -0,0 +1,111 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 pipeline implementation.""" + +import torch + +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + InputValidationStage, + PipelineStage, + TextEncodingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class SD3ConditioningStage(PipelineStage): + """Merge CLIP-T, CLIP-G and T5 embeddings into unified prompt/pooled tensors.""" + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + batch.prompt_embeds, batch.pooled_embeds = self._merge( + batch.prompt_embeds, batch.pooled_embeds + ) + if batch.do_classifier_free_guidance: + batch.negative_prompt_embeds, batch.neg_pooled_embeds = self._merge( + batch.negative_prompt_embeds, batch.neg_pooled_embeds + ) + return batch + + @staticmethod + def _merge( + embeds_list: list[torch.Tensor], + pooled_list: list[torch.Tensor], + ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + """Merge 3 encoder outputs into unified prompt/pooled tensors. + + SD3-medium uses exactly 3 text encoders (CLIP-L, CLIP-G, T5). + Returns single-element lists to match the batch field format expected + by downstream stages (get_pos_prompt_embeds accesses index [0]). + """ + if len(embeds_list) != 3: + raise ValueError( + f"SD3 requires exactly 3 prompt embedding tensors, got {len(embeds_list)}." + ) + if len(pooled_list) < 2: + raise ValueError( + f"SD3 requires at least 2 pooled embedding tensors, got {len(pooled_list)}." + ) + + clipt, clipg, t5 = embeds_list + clip_merged = torch.cat([clipt, clipg], dim=-1) + clip_merged = torch.nn.functional.pad( + clip_merged, (0, t5.shape[-1] - clip_merged.shape[-1]) + ) + merged_embeds = [torch.cat([clip_merged, t5], dim=-2)] + merged_pooled = [torch.cat([pooled_list[0], pooled_list[1]], dim=-1)] + return merged_embeds, merged_pooled + + +class StableDiffusion3Pipeline(ComposedPipelineBase): + """StableDiffusion3 pipeline implementation.""" + + pipeline_name = "StableDiffusion3Pipeline" + + _required_config_modules = [ + "text_encoder", + "text_encoder_2", + "text_encoder_3", + "tokenizer", + "tokenizer_2", + "tokenizer_3", + "vae", + "transformer", + "scheduler", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + self.add_stage(InputValidationStage()) + + self.add_stage( + TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + self.get_module("text_encoder_2"), + self.get_module("text_encoder_3"), + ], + tokenizers=[ + self.get_module("tokenizer"), + self.get_module("tokenizer_2"), + self.get_module("tokenizer_3"), + ], + ), + "prompt_encoding_stage_primary", + ) + + self.add_stage(SD3ConditioningStage()) + + self.add_standard_timestep_preparation_stage() + self.add_standard_latent_preparation_stage() + self.add_standard_denoising_stage() + self.add_standard_decoding_stage() + + +EntryClass = StableDiffusion3Pipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index b5d51ab6e..efb409e50 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -12,8 +12,6 @@ import inspect import torch from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput -from sglang.multimodal_gen.configs.pipeline_configs import FluxPipelineConfig -from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req @@ -263,11 +261,11 @@ class TextEncodingStage(PipelineStage): ).to(target_device) input_ids = text_inputs["input_ids"] - is_flux_v1 = isinstance( - server_args.pipeline_config, FluxPipelineConfig - ) and not isinstance(server_args.pipeline_config, Flux2PipelineConfig) - - attention_mask = None if is_flux_v1 else text_inputs["attention_mask"] + attention_mask = ( + server_args.pipeline_config.get_text_encoder_attention_mask( + text_inputs, i + ) + ) encoder_forward_kwargs = { "input_ids": input_ids, "output_hidden_states": True, @@ -292,12 +290,13 @@ class TextEncodingStage(PipelineStage): prompt_embeds = prompt_embeds.to(device=target_device) embeds_list.append(prompt_embeds) - if is_flux_v1 and outputs.pooler_output is not None: - # FLUX.1 only consumes the pooled CLIP projection. The T5 - # encoder in the same pipeline has no pooler output. - pooled_embeds_list.append( - outputs.pooler_output.to(device=target_device) - ) + + pooled_output = server_args.pipeline_config.get_text_encoder_pooler_output( + outputs, i + ) + if pooled_output is not None: + pooled_embeds_list.append(pooled_output.to(device=target_device)) + if return_attention_mask: mask_to_store = ( attention_mask.to(device=target_device)