diff --git a/python/sglang/multimodal_gen/configs/models/dits/sana.py b/python/sglang/multimodal_gen/configs/models/dits/sana.py new file mode 100644 index 000000000..fa2a3f086 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/sana.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Architecture and model configuration for SANA DiT (Diffusion Transformer). +# +# SANA uses a linear-attention-based transformer that replaces standard +# quadratic self-attention with ReLU-based linear attention, enabling +# efficient high-resolution image synthesis. Cross-attention (standard SDPA) +# is used for text conditioning via Gemma2 embeddings. +# +# Defaults below correspond to the SANA-1.6B / 1024px variant. +# For 4.8B, override num_layers=36, num_attention_heads=64, etc. +# +# Reference: https://arxiv.org/abs/2410.10629 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class SanaArchConfig(DiTArchConfig): + patch_size: int = 1 + in_channels: int = 32 + out_channels: int = 32 + num_layers: int = 20 + attention_head_dim: int = 32 + num_attention_heads: int = 70 + num_cross_attention_heads: int = 20 + cross_attention_head_dim: int = 112 + cross_attention_dim: int = 2240 + caption_channels: int = 2304 + + mlp_ratio: float = 2.5 + # "rms_norm_across_heads" applies RMSNorm over the full (num_heads * head_dim) + + qk_norm: str = "rms_norm_across_heads" + norm_elementwise_affine: bool = False + norm_eps: float = 1e-6 + sample_size: int = 32 + guidance_embeds: bool = False + + param_names_mapping: dict = field( + default_factory=lambda: { + r"^transformer\.(.*)$": r"\1", + } + ) + + def __post_init__(self): + super().__post_init__() + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class SanaConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=SanaArchConfig) + prefix: str = "Sana" diff --git a/python/sglang/multimodal_gen/configs/models/encoders/__init__.py b/python/sglang/multimodal_gen/configs/models/encoders/__init__.py index f0b25420e..49d8430ee 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/encoders/__init__.py @@ -10,6 +10,7 @@ from sglang.multimodal_gen.configs.models.encoders.clip import ( CLIPTextConfig, CLIPVisionConfig, ) +from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig @@ -25,5 +26,6 @@ __all__ = [ "LlamaConfig", "Qwen3TextConfig", "T5Config", + "Gemma2Config", "Gemma3Config", ] diff --git a/python/sglang/multimodal_gen/configs/models/encoders/gemma2.py b/python/sglang/multimodal_gen/configs/models/encoders/gemma2.py new file mode 100644 index 000000000..278479f56 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/encoders/gemma2.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Text encoder configuration for Gemma2 2B, used by SANA for text conditioning. +# +# SANA uses the hidden states from Gemma2 (not logits) as the conditioning +# signal for cross-attention in the DiT. The encoder output dimension (2304) +# is projected to the DiT's inner_dim via caption_projection. +# +# Defaults match google/gemma-2-2b-it (the model used in SANA HF checkpoints). + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.encoders.base import ( + TextEncoderArchConfig, + TextEncoderConfig, +) + + +def _is_transformer_layer(n: str, m) -> bool: + return "layers" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("embed_tokens") + + +def _is_final_norm(n: str, m) -> bool: + return n.endswith("norm") + + +@dataclass +class Gemma2ArchConfig(TextEncoderArchConfig): + vocab_size: int = 256000 + hidden_size: int = 2304 + intermediate_size: int = 9216 + num_hidden_layers: int = 26 + num_attention_heads: int = 8 + num_key_value_heads: int = 4 + head_dim: int = 256 + hidden_act: str = "gelu_pytorch_tanh" + hidden_activation: str = "gelu_pytorch_tanh" + max_position_embeddings: int = 8192 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int = 0 + eos_token_id: int = 1 + bos_token_id: int = 2 + tie_word_embeddings: bool = True + rope_theta: float = 10000.0 + attention_bias: bool = False + attention_dropout: float = 0.0 + + # Gemma2 alternates between global and sliding-window attention + # on odd/even layers, respectively. + sliding_window: int = 4096 + + # query_pre_attn_scalar replaces the standard 1/sqrt(head_dim) scaling. + query_pre_attn_scalar: int = 256 + + # Softcapping bounds raw attention logits via tanh(logits/cap)*cap. + # NOTE: SDPA does not natively support softcapping; the runtime model + # currently skips this (see Gemma2Attention.forward). Quality impact + # is minimal for short text-encoder sequences but should be revisited + # for longer context. + attn_logit_softcapping: float = 50.0 + final_logit_softcapping: float = 30.0 + + text_len: int = 300 + + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=lambda: [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", "0"), + (".gate_up_proj", ".up_proj", "1"), + ] + ) + _fsdp_shard_conditions: list = field( + default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm] + ) + + +@dataclass +class Gemma2Config(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field(default_factory=Gemma2ArchConfig) + prefix: str = "gemma_2" diff --git a/python/sglang/multimodal_gen/configs/models/vaes/sana.py b/python/sglang/multimodal_gen/configs/models/vaes/sana.py new file mode 100644 index 000000000..c3c52e50e --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/vaes/sana.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# VAE configuration for SANA's DC-AE (Deep Compression AutoEncoder). +# +# DC-AE achieves a 32x spatial compression ratio (vs. 8x for standard SD VAEs), +# which means a 1024x1024 image becomes 32x32 latents with 32 channels. +# This aggressive compression is what allows SANA to run efficiently at +# high resolutions despite having a relatively small DiT. +# +# Reference: https://arxiv.org/abs/2405.17811 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class SanaVAEArchConfig(VAEArchConfig): + spatial_compression_ratio: int = 32 + # DC-AE uses a different scaling factor than standard VAEs; + # this value must match the pretrained checkpoint. + scaling_factor: float = 0.41407 + latent_channels: int = 32 + in_channels: int = 3 + + +@dataclass +class SanaVAEConfig(VAEConfig): + arch_config: SanaVAEArchConfig = field(default_factory=SanaVAEArchConfig) + + # DC-AE does not currently support tiling in our wrapper. + # Enable these once the diffusers AutoencoderDC adds tiling support. + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + + def post_init(self): + # Called by VAELoader AFTER update_model_arch() merges the HF config.json + # values into arch_config. Must be post_init() (not __post_init__) because + # __post_init__ fires at dataclass creation time, before the HF config merge. + # + # The base VAEConfig.get_vae_scale_factor() derives from block_out_channels, + # which DC-AE doesn't have. Set vae_scale_factor directly from the + # spatial_compression_ratio (32x for DC-AE). + self.arch_config.vae_scale_factor = self.arch_config.spatial_compression_ratio diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index f54b0ae9a..e2e7c36dc 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -29,6 +29,7 @@ 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.wan import ( SelfForcingWanT2V480PConfig, WanI2V480PConfig, @@ -51,6 +52,7 @@ __all__ = [ "Flux2KleinPipelineConfig", "Flux2FinetunedPipelineConfig", "PipelineConfig", + "SanaPipelineConfig", "SlidingTileAttnConfig", "MOVAPipelineConfig", "WanT2V480PConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py new file mode 100644 index 000000000..4815ab76a --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Pipeline configuration for SANA text-to-image generation. +# +# SANA produces 4D spatial latents (B, C, H', W') directly — unlike Flux/QwenImage +# which use packed token-style latents (B, S, D). This means: +# - We inherit SpatialImagePipelineConfig (not ImagePipelineConfig) +# - prepare_latent_shape returns 4D, not 5D +# - post_denoising_loop is a no-op (no un-packing needed) +# - shard_latents_for_sp shards along the H' dimension +# +# SANA does NOT use rotary position embeddings, so prepare_pos/neg_cond_kwargs +# return empty dicts (the DiT only needs hidden_states + encoder_hidden_states + timestep). +# +# CFG is handled by the denoising stage via guidance_scale in sampling params. +# should_use_guidance=False means no embedded guidance (no extra guidance token in forward), +# but negative_prompt + guidance_scale > 1.0 still enables standard classifier-free guidance. + +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig +from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput +from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig +from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config +from sglang.multimodal_gen.configs.models.vaes.sana import SanaVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + SpatialImagePipelineConfig, + preprocess_text, +) + + +def sana_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: + # SANA uses the final hidden state from Gemma2 directly as text conditioning. + # No intermediate-layer extraction or masking needed (unlike QwenImage/ZImage). + return outputs.last_hidden_state + + +@dataclass +class SanaPipelineConfig(SpatialImagePipelineConfig): + + task_type: ModelTaskType = ModelTaskType.T2I + + # should_use_guidance=False disables *embedded* guidance (timestep-conditioned + # guidance token). Standard CFG via guidance_scale is still active. + should_use_guidance: bool = False + enable_autocast: bool = False + + # DC-AE does not support tiling or SP VAE decode yet. + vae_tiling: bool = False + vae_sp: bool = False + vae_precision: str = "bf16" + + dit_config: DiTConfig = field(default_factory=SanaConfig) + vae_config: VAEConfig = field(default_factory=SanaVAEConfig) + + # Single text encoder: Gemma2 (unlike Flux which uses CLIP + T5) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (Gemma2Config(),) + ) + + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) + + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (preprocess_text,), + ) + + postprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (sana_postprocess_text,) + ) + + def prepare_latent_shape(self, batch, batch_size, num_frames): + # 4D latent shape: (B, C, H', W') — no temporal dim for T2I. + # DC-AE compresses 1024x1024 -> 32x32 with 32 channels. + compression = self.vae_config.arch_config.spatial_compression_ratio + height = batch.height // compression + width = batch.width // compression + num_channels = self.dit_config.arch_config.num_channels_latents + shape = (batch_size, num_channels, height, width) + return shape + + def get_pos_prompt_embeds(self, batch): + # Single encoder -> index [0] (Flux uses [1] because T5 is encoder #2) + 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): + # encoder_attention_mask: batch stores list-of-tensors; diffusers' SanaTransformer + # expects a single tensor (sglang's has list handling). Override with [0]. + out = {} + m = batch.prompt_attention_mask + if isinstance(m, (list, tuple)): + out["encoder_attention_mask"] = m[0] if m else None + elif m is not None: + out["encoder_attention_mask"] = m + return out + + def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): + out = {} + m = batch.negative_attention_mask + if isinstance(m, (list, tuple)): + out["encoder_attention_mask"] = m[0] if m else None + elif m is not None: + out["encoder_attention_mask"] = m + return out + + def post_denoising_loop(self, latents, batch): + return latents diff --git a/python/sglang/multimodal_gen/configs/sample/sana.py b/python/sglang/multimodal_gen/configs/sample/sana.py new file mode 100644 index 000000000..209ec8e8e --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/sana.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sampling parameters for SANA image generation (T2I).""" + +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) + + +@dataclass +class SanaSamplingParams(SamplingParams): + """Defaults for SANA 1.5 1024px variant. + + guidance_scale=4.5 enables standard classifier-free guidance. + """ + + data_type: DataType = DataType.IMAGE + num_frames: int = 1 + guidance_scale: float = 4.5 + num_inference_steps: int = 20 + height: int = 1024 + width: int = 1024 + negative_prompt: str = ( + "low quality, low resolution, blurry, overexposed, underexposed, " + "distorted, deformed, disfigured, bad anatomy, extra limbs, " + "watermark, text, signature, ugly, noisy, artifacts" + ) diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index d61fcb2ee..a8e810a68 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -63,6 +63,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageLayeredPipelineConfig, QwenImagePipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.wan import ( FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, @@ -98,6 +99,7 @@ from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImageLayeredSamplingParams, QwenImageSamplingParams, ) +from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams from sglang.multimodal_gen.configs.sample.wan import ( FastWanT2V480PConfig, Turbo_Wan2_2_I2V_A14B_SamplingParam, @@ -796,6 +798,21 @@ def _register_configs(): ], ) + # SANA + register_configs( + sampling_param_cls=SanaSamplingParams, + pipeline_config_cls=SanaPipelineConfig, + hf_model_paths=[ + "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers", + "Efficient-Large-Model/SANA1.5_4.8B_1024px_diffusers", + "Efficient-Large-Model/Sana_1600M_1024px_diffusers", + "Efficient-Large-Model/Sana_600M_1024px_diffusers", + "Efficient-Large-Model/Sana_1600M_512px_diffusers", + "Efficient-Large-Model/Sana_600M_512px_diffusers", + ], + model_detectors=[lambda hf_id: "sana" in hf_id.lower()], + ) + _register_configs() 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 e3a9dadf4..8a86f7989 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 @@ -128,9 +128,11 @@ class VAELoader(ComponentLoader): safetensors_list = _list_safetensors_files(component_model_path) assert ( - len(safetensors_list) == 1 - ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}" - loaded = safetensors_load_file(safetensors_list[0]) + len(safetensors_list) >= 1 + ), f"Found no safetensors files in {component_model_path}" + loaded = {} + for sf_path in safetensors_list: + loaded.update(safetensors_load_file(sf_path)) vae.load_state_dict(loaded, strict=False) state_keys = set(vae.state_dict().keys()) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 7aea2060d..2b2630fb1 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -369,15 +369,16 @@ def load_model_from_full_model_state_dict( if unused_keys: logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys) - # for nunchaku + # for nunchaku; norm_q/norm_k for SANA QK normalization layers ALLOWED_NEW_PARAM_PATTERNS = [ "gate_compress", "wcscales", "wtscale", "bias", + "norm_q", + "norm_k", ] for new_param_name in unused_keys: - # check unallowed missing params if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS): logger.error( "Unsupported new parameter: %s. Allowed patterns: %s", @@ -392,7 +393,9 @@ def load_model_from_full_model_state_dict( meta_sharded_param = meta_sd.get(new_param_name) meta_sharded_param_dtype = meta_sharded_param.dtype - if "wcscales" in new_param_name or "wtscale" in new_param_name: + if any( + p in new_param_name for p in ("wcscales", "wtscale", "norm_q", "norm_k") + ): init_like = torch.ones_like else: init_like = torch.zeros_like diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py new file mode 100644 index 000000000..d33ba49ba --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding + +from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig +from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm +from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class SanaCombinedTimestepSizeEmbeddings(nn.Module): + def __init__(self, embedding_dim): + super().__init__() + self.time_proj = Timesteps( + num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim + ) + + def forward(self, timestep, hidden_dtype=None): + timesteps_proj = self.time_proj(timestep) + if hidden_dtype is not None: + timesteps_proj = timesteps_proj.to(dtype=hidden_dtype) + timesteps_emb = self.timestep_embedder(timesteps_proj) + return timesteps_emb + + +class SanaAdaLayerNormSingle(nn.Module): + def __init__(self, embedding_dim): + super().__init__() + self.emb = SanaCombinedTimestepSizeEmbeddings(embedding_dim) + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + def forward(self, timestep, hidden_dtype=None): + embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype) + out = self.linear(self.silu(embedded_timestep)) + return out, embedded_timestep + + +class SanaModulatedNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + + def forward(self, x, temb, scale_shift_table): + x = self.norm(x) + shift, scale = (scale_shift_table[None] + temb[:, None]).chunk(2, dim=1) + x = x * (1 + scale) + shift + return x + + +class GLUMBConv(nn.Module): + """Gated Linear Unit with Multi-Branch Convolution.""" + + def __init__(self, in_channels, out_channels, expand_ratio=2.5): + super().__init__() + hidden_channels = int(expand_ratio * in_channels) + self.nonlinearity = nn.SiLU() + self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0) + self.conv_depth = nn.Conv2d( + hidden_channels * 2, + hidden_channels * 2, + 3, + 1, + 1, + groups=hidden_channels * 2, + ) + self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False) + + def forward(self, hidden_states): + hidden_states = self.conv_inverted(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv_depth(hidden_states) + hidden_states, gate = torch.chunk(hidden_states, 2, dim=1) + hidden_states = hidden_states * self.nonlinearity(gate) + hidden_states = self.conv_point(hidden_states) + return hidden_states + + +class SanaLinearAttention(nn.Module): + """Linear attention with O(N*D^2) complexity instead of O(N^2*D).""" + + def __init__(self, query_dim, num_heads, head_dim, qk_norm_dim, bias=False): + super().__init__() + inner_dim = num_heads * head_dim + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = nn.Linear(query_dim, inner_dim, bias=bias) + self.to_k = nn.Linear(query_dim, inner_dim, bias=bias) + self.to_v = nn.Linear(query_dim, inner_dim, bias=bias) + self.to_out = nn.ModuleList( + [nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()] + ) + self.norm_q = RMSNorm(qk_norm_dim) + self.norm_k = RMSNorm(qk_norm_dim) + + def forward(self, hidden_states): + B, S, _ = hidden_states.shape + + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + query = self.norm_q(query) + key = self.norm_k(key) + + query = query.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) + key = key.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) + value = value.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) + query = F.relu(query) + key = F.relu(key) + + kv = torch.matmul(key.transpose(-2, -1), value) # (B, H, D, D) + qkv = torch.matmul(query, kv) # (B, H, S, D) + + key_sum = key.sum(dim=-2, keepdim=True) # (B, H, 1, D) + normalizer = torch.matmul(query, key_sum.transpose(-2, -1)).clamp(min=1e-6) + hidden_states = qkv / normalizer + + hidden_states = hidden_states.transpose(1, 2).reshape(B, S, -1) + hidden_states = self.to_out[0](hidden_states) + return hidden_states + + +class SanaCrossAttention(nn.Module): + def __init__(self, query_dim, cross_attention_dim, num_heads, head_dim, bias=False): + super().__init__() + inner_dim = num_heads * head_dim + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = nn.Linear(query_dim, inner_dim, bias=bias) + self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias) + self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias) + self.to_out = nn.ModuleList( + [nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()] + ) + + self.norm_q = RMSNorm(inner_dim) + self.norm_k = RMSNorm(inner_dim) + + def forward( + self, hidden_states, encoder_hidden_states, encoder_attention_mask=None + ): + B, S, _ = hidden_states.shape + T = encoder_hidden_states.shape[1] + + query = self.to_q(hidden_states) + key = self.to_k(encoder_hidden_states) + value = self.to_v(encoder_hidden_states) + + query = self.norm_q(query) + key = self.norm_k(key) + + query = query.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) + key = key.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + value = value.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + + attn_mask = None + if encoder_attention_mask is not None: + attn_mask = encoder_attention_mask.bool() + attn_mask = attn_mask[:, None, None, :].expand(B, self.num_heads, S, T) + + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attn_mask + ) + hidden_states = hidden_states.transpose(1, 2).reshape(B, S, -1) + hidden_states = self.to_out[0](hidden_states) + return hidden_states + + +class SanaTransformerBlock(nn.Module): + def __init__( + self, + dim, + num_attention_heads, + attention_head_dim, + num_cross_attention_heads, + cross_attention_head_dim, + cross_attention_dim, + mlp_ratio, + norm_eps, + attention_bias=False, + ): + super().__init__() + + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) + self.attn1 = SanaLinearAttention( + query_dim=dim, + num_heads=num_attention_heads, + head_dim=attention_head_dim, + qk_norm_dim=num_attention_heads * attention_head_dim, + bias=attention_bias, + ) + + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) + self.attn2 = SanaCrossAttention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + num_heads=num_cross_attention_heads, + head_dim=cross_attention_head_dim, + bias=True, + ) + + self.ff = GLUMBConv(in_channels=dim, out_channels=dim, expand_ratio=mlp_ratio) + + def forward( + self, + hidden_states, + encoder_hidden_states, + timestep, + height, + width, + encoder_attention_mask=None, + ): + batch_size = hidden_states.shape[0] + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + + norm_hidden = self.norm1(hidden_states) + norm_hidden = norm_hidden * (1 + scale_msa) + shift_msa + attn_output = self.attn1(norm_hidden) + hidden_states = hidden_states + gate_msa * attn_output + + attn_output = self.attn2( + hidden_states, encoder_hidden_states, encoder_attention_mask + ) + hidden_states = hidden_states + attn_output + + norm_hidden = self.norm2(hidden_states) + norm_hidden = norm_hidden * (1 + scale_mlp) + shift_mlp + norm_hidden = norm_hidden.unflatten(1, (height, width)).permute(0, 3, 1, 2) + ff_output = self.ff(norm_hidden) + ff_output = ff_output.flatten(2, 3).permute(0, 2, 1) + hidden_states = hidden_states + gate_mlp * ff_output + + return hidden_states + + +class SanaTransformer2DModel(CachableDiT, OffloadableDiTMixin): + + _fsdp_shard_conditions = [ + lambda n, m: isinstance(m, SanaTransformerBlock), + ] + _compile_conditions = [ + lambda n, m: isinstance(m, SanaTransformerBlock), + ] + param_names_mapping = SanaConfig().arch_config.param_names_mapping + reverse_param_names_mapping = {} + + def __init__(self, config: SanaConfig, hf_config=None, **kwargs): + super().__init__(config, hf_config=hf_config or {}, **kwargs) + + arch = config.arch_config + self.out_channels = arch.out_channels + self.patch_size = arch.patch_size + self.inner_dim = arch.num_attention_heads * arch.attention_head_dim + + self.hidden_size = self.inner_dim + self.num_attention_heads = arch.num_attention_heads + self.num_channels_latents = arch.num_channels_latents + + self.patch_embed = nn.ModuleDict( + { + "proj": nn.Conv2d( + arch.in_channels, + self.inner_dim, + kernel_size=arch.patch_size, + stride=arch.patch_size, + bias=True, + ), + } + ) + self.time_embed = SanaAdaLayerNormSingle(self.inner_dim) + self.caption_projection = PixArtAlphaTextProjection( + in_features=arch.caption_channels, + hidden_size=self.inner_dim, + ) + + self.caption_norm = RMSNorm(self.inner_dim) + + self.transformer_blocks = nn.ModuleList( + [ + SanaTransformerBlock( + dim=self.inner_dim, + num_attention_heads=arch.num_attention_heads, + attention_head_dim=arch.attention_head_dim, + num_cross_attention_heads=arch.num_cross_attention_heads, + cross_attention_head_dim=arch.cross_attention_head_dim, + cross_attention_dim=arch.cross_attention_dim, + mlp_ratio=arch.mlp_ratio, + norm_eps=arch.norm_eps, + attention_bias=False, + ) + for _ in range(arch.num_layers) + ] + ) + self.scale_shift_table = nn.Parameter( + torch.randn(2, self.inner_dim) / self.inner_dim**0.5 + ) + + self.norm_out = SanaModulatedNorm(self.inner_dim, eps=arch.norm_eps) + + self.proj_out = nn.Linear( + self.inner_dim, + arch.patch_size * arch.patch_size * self.out_channels, + bias=True, + ) + + self.layer_names = ["transformer_blocks"] + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + timestep: torch.LongTensor = None, + guidance: torch.Tensor = None, + encoder_attention_mask: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + + # Input validation - fail fast + if encoder_hidden_states is None: + raise ValueError("SANA forward pass requires encoder_hidden_states") + + batch_size, channels, height, width = hidden_states.shape + p = self.patch_size + post_patch_height = height // p + post_patch_width = width // p + + hidden_states = self.patch_embed["proj"](hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + timestep_emb, embedded_timestep = self.time_embed( + timestep, hidden_dtype=hidden_states.dtype + ) + + if isinstance(encoder_attention_mask, (list, tuple)): + encoder_attention_mask = encoder_attention_mask[0] + + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + if encoder_hidden_states.shape[0] != batch_size: + encoder_hidden_states = encoder_hidden_states.expand( + batch_size, -1, -1 + ).contiguous() + encoder_hidden_states = encoder_hidden_states.view( + batch_size, -1, hidden_states.shape[-1] + ) + encoder_hidden_states = self.caption_norm(encoder_hidden_states) + + if ( + encoder_attention_mask is not None + and encoder_attention_mask.shape[0] != batch_size + ): + encoder_attention_mask = encoder_attention_mask.expand( + batch_size, -1 + ).contiguous() + + for block in self.transformer_blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_emb, + post_patch_height, + post_patch_width, + encoder_attention_mask=encoder_attention_mask, + ) + hidden_states = self.norm_out( + hidden_states, embedded_timestep, self.scale_shift_table + ) + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.reshape( + batch_size, post_patch_height, post_patch_width, p, p, self.out_channels + ) + hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4) + hidden_states = hidden_states.reshape( + batch_size, self.out_channels, height, width + ) + + return hidden_states + + +EntryClass = SanaTransformer2DModel diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/gemma2.py b/python/sglang/multimodal_gen/runtime/models/encoders/gemma2.py new file mode 100644 index 000000000..af18c4f2a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/encoders/gemma2.py @@ -0,0 +1,451 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Gemma2 2B text encoder for SANA. +# +# This is a decoder-only language model used as a text encoder: we feed +# in tokenized text and extract the final hidden states (not logits) as +# the conditioning signal for SANA's cross-attention layers. +# +# Architecture follows google/gemma-2-2b-it: +# - 26 layers, alternating global / sliding-window attention +# - GQA with 8 query heads, 4 KV heads, head_dim=256 +# - Pre/post attention + pre/post feedforward LayerNorm (Gemma2-style) +# - GeGLU activation (gelu_pytorch_tanh) +# +# Adapted from the Gemma3 text model implementation in this codebase. + +import logging +from typing import Any, Iterable + +import torch +from torch import nn + +from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput +from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config +from sglang.multimodal_gen.runtime.distributed import get_tp_world_size +from sglang.multimodal_gen.runtime.layers.activation import GeluAndMul +from sglang.multimodal_gen.runtime.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig +from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope +from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader + +logger = logging.getLogger(__name__) + + +class Gemma2RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()) + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + +class Gemma2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "gelu_pytorch_tanh": + raise ValueError( + "Gemma2 uses `gelu_pytorch_tanh` as the hidden activation. " + f"Got: {hidden_act}" + ) + self.act_fn = GeluAndMul(approximate="tanh") + + def forward(self, x): + x, _ = self.gate_up_proj(x) + x = self.act_fn(x) + x, _ = self.down_proj(x) + return x + + +class Gemma2Attention(nn.Module): + def __init__( + self, + layer_id: int, + config: Gemma2Config, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = layer_id + self.hidden_size = hidden_size + tp_size = get_tp_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + arch = config.arch_config + self.head_dim = arch.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = arch.query_pre_attn_scalar**-0.5 + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=arch.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=arch.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # Gemma2 interleaves global (even layers) and sliding-window (odd layers) + # attention. This pattern reduces memory for long sequences while + # maintaining global context every other layer. + self.is_sliding = (layer_id % 2) == 1 + if self.is_sliding: + self.sliding_window = arch.sliding_window + else: + self.sliding_window = None + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=arch.max_position_embeddings, + base=arch.rope_theta, + is_neox_style=True, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + batch_size, seq_len, _ = q.shape + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + + q, k = self.rotary_emb(positions, q, k) + + query = q.transpose(1, 2) + key = k.transpose(1, 2) + value = v.transpose(1, 2) + + attn_mask = torch.zeros( + (seq_len, seq_len), device=hidden_states.device, dtype=torch.float32 + ) + causal = torch.triu( + torch.ones( + (seq_len, seq_len), device=hidden_states.device, dtype=torch.bool + ), + diagonal=1, + ) + attn_mask = attn_mask.masked_fill(causal, float("-inf")) + if self.is_sliding and self.sliding_window is not None: + idx = torch.arange(seq_len, device=hidden_states.device) + dist = idx[None, :] - idx[:, None] + too_far = dist > self.sliding_window + attn_mask = attn_mask.masked_fill(too_far, float("-inf")) + + if attention_mask is not None: + key_pad = ~attention_mask.to(torch.bool) + attn_mask = attn_mask[None, None, :, :].expand( + batch_size, 1, seq_len, seq_len + ) + attn_mask = attn_mask.masked_fill( + key_pad[:, None, None, :].expand(batch_size, 1, seq_len, seq_len), + float("-inf"), + ) + + attn_kwargs = { + "attn_mask": attn_mask, + "dropout_p": 0.0, + "is_causal": False, + "scale": self.scaling, + } + if query.shape[1] != key.shape[1]: + attn_kwargs["enable_gqa"] = True + attn_output = torch.nn.functional.scaled_dot_product_attention( + query, key, value, **attn_kwargs + ) + + # NOTE: Gemma2 specifies attn_logit_softcapping (tanh(logits/cap)*cap) but + # PyTorch's scaled_dot_product_attention does not support it natively. + # For short text-encoder sequences (~300 tokens), the quality impact is + # negligible. A custom attention kernel would be needed for full fidelity. + + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape( + batch_size, seq_len, self.num_heads * self.head_dim + ) + + output, _ = self.o_proj(attn_output) + return output + + +class Gemma2DecoderLayer(nn.Module): + def __init__( + self, + layer_id: int, + config: Gemma2Config, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + arch = config.arch_config + self.hidden_size = arch.hidden_size + self.self_attn = Gemma2Attention( + layer_id=layer_id, + config=config, + hidden_size=self.hidden_size, + num_heads=arch.num_attention_heads, + num_kv_heads=arch.num_key_value_heads, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.mlp = Gemma2MLP( + hidden_size=self.hidden_size, + intermediate_size=arch.intermediate_size, + hidden_act=arch.hidden_activation, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = Gemma2RMSNorm(self.hidden_size, eps=arch.rms_norm_eps) + self.post_attention_layernorm = Gemma2RMSNorm( + self.hidden_size, eps=arch.rms_norm_eps + ) + self.pre_feedforward_layernorm = Gemma2RMSNorm( + self.hidden_size, eps=arch.rms_norm_eps + ) + self.post_feedforward_layernorm = Gemma2RMSNorm( + self.hidden_size, eps=arch.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(positions, hidden_states, attention_mask) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class Gemma2Model(nn.Module): + """Gemma2 text encoder model for SANA pipeline.""" + + _fsdp_shard_conditions = [] + + def __init__(self, config: Gemma2Config, **kwargs): + super().__init__() + self.config = config + arch = config.arch_config + self.quant_config = None + + self.vocab_size = arch.vocab_size + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + arch.hidden_size, + org_num_embeddings=arch.vocab_size, + quant_config=self.quant_config, + ) + self.embed_scale = arch.hidden_size**0.5 + + self.layers = nn.ModuleList( + [ + Gemma2DecoderLayer( + layer_id=i, + config=config, + quant_config=self.quant_config, + prefix=f"model.layers.{i}", + ) + for i in range(arch.num_hidden_layers) + ] + ) + + self.norm = Gemma2RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) * self.embed_scale + + def forward( + self, + input_ids: torch.Tensor | None = 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: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else getattr(self.config.arch_config, "output_hidden_states", False) + ) + + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + + if position_ids is None: + position_ids = torch.arange( + 0, hidden_states.shape[1], device=hidden_states.device + ).unsqueeze(0) + + all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None + + for layer in self.layers: + if all_hidden_states is not None: + all_hidden_states += (hidden_states,) + + hidden_states = layer(position_ids, hidden_states, attention_mask) + + hidden_states = self.norm(hidden_states) + + if all_hidden_states is not None: + all_hidden_states += (hidden_states,) + + return BaseEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + stacked_params_mapping = getattr( + self.config.arch_config, "stacked_params_mapping", None + ) + if stacked_params_mapping is None: + stacked_params_mapping = [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", "0"), + (".gate_up_proj", ".up_proj", "1"), + ] + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + # HF Gemma2Model stores weights as model.layers.X... / model.embed_tokens... + # Strip "model." prefix if present to match our naming + if name.startswith("model."): + name = name[len("model.") :] + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + self._load_with_shard_id(weight_loader, param, loaded_weight, shard_id) + break + else: + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + loaded_params.add(name) + return loaded_params + + @staticmethod + def _load_with_shard_id(weight_loader, param, loaded_weight, shard_id): + try: + weight_loader(param, loaded_weight, shard_id) + return + except (AssertionError, TypeError): + pass + + if isinstance(shard_id, str): + mapping = {"q": 0, "k": 1, "v": 2} + if shard_id in mapping: + weight_loader(param, loaded_weight, mapping[shard_id]) + return + if shard_id.isdigit(): + weight_loader(param, loaded_weight, int(shard_id)) + return + elif isinstance(shard_id, int): + mapping = {0: "q", 1: "k", 2: "v"} + if shard_id in mapping: + weight_loader(param, loaded_weight, mapping[shard_id]) + return + + raise TypeError( + f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader}" + ) + + +EntryClass = Gemma2Model diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_dpm_solver_multistep.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_dpm_solver_multistep.py new file mode 100644 index 000000000..8f72f22f0 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_dpm_solver_multistep.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# DPM-Solver++ multistep scheduler wrapper for SANA. +# +# SANA uses DPM-Solver++ (Lu et al., 2022) as its noise scheduler, which +# is a high-order ODE solver that converges in fewer steps than DDIM. +# With solver_order=2 and 20 steps, SANA achieves high-quality results. +# +# This wrapper delegates all numerical work to diffusers' implementation +# and only adapts the interface for sglang's denoising stage. + +import torch +from diffusers import ( + DPMSolverMultistepScheduler as DiffusersDPMSolverMultistepScheduler, +) +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin + +from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler + + +class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): + """DPM-Solver++ multistep scheduler wrapper for sglang's BaseScheduler interface.""" + + order = 1 + num_train_timesteps = 1000 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + beta_start: float = 0.0001, + beta_end: float = 0.02, + beta_schedule: str = "scaled_linear", + trained_betas=None, + solver_order: int = 2, + prediction_type: str = "epsilon", + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + algorithm_type: str = "dpmsolver++", + solver_type: str = "midpoint", + lower_order_final: bool = True, + euler_at_final: bool = False, + use_karras_sigmas: bool = False, + use_lu_lambdas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + use_flow_sigmas: bool = False, + final_sigmas_type: str = "zero", + lambda_min_clipped: float = -float("inf"), + variance_type: str | None = None, + timestep_spacing: str = "linspace", + steps_offset: int = 0, + rescale_betas_zero_snr: bool = False, + flow_shift: float | None = None, + **kwargs, + ): + self.num_train_timesteps = num_train_timesteps + self._inner = DiffusersDPMSolverMultistepScheduler( + num_train_timesteps=num_train_timesteps, + beta_start=beta_start, + beta_end=beta_end, + beta_schedule=beta_schedule, + trained_betas=trained_betas, + solver_order=solver_order, + prediction_type=prediction_type, + thresholding=thresholding, + dynamic_thresholding_ratio=dynamic_thresholding_ratio, + sample_max_value=sample_max_value, + algorithm_type=algorithm_type, + solver_type=solver_type, + lower_order_final=lower_order_final, + euler_at_final=euler_at_final, + use_karras_sigmas=use_karras_sigmas, + use_lu_lambdas=use_lu_lambdas, + use_exponential_sigmas=use_exponential_sigmas, + use_beta_sigmas=use_beta_sigmas, + use_flow_sigmas=use_flow_sigmas, + flow_shift=flow_shift, + final_sigmas_type=final_sigmas_type, + lambda_min_clipped=lambda_min_clipped, + variance_type=variance_type, + timestep_spacing=timestep_spacing, + steps_offset=steps_offset, + rescale_betas_zero_snr=rescale_betas_zero_snr, + ) + self.timesteps = self._inner.timesteps + self.order = solver_order + self._flow_shift = flow_shift + self._begin_index: int | None = None + BaseScheduler.__init__(self) + + def set_shift(self, shift: float) -> None: + self._flow_shift = shift + + def set_begin_index(self, begin_index: int = 0) -> None: + self._begin_index = begin_index + + @property + def begin_index(self) -> int | None: + return self._begin_index + + def set_timesteps(self, num_inference_steps: int, device=None, **kwargs): + self._inner.set_timesteps(num_inference_steps, device=device, **kwargs) + self.timesteps = self._inner.timesteps + + def scale_model_input( + self, sample: torch.Tensor, timestep: int | None = None + ) -> torch.Tensor: + return self._inner.scale_model_input(sample, timestep) + + def step( + self, + model_output: torch.Tensor, + timestep: int, + sample: torch.Tensor, + **kwargs, + ): + return self._inner.step(model_output, timestep, sample, **kwargs) + + @property + def sigmas(self): + return getattr(self._inner, "sigmas", None) + + @property + def init_noise_sigma(self): + return self._inner.init_noise_sigma + + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.Tensor, + ) -> torch.Tensor: + return self._inner.add_noise(original_samples, noise, timesteps) + + +EntryClass = DPMSolverMultistepScheduler diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_dc.py b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_dc.py new file mode 100644 index 000000000..f8434e949 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_dc.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Iterable + +import torch +from torch import nn + +from sglang.multimodal_gen.configs.models.vaes.sana import SanaVAEConfig +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class AutoencoderDC(nn.Module): + """Deep Compression Autoencoder wrapper with 32x spatial compression.""" + + def __init__(self, config: SanaVAEConfig = None, **kwargs): + super().__init__() + self._config = config + self._inner_model = None + self._loaded_state_dict: dict[str, torch.Tensor] = {} + + def _ensure_inner_model(self, state_dict: dict[str, torch.Tensor] | None = None): + if self._inner_model is not None: + return + + from diffusers import AutoencoderDC as DiffusersAutoencoderDC + + device = "cpu" + state_to_load = ( + state_dict if state_dict is not None else self._loaded_state_dict + ) + if state_to_load: + first_tensor = next(iter(state_to_load.values())) + device = first_tensor.device + hf_config = {} + if self._config is not None: + arch = self._config.arch_config + for key, value in vars(arch).items(): + if key == "extra_attrs" and isinstance(value, dict): + for ek, ev in value.items(): + hf_config[ek] = ev + elif not key.startswith("_") and not callable(value): + hf_config[key] = value + + self._inner_model = DiffusersAutoencoderDC.from_config(hf_config) + + if state_to_load: + missing, unexpected = self._inner_model.load_state_dict( + state_to_load, strict=False + ) + if missing: + logger.warning( + "AutoencoderDC missing keys when loading: %d keys", len(missing) + ) + if len(missing) > 10: + logger.debug("First 10 missing keys: %s", list(missing)[:10]) + else: + logger.debug("Missing keys: %s", list(missing)) + if unexpected: + logger.debug( + "AutoencoderDC unexpected keys when loading: %d keys", + len(unexpected), + ) + if state_dict is None: + self._loaded_state_dict.clear() + + self._inner_model = self._inner_model.to(device) + + @property + def config(self): + if self._inner_model is not None: + return self._inner_model.config + return self._config + + @property + def dtype(self): + if self._inner_model is not None: + return next(self._inner_model.parameters()).dtype + return torch.float32 + + @property + def device(self): + if self._inner_model is not None: + return next(self._inner_model.parameters()).device + return torch.device("cpu") + + def encode(self, x: torch.Tensor, **kwargs): + self._ensure_inner_model() + return self._inner_model.encode(x, **kwargs) + + def decode(self, z: torch.Tensor, **kwargs): + self._ensure_inner_model() + z = z.to(dtype=self.dtype) + return self._inner_model.decode(z, **kwargs) + + def forward(self, x: torch.Tensor, **kwargs): + self._ensure_inner_model() + return self._inner_model(x, **kwargs) + + def load_state_dict( + self, + state_dict: dict[str, torch.Tensor], + strict: bool = True, + assign: bool = False, + ): + """Intercept load_state_dict to route weights into the inner diffusers model.""" + self._ensure_inner_model(state_dict=state_dict) + + def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]: + self._ensure_inner_model() + return self._inner_model.state_dict(*args, **kwargs) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Buffer weights for deferred loading. The inner model is built lazily.""" + loaded_params: set[str] = set() + for name, weight in weights: + self._loaded_state_dict[name] = weight + loaded_params.add(name) + return loaded_params + + def to(self, *args, **kwargs): + if self._inner_model is not None: + self._inner_model = self._inner_model.to(*args, **kwargs) + return super().to(*args, **kwargs) + + +EntryClass = AutoencoderDC diff --git a/python/sglang/multimodal_gen/runtime/pipelines/sana.py b/python/sglang/multimodal_gen/runtime/pipelines/sana.py new file mode 100644 index 000000000..e42218bd1 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/sana.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# SANA text-to-image pipeline. +# +# Stage order matches Flux (InputValidation -> TextEncoding -> TimestepPrep -> +# LatentPrep -> Denoising -> Decoding) rather than the add_standard_t2i_stages +# helper (which puts LatentPrep before TimestepPrep). Both orderings are +# functionally equivalent since these stages are independent. +# +# SANA uses a single text encoder (Gemma2), so only one text_encoder + tokenizer +# pair is registered — unlike Flux which has text_encoder + text_encoder_2. +# The pipeline_name must match the _class_name in HF model_index.json. + +from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + InputValidationStage, + 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 SanaPipeline(LoRAPipeline, ComposedPipelineBase): + pipeline_name = "SanaPipeline" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "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")], + tokenizers=[self.get_module("tokenizer")], + ), + "prompt_encoding_stage_primary", + ) + + self.add_standard_timestep_preparation_stage() + self.add_standard_latent_preparation_stage() + self.add_standard_denoising_stage() + self.add_standard_decoding_stage() + + +EntryClass = SanaPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index 980d15210..39c7c9d41 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -118,10 +118,9 @@ class DecodingStage(PipelineStage): Decoded video tensor with shape (batch, channels, frames, height, width), normalized to [0, 1] range and moved to CPU as float32 """ - self.vae = self.vae.to(get_local_torch_device()) - latents = latents.to(get_local_torch_device()) - # Setup VAE precision vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + self.vae = self.vae.to(device=get_local_torch_device(), dtype=vae_dtype) + latents = latents.to(get_local_torch_device()) vae_autocast_enabled = ( vae_dtype != torch.float32 ) and not server_args.disable_autocast diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index a840408ce..8cea589b6 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -459,6 +459,15 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [ ), T2I_sampling_params, ), + DiffusionTestCase( + "sana_image_t2i", + DiffusionServerArgs( + model_path="Efficient-Large-Model/Sana_600M_1024px_diffusers", + modality="image", + ), + T2I_sampling_params, + run_perf_check=False, + ), # === Text and Image to Image (TI2I) === DiffusionTestCase( "qwen_image_edit_ti2i",