diff --git a/python/sglang/multimodal_gen/configs/models/dits/glmimage.py b/python/sglang/multimodal_gen/configs/models/dits/glmimage.py new file mode 100644 index 000000000..5fc2553e4 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/glmimage.py @@ -0,0 +1,39 @@ +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class GlmImageArchConfig(DiTArchConfig): + patch_size: int = 2 + in_channels: int = 16 + out_channels: int | None = 16 + num_layers: int = 30 + attention_head_dim: int = 128 + num_attention_heads: int = 32 + condition_dim: int = 256 + prior_vq_quantizer_codebook_size: int = 16384 + text_embed_dim: int = 1472 + time_embed_dim: int = 512 + + stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list) + + param_names_mapping: dict = field( + default_factory=lambda: { + # LoRA mappings + r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1", + } + ) + + def __post_init__(self): + super().__post_init__() + self.out_channels = self.out_channels or self.in_channels + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class GlmImageDitConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=GlmImageArchConfig) + + prefix: str = "glmimage" diff --git a/python/sglang/multimodal_gen/configs/models/vaes/glmimage.py b/python/sglang/multimodal_gen/configs/models/vaes/glmimage.py new file mode 100644 index 000000000..3c90cfef4 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/vaes/glmimage.py @@ -0,0 +1,60 @@ +from dataclasses import dataclass, field + +import torch + +from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class GlmImageVAEArchConfig(VAEArchConfig): + spatial_compression_ratio: int = 1 + + base_dim: int = 96 + decoder_base_dim: int | None = None + z_dim: int = 16 + dim_mult: tuple[int, ...] = (1, 2, 4, 4) + num_res_blocks: int = 2 + attn_scales: tuple[float, ...] = () + temperal_downsample: tuple[bool, ...] = (False, True, True) + dropout: float = 0.0 + + is_residual: bool = False + input_channels: int = 3 + out_channels: int = 3 + patch_size: int | None = None + scale_factor_temporal: int = 4 + scale_factor_spatial: int = 8 + clip_output: bool = True + + scaling_factor: float | torch.Tensor = 0 + + latents_mean: tuple[float, ...] | None = None + latents_std: tuple[float, ...] | None = None + shift_factor: float | None = None + latent_channels: int = 16 + in_channels: int = 16 + + +@dataclass +class GlmImageVAEConfig(VAEConfig): + arch_config: GlmImageVAEArchConfig = field(default_factory=GlmImageVAEArchConfig) + + use_feature_cache: bool = True + + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + + def get_vae_scale_factor(self): + return 2 ** len(self.arch_config.temperal_downsample) + + def __post_init__(self): + self.blend_num_frames = ( + self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + ) * 2 + + def post_init(self): + self.arch_config.vae_scale_factor = 2 ** ( + len(self.arch_config.temperal_downsample) + ) + self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 361551a30..82027f98e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -366,6 +366,9 @@ class PipelineConfig: latents = maybe_unpad_latents(latents, batch) return latents + def post_decoding(self, frames, server_args): + return frames + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return {} diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py new file mode 100644 index 000000000..0aa6903bb --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py @@ -0,0 +1,86 @@ +from dataclasses import dataclass, field + +import torch +from diffusers.image_processor import VaeImageProcessor + +from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits.glmimage import GlmImageDitConfig +from sglang.multimodal_gen.configs.models.vaes.glmimage import GlmImageVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ImagePipelineConfig, + ModelTaskType, +) + + +@dataclass +class GlmImagePipelineConfig(ImagePipelineConfig): + """Configuration for the GlmImage pipeline.""" + + vae_precision: str = "bf16" + + should_use_guidance: bool = False + task_type: ModelTaskType = ModelTaskType.T2I + + vae_tiling: bool = False + + vae_sp: bool = False + + dit_config: DiTConfig = field(default_factory=GlmImageDitConfig) + # VAE + vae_config: VAEConfig = field(default_factory=GlmImageVAEConfig) + + enable_autocast: bool = False + + def __post_init__(self): + self.vae_scale_factor = self.vae_config.get_vae_scale_factor() + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + def get_freqs_cis(self, batch, device, rotary_emb, dtype): + height = batch.height // self.vae_scale_factor + width = batch.width // self.vae_scale_factor + hidden_states = torch.empty(1, 1, height, width, device=device, dtype=dtype) + freqs_cis = rotary_emb(hidden_states) + return freqs_cis + + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): + return { + "prior_token_id": batch.prior_token_id, + "prior_token_drop": batch.prior_token_drop_cond, + "crop_coords": batch.crop_coords, + "target_size": batch.target_size, + "kv_caches": batch.kv_caches, + "kv_caches_mode": "read", + "freqs_cis": self.get_freqs_cis(batch, device, rotary_emb, dtype), + } + + def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): + return { + "prior_token_id": batch.prior_token_id, + "prior_token_drop": batch.prior_token_drop_uncond, + "crop_coords": batch.crop_coords, + "target_size": batch.target_size, + "kv_caches": batch.kv_caches, + "kv_caches_mode": "skip", + "freqs_cis": self.get_freqs_cis(batch, device, rotary_emb, dtype), + } + + def get_decode_scale_and_shift(self, device, dtype, vae): + latents_mean = ( + torch.tensor(self.vae_config.latents_mean) + .view(1, self.vae_config.latent_channels, 1, 1) + .to(device, dtype) + ) + latents_std = ( + torch.tensor(self.vae_config.latents_std) + .view(1, self.vae_config.latent_channels, 1, 1) + .to(device, dtype) + ) + return 1.0 / latents_std, latents_mean + + def post_denoising_loop(self, latents, batch): + if getattr(batch, "kv_caches", None) is not None: + batch.kv_caches.clear() + return latents.bfloat16() + + def post_decoding(self, frames, server_args): + return self.image_processor.postprocess(frames, output_type="latent") diff --git a/python/sglang/multimodal_gen/configs/sample/glmimage.py b/python/sglang/multimodal_gen/configs/sample/glmimage.py new file mode 100644 index 000000000..27ff3c741 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/glmimage.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + + +@dataclass +class GlmImageSamplingParams(SamplingParams): + negative_prompt = "" + + num_frames: int = 1 + guidance_scale: float = 1.5 + num_inference_steps: int = 30 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 5bd084f93..c244bc8b5 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -39,6 +39,9 @@ from sglang.multimodal_gen.configs.pipeline_configs import ( ) from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.glm_image import ( + GlmImagePipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageEditPipelineConfig, QwenImageEditPlus_2511_PipelineConfig, @@ -56,6 +59,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import ( Wan2_2_TI2V_5B_Config, ) from sglang.multimodal_gen.configs.sample.flux import FluxSamplingParams +from sglang.multimodal_gen.configs.sample.glmimage import GlmImageSamplingParams from sglang.multimodal_gen.configs.sample.hunyuan import ( FastHunyuanSamplingParam, HunyuanSamplingParams, @@ -581,5 +585,11 @@ def _register_configs(): hf_model_paths=["Qwen/Qwen-Image-Layered"], ) + register_configs( + sampling_param_cls=GlmImageSamplingParams, + pipeline_config_cls=GlmImagePipelineConfig, + model_detectors=[lambda hf_id: "glm-image" in hf_id.lower()], + ) + _register_configs() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index a4b0bd660..dd8bcae29 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -146,6 +146,7 @@ async def generations( ) resp_format = (request.response_format or "b64_json").lower() + if resp_format == "b64_json": with open(save_file_path, "rb") as f: b64 = base64.b64encode(f.read()).decode("utf-8") diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 5b7cf0cf8..da1dde89c 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -266,6 +266,7 @@ class ComponentLoader(ABC): "image_processor": (ImageProcessorLoader, "transformers"), "image_encoder": (ImageEncoderLoader, "transformers"), "processor": (AutoProcessorLoader, "transformers"), + "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"), } if module_type in module_loaders: @@ -779,6 +780,36 @@ class GenericComponentLoader(ComponentLoader): self.library = library +class VisionLanguageEncoderLoader(ComponentLoader): + """Loader for vision language encoder (typically Causal LM or Vision2Seq).""" + + def load_customized( + self, + component_model_path: str, + server_args: ServerArgs, + transformers_or_diffusers: str = "vision_language_encoder", + ) -> Any: + if transformers_or_diffusers == "vision_language_encoder": + from transformers import GlmImageForConditionalGeneration + + config = get_hf_config( + component_model_path, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + ) + model = GlmImageForConditionalGeneration.from_pretrained( + component_model_path, + config=config, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + ).to(get_local_torch_device()) + return model + else: + raise ValueError( + f"Unsupported library for VisionLanguageEncoder: {transformers_or_diffusers}" + ) + + class PipelineComponentLoader: """ Utility class for loading pipeline components. diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py new file mode 100644 index 000000000..439cc7242 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -0,0 +1,820 @@ +# Copyright 2025 The CogView team, Tsinghua University & ZhipuAI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.attention import FeedForward + +from sglang.multimodal_gen.configs.models.dits.glmimage import GlmImageDitConfig +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import ( + ScaleResidualLayerNormScaleShift, +) +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb +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.platforms import AttentionBackendEnum +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 GlmImageLayerKVCache: + """KV cache for GlmImage model.""" + + def __init__(self): + self.k_cache = None + self.v_cache = None + self.mode: Optional[str] = None # "write", "read", "skip" + + def store(self, k: torch.Tensor, v: torch.Tensor): + if self.k_cache is None: + self.k_cache = k + self.v_cache = v + else: + self.k_cache = torch.cat([self.k_cache, k], dim=2) + self.v_cache = torch.cat([self.v_cache, v], dim=2) + + def get(self): + return self.k_cache, self.v_cache + + def clear(self): + self.k_cache = None + self.v_cache = None + self.mode = None + + +class GlmImageKVCache: + """Container for all layers' KV caches.""" + + def __init__(self, num_layers: int): + self.num_layers = num_layers + self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)] + + def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache: + return self.caches[layer_idx] + + def set_mode(self, mode: Optional[str]): + if mode is not None and mode not in ["write", "read", "skip"]: + raise ValueError( + f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'" + ) + for cache in self.caches: + cache.mode = mode + + def clear(self): + for cache in self.caches: + cache.clear() + + +class GlmImageTimestepEmbedding(nn.Module): + """ + Replacement for diffusers TimestepEmbedding using ReplicatedLinear. + Structure: linear_1 -> act(silu) -> linear_2 + """ + + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + ): + super().__init__() + if out_dim is None: + out_dim = time_embed_dim + self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True) + if act_fn == "silu": + self.act = nn.SiLU() + elif act_fn == "gelu": + self.act = nn.GELU(approximate="tanh") + else: + self.act = nn.SiLU() + self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + sample, _ = self.linear_1(sample) + sample = self.act(sample) + sample, _ = self.linear_2(sample) + return sample + + +class GlmImageTextProjection(nn.Module): + """ + Replacement for diffusers PixArtAlphaTextProjection using ReplicatedLinear. + Structure: linear_1 -> act_1 -> linear_2 + """ + + def __init__( + self, + in_features: int, + hidden_size: int, + out_features: int = None, + act_fn: str = "silu", + ): + super().__init__() + if out_features is None: + out_features = hidden_size + self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True) + if act_fn == "silu": + self.act_1 = nn.SiLU() + elif act_fn == "gelu_tanh": + self.act_1 = nn.GELU(approximate="tanh") + else: + self.act_1 = nn.SiLU() + self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True) + + def forward(self, caption: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states, _ = self.linear_2(hidden_states) + return hidden_states + + +class GlmImageCombinedTimestepSizeEmbeddings(nn.Module): + def __init__( + self, + embedding_dim: int, + condition_dim: int, + pooled_projection_dim: int, + timesteps_dim: int = 256, + ): + super().__init__() + + self.time_proj = Timesteps( + num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.condition_proj = Timesteps( + num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.timestep_embedder = GlmImageTimestepEmbedding( + in_channels=timesteps_dim, time_embed_dim=embedding_dim + ) + self.condition_embedder = GlmImageTextProjection( + pooled_projection_dim, embedding_dim, act_fn="silu" + ) + + def forward( + self, + timestep: torch.Tensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + hidden_dtype: torch.dtype, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + + crop_coords_proj = self.condition_proj(crop_coords.flatten()).view( + crop_coords.size(0), -1 + ) + target_size_proj = self.condition_proj(target_size.flatten()).view( + target_size.size(0), -1 + ) + + # (B, 2 * condition_dim) + condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1) + + timesteps_emb = self.timestep_embedder( + timesteps_proj.to(dtype=hidden_dtype) + ) # (B, embedding_dim) + condition_emb = self.condition_embedder( + condition_proj.to(dtype=hidden_dtype) + ) # (B, embedding_dim) + + conditioning = timesteps_emb + condition_emb + return conditioning + + +class GlmImageImageProjector(nn.Module): + def __init__( + self, + in_channels: int = 16, + hidden_size: int = 2560, + patch_size: int = 2, + ): + super().__init__() + self.patch_size = patch_size + + self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, channel, height, width = hidden_states.shape + post_patch_height = height // self.patch_size + post_patch_width = width // self.patch_size + + hidden_states = hidden_states.reshape( + batch_size, + channel, + post_patch_height, + self.patch_size, + post_patch_width, + self.patch_size, + ) + hidden_states = ( + hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + ) + hidden_states = self.proj(hidden_states) + + return hidden_states + + +class GlmImageAdaLayerNormZero(nn.Module): + def __init__(self, embedding_dim: int, dim: int) -> None: + super().__init__() + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + dtype = hidden_states.dtype + norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) + norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to( + dtype=dtype + ) + + emb, _ = self.linear(temb) + ( + shift_msa, + c_shift_msa, + scale_msa, + c_scale_msa, + gate_msa, + c_gate_msa, + shift_mlp, + c_shift_mlp, + scale_mlp, + c_scale_mlp, + gate_mlp, + c_gate_mlp, + ) = emb.chunk(12, dim=1) + + hidden_states = norm_hidden_states * ( + 1 + scale_msa.unsqueeze(1) + ) + shift_msa.unsqueeze(1) + encoder_hidden_states = norm_encoder_hidden_states * ( + 1 + c_scale_msa.unsqueeze(1) + ) + c_shift_msa.unsqueeze(1) + + return ( + hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + encoder_hidden_states, + c_gate_msa, + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) + + +class GlmImageAttention(torch.nn.Module): + def __init__( + self, + query_dim, + heads, + dim_head, + out_dim, + bias, + qk_norm, + elementwise_affine, + eps, + supported_attention_backends: set[AttentionBackendEnum] | None = None, + prefix: str = "", + ): + super().__init__() + + self.k_cache = None + self.v_cache = None + + self.heads = out_dim // dim_head if out_dim is not None else heads + self.dim_head = dim_head + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.inner_kv_dim = self.inner_dim + self.out_dim = out_dim if out_dim is not None else query_dim + + self.num_kv_heads = self.dim_head // self.inner_kv_dim + + self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias) + self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias) + self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias) + + # (dropout omitted) + self.to_out = nn.ModuleList( + [ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)] + ) + + if qk_norm is None: + self.norm_q = None + self.norm_k = None + elif qk_norm == "layer_norm": + self.norm_q = nn.LayerNorm( + dim_head, eps=eps, elementwise_affine=elementwise_affine + ) + self.norm_k = nn.LayerNorm( + dim_head, eps=eps, elementwise_affine=elementwise_affine + ) + else: + raise ValueError( + f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'fp32_layer_norm', 'layer_norm_across_heads', 'rms_norm', 'rms_norm_across_heads', 'l2'." + ) + + self.attn = USPAttention( + num_heads=self.heads, + head_size=dim_head, + num_kv_heads=self.num_kv_heads, + dropout_rate=0, + softmax_scale=None, + causal=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + kv_cache: Optional[GlmImageLayerKVCache] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + dtype = encoder_hidden_states.dtype + + batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape + batch_size, image_seq_length, embed_dim = hidden_states.shape + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + # 1. QKV projections + query, _ = self.to_q(hidden_states) + key, _ = self.to_k(hidden_states) + value, _ = self.to_v(hidden_states) + + query = query.unflatten(2, (self.heads, -1)) + key = key.unflatten(2, (self.heads, -1)) + value = value.unflatten(2, (self.heads, -1)) + + # 2. QK normalization + if self.norm_q is not None: + query = self.norm_q(query).to(dtype=dtype) + if self.norm_k is not None: + key = self.norm_k(key).to(dtype=dtype) + + # 3. Rotational positional embeddings applied to latent stream + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + + query[:, text_seq_length:, :, :] = _apply_rotary_emb( + query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True + ) + key[:, text_seq_length:, :, :] = _apply_rotary_emb( + key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True + ) + + if kv_cache is not None: + if kv_cache.mode == "write": + kv_cache.store(key, value) + elif kv_cache.mode == "read": + k_cache, v_cache = kv_cache.get() + key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key + value = ( + torch.cat([v_cache, value], dim=1) if v_cache is not None else value + ) + elif kv_cache.mode == "skip": + pass + + # 4. Attention + if attention_mask is not None: + text_attn_mask = attention_mask + assert ( + text_attn_mask.dim() == 2 + ), "the shape of text_attn_mask should be (batch_size, text_seq_length)" + text_attn_mask = text_attn_mask.float().to(query.device) + mix_attn_mask = torch.ones( + (batch_size, text_seq_length + image_seq_length), device=query.device + ) + mix_attn_mask[:, :text_seq_length] = text_attn_mask + mix_attn_mask = mix_attn_mask.unsqueeze(2) + attn_mask_matrix = mix_attn_mask @ mix_attn_mask.transpose(1, 2) + attention_mask = (attn_mask_matrix > 0).unsqueeze(1).to(query.dtype) + hidden_states = self.attn(query, key, value) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + # 5. Output projection + hidden_states, _ = self.to_out[0](hidden_states) + # hidden_states = self.to_out[1](hidden_states) # (dropout omitted) + + encoder_hidden_states, hidden_states = hidden_states.split( + [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1 + ) + return hidden_states, encoder_hidden_states + + +class GlmImageTransformerBlock(nn.Module): + def __init__( + self, + dim: int = 2560, + num_attention_heads: int = 64, + attention_head_dim: int = 40, + time_embed_dim: int = 512, + supported_attention_backends: set[AttentionBackendEnum] | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + # 1. Attention + self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim) + + self.attn1 = GlmImageAttention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + out_dim=dim, + bias=True, + qk_norm="layer_norm", + elementwise_affine=False, + eps=1e-5, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1", + ) + + # 2. Feedforward + self.norm2 = ScaleResidualLayerNormScaleShift( + dim, norm_type="layer", eps=1e-5, elementwise_affine=False + ) + self.norm2_context = ScaleResidualLayerNormScaleShift( + dim, norm_type="layer", eps=1e-5, elementwise_affine=False + ) + self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[ + Union[ + Tuple[torch.Tensor, torch.Tensor], + List[Tuple[torch.Tensor, torch.Tensor]], + ] + ] = None, + attention_mask: Optional[Dict[str, torch.Tensor]] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + kv_cache: Optional[GlmImageLayerKVCache] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # 1. Timestep conditioning + ( + norm_hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + norm_encoder_hidden_states, + c_gate_msa, + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) = self.norm1(hidden_states, encoder_hidden_states, temb) + + # 2. Attention + if attention_kwargs is None: + attention_kwargs = {} + + attn_hidden_states, attn_encoder_hidden_states = self.attn1( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + kv_cache=kv_cache, + **attention_kwargs, + ) + + # 3. Feedforward (fused residual + norm + scale/shift) + norm_hidden_states, hidden_states = self.norm2( + hidden_states, + attn_hidden_states, + gate_msa.unsqueeze(1), + shift_mlp.unsqueeze(1), + scale_mlp.unsqueeze(1), + ) + norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context( + encoder_hidden_states, + attn_encoder_hidden_states, + c_gate_msa.unsqueeze(1), + c_shift_mlp.unsqueeze(1), + c_scale_mlp.unsqueeze(1), + ) + + ff_output = self.ff(norm_hidden_states) + ff_output_context = self.ff(norm_encoder_hidden_states) + hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1) + encoder_hidden_states = ( + encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1) + ) + + return hidden_states, encoder_hidden_states + + +class GlmImageRotaryPosEmbed(nn.Module): + def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None: + super().__init__() + + self.dim = dim + self.patch_size = patch_size + self.theta = theta + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + batch_size, num_channels, height, width = hidden_states.shape + height, width = height // self.patch_size, width // self.patch_size + device = hidden_states.device + + dim_h, dim_w = self.dim // 2, self.dim // 2 + h_inv_freq = 1.0 / ( + self.theta + ** ( + torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[ + : (dim_h // 2) + ].float() + / dim_h + ) + ) + w_inv_freq = 1.0 / ( + self.theta + ** ( + torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[ + : (dim_w // 2) + ].float() + / dim_w + ) + ) + h_seq = torch.arange(height, device=device) + w_seq = torch.arange(width, device=device) + freqs_h = torch.outer(h_seq, h_inv_freq) + freqs_w = torch.outer(w_seq, w_inv_freq) + + # Create position matrices for height and width + # [height, 1, dim//4] and [1, width, dim//4] + freqs_h = freqs_h.unsqueeze(1) + freqs_w = freqs_w.unsqueeze(0) + # Broadcast freqs_h and freqs_w to [height, width, dim//4] + freqs_h = freqs_h.expand(height, width, -1) + freqs_w = freqs_w.expand(height, width, -1) + + # Concatenate along last dimension to get [height, width, dim//2] + freqs = torch.cat([freqs_h, freqs_w], dim=-1) + freqs = freqs.reshape(height * width, -1) # [height * width, dim//2] + return (freqs.cos(), freqs.sin()) + + +class GlmImageAdaLayerNormContinuous(nn.Module): + """ + GlmImage-only final AdaLN: LN(x) -> Linear(cond) -> chunk -> affine. Matches Megatron: **no activation** before the + Linear on conditioning embedding. + """ + + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + bias: bool = True, + norm_type: str = "layer_norm", + ): + super().__init__() + self.linear = nn.Linear( + conditioning_embedding_dim, embedding_dim * 2, bias=bias + ) + if norm_type == "layer_norm": + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) + # For now, don’t replace this with sglang’s LayerNorm + # because the model doesn’t have this parameter and it will break model loading + elif norm_type == "rms_norm": + self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward( + self, x: torch.Tensor, conditioning_embedding: torch.Tensor + ) -> torch.Tensor: + # *** NO SiLU here *** + emb = self.linear(conditioning_embedding.to(x.dtype)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +class GlmImageTransformer2DModel(CachableDiT, OffloadableDiTMixin): + r""" + Args: + patch_size (`int`, defaults to `2`): + The size of the patches to use in the patch embedding layer. + in_channels (`int`, defaults to `16`): + The number of channels in the input. + num_layers (`int`, defaults to `30`): + The number of layers of Transformer blocks to use. + attention_head_dim (`int`, defaults to `40`): + The number of channels in each head. + num_attention_heads (`int`, defaults to `64`): + The number of heads to use for multi-head attention. + out_channels (`int`, defaults to `16`): + The number of channels in the output. + text_embed_dim (`int`, defaults to `1472`): + Input dimension of text embeddings from the text encoder. + time_embed_dim (`int`, defaults to `512`): + Output dimension of timestep embeddings. + condition_dim (`int`, defaults to `256`): + The embedding dimension of the input SDXL-style resolution conditions (original_size, target_size, + crop_coords). + pos_embed_max_size (`int`, defaults to `128`): + The maximum resolution of the positional embeddings, from which slices of shape `H x W` are taken and added + to input patched latents, where `H` and `W` are the latent height and width respectively. A value of 128 + means that the maximum supported height and width for image generation is `128 * vae_scale_factor * + patch_size => 128 * 8 * 2 => 2048`. + sample_size (`int`, defaults to `128`): + The base resolution of input latents. If height/width is not provided during generation, this value is used + to determine the resolution as `sample_size * vae_scale_factor => 128 * 8 => 1024` + """ + + def __init__( + self, + config: GlmImageDitConfig, + hf_config: dict[str, Any], + ): + super().__init__(config=config, hf_config=hf_config) + + self.config_data = config # Store config + arch_config = config.arch_config + + self.in_channels = arch_config.in_channels + self.out_channels = arch_config.out_channels + self.patch_size = arch_config.patch_size + self.num_layers = arch_config.num_layers + self.attention_head_dim = arch_config.attention_head_dim + self.num_attention_heads = arch_config.num_attention_heads + self.text_embed_dim = arch_config.text_embed_dim + self.time_embed_dim = arch_config.time_embed_dim + + # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords + # Each of these are sincos embeddings of shape 2 * condition_dim + pooled_projection_dim = 2 * 2 * arch_config.condition_dim + inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim + + # 1. RoPE + self.rotary_emb = GlmImageRotaryPosEmbed( + arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0 + ) + + # 2. Patch & Text-timestep embedding + self.image_projector = GlmImageImageProjector( + arch_config.in_channels, inner_dim, arch_config.patch_size + ) + self.glyph_projector = FeedForward( + arch_config.text_embed_dim, + inner_dim, + inner_dim=inner_dim, + activation_fn="gelu", + ) + self.prior_token_embedding = nn.Embedding( + arch_config.prior_vq_quantizer_codebook_size, inner_dim + ) + self.prior_projector = FeedForward( + inner_dim, inner_dim, inner_dim=inner_dim, activation_fn="linear-silu" + ) + + self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( + embedding_dim=arch_config.time_embed_dim, + condition_dim=arch_config.condition_dim, + pooled_projection_dim=pooled_projection_dim, + timesteps_dim=arch_config.time_embed_dim, + ) + + # 3. Transformer blocks + self._supported_attention_backends = arch_config._supported_attention_backends + self.transformer_blocks = nn.ModuleList( + [ + GlmImageTransformerBlock( + inner_dim, + arch_config.num_attention_heads, + arch_config.attention_head_dim, + arch_config.time_embed_dim, + supported_attention_backends=self._supported_attention_backends, + prefix=f"transformer_blocks.{i}", + ) + for i in range(arch_config.num_layers) + ] + ) + + # 4. Output projection + self.norm_out = GlmImageAdaLayerNormContinuous( + inner_dim, arch_config.time_embed_dim, elementwise_affine=False + ) + self.proj_out = nn.Linear( + inner_dim, + arch_config.patch_size * arch_config.patch_size * arch_config.out_channels, + bias=True, + ) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + prior_token_id: torch.Tensor, + prior_token_drop: torch.Tensor, + timestep: torch.LongTensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + attention_kwargs: Optional[Dict[str, Any]] = None, + attention_mask: Optional[torch.Tensor] = None, + kv_caches: Optional[GlmImageKVCache] = None, + kv_caches_mode: Optional[str] = None, + freqs_cis: Optional[ + Union[ + Tuple[torch.Tensor, torch.Tensor], + List[Tuple[torch.Tensor, torch.Tensor]], + ] + ] = None, + ### + guidance: torch.Tensor = None, # TODO: this should probably be removed + ) -> Tuple[torch.Tensor]: + if kv_caches is not None: + kv_caches.set_mode(kv_caches_mode) + + batch_size, num_channels, height, width = hidden_states.shape + + timestep -= 1.0 + + if isinstance(encoder_hidden_states, list): + encoder_hidden_states = encoder_hidden_states[0] + + # 1. RoPE + image_rotary_emb = freqs_cis + if image_rotary_emb is None: + image_rotary_emb = self.rotary_emb(hidden_states) + # 2. Patch & Timestep embeddings + p = self.config.patch_size + post_patch_height = height // p + post_patch_width = width // p + + hidden_states = self.image_projector(hidden_states) + encoder_hidden_states = self.glyph_projector(encoder_hidden_states) + prior_embedding = self.prior_token_embedding(prior_token_id) + prior_embedding[prior_token_drop] *= 0.0 + prior_hidden_states = self.prior_projector(prior_embedding) + hidden_states = hidden_states + prior_hidden_states + + temb = self.time_condition_embed( + timestep, target_size, crop_coords, hidden_states.dtype + ) + temb = F.silu(temb) + + # 3. Transformer blocks + for idx, block in enumerate(self.transformer_blocks): + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + attention_mask, + attention_kwargs, + kv_cache=kv_caches[idx] if kv_caches is not None else None, + ) + + # 4. Output norm & projection + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, post_patch_height, post_patch_width, -1, p, p + ) + output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + + return output.float() + # float() + # reference: https://github.com/zRzRzRzRzRzRzR/diffusers/blob/6cfc83b4abc5b083fef56a18ec4700f48ba3aaba/src/diffusers/pipelines/glm_image/pipeline_glm_image.py#L737 + + +EntryClass = GlmImageTransformer2DModel diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/t5.py b/python/sglang/multimodal_gen/runtime/models/encoders/t5.py index 048308ad1..c502c0340 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/t5.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/t5.py @@ -189,7 +189,7 @@ class T5Attention(nn.Module): self.qkv_proj = QKVParallelLinear( self.d_model, - self.d_model // self.total_num_heads, + self.key_value_proj_dim, self.total_num_heads, self.total_num_kv_heads, bias=False, @@ -208,7 +208,7 @@ class T5Attention(nn.Module): quant_config=quant_config, ) self.o = RowParallelLinear( - self.d_model, + self.total_num_heads * self.key_value_proj_dim, self.d_model, bias=False, quant_config=quant_config, @@ -307,7 +307,10 @@ class T5Attention(nn.Module): ) -> torch.Tensor: bs, seq_len, _ = hidden_states.shape num_seqs = bs - n, c = self.n_heads, self.d_model // self.total_num_heads + n, c = ( + self.n_heads, + self.key_value_proj_dim, + ) qkv, _ = self.qkv_proj(hidden_states) # Projection of 'own' hidden state (self-attention). No GQA here. q, k, v = qkv.split(self.inner_dim, dim=-1) diff --git a/python/sglang/multimodal_gen/runtime/models/model_stages/glm_image.py b/python/sglang/multimodal_gen/runtime/models/model_stages/glm_image.py new file mode 100644 index 000000000..48476a881 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/model_stages/glm_image.py @@ -0,0 +1,823 @@ +import inspect +import re +import time +from math import sqrt +from typing import List, Optional, Tuple, Union + +import numpy as np +import PIL +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.utils.torch_utils import randn_tensor + +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.models.dits.glm_image import GlmImageKVCache +from sglang.multimodal_gen.runtime.models.vision_utils import load_image +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + base_shift: float = 0.25, + max_shift: float = 0.75, +) -> float: + m = (image_seq_len / base_seq_len) ** 0.5 + mu = m * max_shift + base_shift + return mu + + +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + """ + accepts_timesteps = "timesteps" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys() + ) + accepts_sigmas = "sigmas" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys() + ) + + if timesteps is not None and sigmas is not None: + if not accepts_timesteps and not accepts_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep or sigma schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps( + timesteps=timesteps, sigmas=sigmas, device=device, **kwargs + ) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif timesteps is not None and sigmas is None: + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif timesteps is None and sigmas is not None: + if not accepts_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, + generator: Optional[torch.Generator] = None, + sample_mode: str = "sample", +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +class GlmImageBeforeDenoisingStage(PipelineStage): + r""" + Pipeline for text-to-image generation using GLM-Image. + + This pipeline integrates both the AR (autoregressive) model for token generation and the DiT (diffusion + transformer) model for image decoding. + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`T5EncoderModel`]): + Frozen text-encoder for glyph embeddings. + tokenizer (`PreTrainedTokenizer`): + Tokenizer for the text encoder. + processor (`AutoProcessor`): + Processor for the AR model to handle chat templates and tokenization. + vision_language_encoder ([`GlmImageForConditionalGeneration`]): + The AR model that generates image tokens from text prompts. + transformer ([`GlmImageTransformer2DModel`]): + A text conditioned transformer to denoise the encoded image latents (DiT). + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + """ + + def __init__( + self, + tokenizer, + processor, + text_encoder, + vision_language_encoder, + vae, + transformer, + scheduler, + ): + super().__init__() + + self.tokenizer = tokenizer + self.processor = processor + self.text_encoder = text_encoder + self.vision_language_encoder = vision_language_encoder + self.vae = vae + self.transformer = transformer + self.scheduler = scheduler + + self.vae_scale_factor = ( + 2 ** (len(self.vae.config.block_out_channels) - 1) + if getattr(self, "vae", None) + else 8 + ) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + self.default_sample_size = ( + self.transformer.config.sample_size + if hasattr(self, "transformer") + and self.transformer is not None + and hasattr(self.transformer.config, "sample_size") + else 128 + ) + + def _parse_and_expand_shape_info( + self, prompt: str + ) -> Tuple[str, int, int, int, int]: + """ + Parse the shape info from prompt and expand it for AR model. + + Args: + prompt: The prompt containing H W shape specification + + Returns: + Tuple of (expanded_prompt, token_h, token_w, prev_token_h, prev_token_w) + """ + match = re.search(r"(\d+)\s+(\d+)", prompt) + if match is None: + raise ValueError( + f"Prompt must contain shape info in format 'H W', got: {prompt}" + ) + + token_h, token_w = int(match.group(1)), int(match.group(2)) + ratio = token_h / token_w + prev_token_h = int(sqrt(ratio) * 16) + prev_token_w = int(sqrt(1 / ratio) * 16) + + old_shape = f"{token_h} {token_w}" + new_shape = ( + f"{token_h} {token_w}{prev_token_h} {prev_token_w}" + ) + expanded_prompt = prompt.replace(old_shape, new_shape) + + return expanded_prompt, token_h, token_w, prev_token_h, prev_token_w + + def _build_image_grid_thw( + self, + token_h: int, + token_w: int, + prev_token_h: int, + prev_token_w: int, + existing_grid: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + ) -> torch.Tensor: + """ + Build image grid tensor for AR model. + + For text-to-image: creates grid for large image + small image For image-to-image: appends new image to existing + grid + """ + if existing_grid is None or existing_grid.numel() == 0: + # Text-to-image: large image + small image + return torch.tensor( + [ + [1, token_h, token_w], + [1, prev_token_h, prev_token_w], + ], + device=device, + ) + else: + # Image-to-image: append to existing + return torch.cat( + [existing_grid, torch.tensor([[1, token_h, token_w]], device=device)], + dim=0, + ) + + def _calculate_ar_generation_params( + self, + token_h: int, + token_w: int, + prev_token_h: int, + prev_token_w: int, + is_text_to_image: bool, + ) -> Tuple[int, int]: + """ + Calculate max_new_tokens and large_image_start_offset for AR generation. + """ + large_image_tokens = token_h * token_w + small_image_tokens = prev_token_h * prev_token_w + + if is_text_to_image: + max_new_tokens = small_image_tokens + large_image_tokens + 1 + large_image_start_offset = small_image_tokens + else: + max_new_tokens = large_image_tokens + 1 + large_image_start_offset = 0 + + return max_new_tokens, large_image_start_offset + + def _extract_large_image_tokens( + self, + outputs: torch.Tensor, + input_length: int, + large_image_start_offset: int, + large_image_tokens: int, + ) -> torch.Tensor: + """ + Extract the large image tokens from AR model output. + """ + generated_tokens = outputs[0][input_length:] + large_image_start = large_image_start_offset + large_image_end = large_image_start + large_image_tokens + return generated_tokens[large_image_start:large_image_end] + + def _upsample_d32_to_d16( + self, token_ids: torch.Tensor, token_h: int, token_w: int + ) -> torch.Tensor: + """ + Upsample token IDs from d32 format to d16 format. + + AR model generates tokens at d32 resolution (each token = 32x32 pixels). DiT expects tokens at d16 resolution + (each token = 16x16 pixels). This function performs 2x nearest-neighbor upsampling. + + Args: + token_ids: Token IDs of shape [N] where N = token_h * token_w + token_h: Height in d32 token units + token_w: Width in d32 token units + + Returns: + Upsampled token IDs of shape [1, N*4] where N*4 = (token_h*2) * (token_w*2) + """ + # Reshape to spatial format: [1, 1, H, W] + token_ids = token_ids.view(1, 1, token_h, token_w) + + # 2x nearest-neighbor upsampling + token_ids = torch.nn.functional.interpolate( + token_ids.float(), scale_factor=2, mode="nearest" + ).to(dtype=torch.long) + + # Flatten back to [1, H*W*4] + token_ids = token_ids.view(1, -1) + + return token_ids + + @staticmethod + def _compute_generation_params( + image_grid_thw, + is_text_to_image: bool, + ): + grid_sizes = [] + grid_hw = [] + + for i in range(image_grid_thw.shape[0]): + t, h, w = image_grid_thw[i].tolist() + grid_sizes.append(int(h * w)) + grid_hw.append((int(h), int(w))) + + if not is_text_to_image: + max_new_tokens = grid_sizes[-1] + 1 + large_image_start_offset = 0 + target_grid_h, target_grid_w = grid_hw[-1] + else: + total_tokens = sum(grid_sizes) + max_new_tokens = total_tokens + 1 + large_image_start_offset = sum(grid_sizes[1:]) + target_grid_h, target_grid_w = grid_hw[0] + return max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w + + @staticmethod + def _upsample_token_ids( + token_ids: torch.Tensor, token_h: int, token_w: int + ) -> torch.Tensor: + token_ids = token_ids.view(1, 1, token_h, token_w) + token_ids = torch.nn.functional.interpolate( + token_ids.float(), scale_factor=2, mode="nearest" + ).to(dtype=torch.long) + token_ids = token_ids.view(1, -1) + return token_ids + + def generate_prior_tokens( + self, + prompt: str, + height: int, + width: int, + image: Optional[List[PIL.Image.Image]] = None, + factor: int = 32, + ) -> Tuple[torch.Tensor, int, int]: + """ + Generate prior tokens using the AR (vision_language_encoder) model. + + Args: + prompt: The text prompt with shape info (e.g., "description36 24") + condition_images: Optional list of condition images for i2i + + Returns: + Tuple of (prior_token_ids, pixel_height, pixel_width) + - prior_token_ids: Upsampled to d16 format, shape [1, token_h*token_w*4] + - pixel_height: Image height in pixels + - pixel_width: Image width in pixels + """ + device = self.vision_language_encoder.device + height = (height // factor) * factor + width = (width // factor) * factor + + is_text_to_image = image is None or len(image) == 0 + # Build messages for processor + content = [] + if image is not None: + for img in image: + content.append({"type": "image", "image": img}) + content.append({"type": "text", "text": prompt}) + messages = [{"role": "user", "content": content}] + + inputs = self.processor.apply_chat_template( + messages, + tokenize=True, + target_h=height, + target_w=width, + return_dict=True, + return_tensors="pt", + ).to(device) + + image_grid_thw = inputs.get("image_grid_thw") + max_new_tokens, large_image_offset, token_h, token_w = ( + self._compute_generation_params( + image_grid_thw=image_grid_thw, is_text_to_image=is_text_to_image + ) + ) + + prior_token_image_ids = None + if image is not None: + prior_token_image_embed = self.vision_language_encoder.get_image_features( + inputs["pixel_values"], image_grid_thw[:-1] + ) + prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0) + prior_token_image_ids = self.vision_language_encoder.get_image_tokens( + prior_token_image_embed, image_grid_thw[:-1] + ) + + # For GLM-Image, greedy decoding is not allowed; it may cause repetitive outputs. + # max_new_tokens must be exactly grid_h * grid_w + 1 (the +1 is for EOS). + outputs = self.vision_language_encoder.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=True, + ) + + prior_token_ids_d32 = self._extract_large_image_tokens( + outputs, + inputs["input_ids"].shape[-1], + large_image_offset, + token_h * token_w, + ) + prior_token_ids = self._upsample_token_ids( + prior_token_ids_d32, token_h, token_w + ) + + return prior_token_ids, prior_token_image_ids + + def get_glyph_texts(self, prompt): + prompt = prompt[0] if isinstance(prompt, list) else prompt + ocr_texts = ( + re.findall(r"'([^']*)'", prompt) + + re.findall(r"“([^“”]*)”", prompt) + + re.findall(r'"([^"]*)"', prompt) + + re.findall(r"「([^「」]*)」", prompt) + ) + return ocr_texts + + def _get_glyph_embeds( + self, + prompt: Union[str, List[str]] = None, + max_sequence_length: int = 2048, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + glyph_texts = self.get_glyph_texts(prompt) + input_ids = self.tokenizer( + glyph_texts if len(glyph_texts) > 0 else [""], + max_length=max_sequence_length, + truncation=True, + ).input_ids + input_ids = [ + [self.tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + input_ids_ + for input_ids_ in input_ids + ] + max_length = max(len(input_ids_) for input_ids_ in input_ids) + attention_mask = torch.tensor( + [ + [1] * len(input_ids_) + [0] * (max_length - len(input_ids_)) + for input_ids_ in input_ids + ], + device=device, + ) + input_ids = torch.tensor( + [ + input_ids_ + + [self.tokenizer.pad_token_id] * (max_length - len(input_ids_)) + for input_ids_ in input_ids + ], + device=device, + ) + outputs = self.text_encoder(input_ids, attention_mask=attention_mask) + glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0) + + return glyph_embeds.to(device=device, dtype=dtype) + + def encode_prompt( + self, + prompt: Union[str, List[str]], + do_classifier_free_guidance: bool = True, + prompt_embeds: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + max_sequence_length: int = 2048, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + max_sequence_length (`int`, defaults to `2048`): + Maximum sequence length in encoded prompt. Can be set to other values but may lead to poorer results. + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds = self._get_glyph_embeds( + prompt, max_sequence_length, device, dtype + ) + + seq_len = prompt_embeds.size(1) + prompt_embeds = prompt_embeds.repeat(1, 1, 1) + prompt_embeds = prompt_embeds.view(1, seq_len, -1) + + negative_prompt_embeds = None + if do_classifier_free_guidance: + negative_prompt = "" + negative_prompt = ( + batch_size * [negative_prompt] + if isinstance(negative_prompt, str) + else negative_prompt + ) + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds = self._get_glyph_embeds( + negative_prompt, max_sequence_length, device, dtype + ) + + seq_len = negative_prompt_embeds.size(1) + negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1) + negative_prompt_embeds = negative_prompt_embeds.view(1, seq_len, -1) + + return prompt_embeds, negative_prompt_embeds + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + ): + + shape = ( + batch_size, + num_channels_latents, + int(height) // self.vae_scale_factor, + int(width) // self.vae_scale_factor, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def check_inputs( + self, + prompt, + height, + width, + callback_on_step_end_tensor_inputs, + prompt_embeds=None, + ): + if ( + height is not None + and height % (self.vae_scale_factor * self.transformer.config.patch_size) + != 0 + or width is not None + and width % (self.transformer.config.patch_size) != 0 + ): + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs + for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and ( + not isinstance(prompt, str) and not isinstance(prompt, list) + ): + raise ValueError( + f"`prompt` has to be of type `str` or `list` but is {type(prompt)}" + ) + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + def forward( + self, + batch: Req, + server_args: ServerArgs, + ) -> Req: + + guidance_scale = batch.guidance_scale + prompt = batch.prompt + num_inference_steps = batch.num_inference_steps + if batch.image_path is not None: + ar_condition_images = [ + load_image(img_path) for img_path in batch.image_path + ] + else: + ar_condition_images = None + + height = batch.height + width = batch.width + + device = get_local_torch_device() + max_sequence_length = 1024 + generator = torch.Generator(device=device).manual_seed(batch.seed) + attention_kwargs = {} + prompt_embeds = None + do_classifier_free_guidance = True + dtype = torch.bfloat16 + + self._guidance_scale = guidance_scale + self._current_timestep = None + self._interrupt = False + + batch_size = 1 + + device = get_local_torch_device() + + if ar_condition_images is not None: + height = height or ar_condition_images[0].height + width = width or ar_condition_images[0].width + time_start = time.time() + prior_token_id, prior_token_image_ids = self.generate_prior_tokens( + prompt=prompt, + image=ar_condition_images, + height=height, + width=width, + ) + prior_token_id = prior_token_id.to(device=device) + time_end = time.time() + logger.info(f"generate_prior_tokens time: {time_end - time_start}") + + # 3. Encode input prompt + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + # 4. process images + if ar_condition_images is not None: + preprocessed_condition_images = [] + for img in ar_condition_images: + image_height, image_width = ( + img.size[::-1] + if isinstance(img, PIL.Image.Image) + else img.shape[:2] + ) + multiple_of = self.vae_scale_factor * self.transformer.config.patch_size + image_height = (image_height // multiple_of) * multiple_of + image_width = (image_width // multiple_of) * multiple_of + img = self.image_processor.preprocess( + img, height=image_height, width=image_width + ) + preprocessed_condition_images.append(img) + ar_condition_images = preprocessed_condition_images + + # 5. Prepare latents and (optional) condition_images kv cache + latent_channels = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size=1, + num_channels_latents=latent_channels, + height=height, + width=width, + dtype=torch.float32, + device=device, + generator=generator, + ) + + kv_caches = GlmImageKVCache(num_layers=self.transformer.config.num_layers) + + if ar_condition_images is not None: + latents_mean = torch.tensor(self.vae.config.latents_mean).view( + 1, self.vae.config.latent_channels, 1, 1 + ) + latents_std = torch.tensor(self.vae.config.latents_std).view( + 1, self.vae.config.latent_channels, 1, 1 + ) + + latents_mean = latents_mean.to(device=device, dtype=prompt_embeds.dtype) + latents_std = latents_std.to(device=device, dtype=prompt_embeds.dtype) + + for condition_image, condition_image_prior_token_id in zip( + ar_condition_images, prior_token_image_ids + ): + condition_image = condition_image.to( + device=device, dtype=prompt_embeds.dtype + ) + + condition_latent = retrieve_latents( + self.vae.encode(condition_image), + generator=generator, + sample_mode="argmax", + ) + condition_latent = (condition_latent - latents_mean) / latents_std + + # Do not remove. + # It would be use to run the reference image through a + # forward pass at timestep 0 and keep the KV cache. + with set_forward_context(current_timestep=1, attn_metadata=None): + _ = self.transformer( + hidden_states=condition_latent, + encoder_hidden_states=torch.zeros_like(prompt_embeds)[ + :1, :0, ... + ], + prior_token_id=condition_image_prior_token_id, + prior_token_drop=torch.full_like( + condition_image_prior_token_id, False, dtype=torch.bool + ), + timestep=torch.zeros((1,), device=device), + target_size=torch.tensor( + [condition_image.shape[-2:]], device=device + ), + crop_coords=torch.zeros((1, 2), device=device), + attention_kwargs=attention_kwargs, + kv_caches=kv_caches, + kv_caches_mode="write", + ) + + # 6. Prepare additional timestep conditions + target_size = (height, width) + target_size = torch.tensor( + [target_size], dtype=prompt_embeds.dtype, device=device + ) + crops_coords_top_left = torch.tensor( + [(0, 0)], dtype=prompt_embeds.dtype, device=device + ) + + # Prepare timesteps + image_seq_len = ( + (height // self.vae_scale_factor) * (width // self.vae_scale_factor) + ) // (self.transformer.config.patch_size**2) + timesteps = np.linspace( + self.scheduler.config.num_train_timesteps, 1.0, num_inference_steps + 1 + )[:-1] + timesteps = timesteps.astype(np.int64).astype(np.float32) + sigmas = timesteps / self.scheduler.config.num_train_timesteps + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("base_shift", 0.25), + self.scheduler.config.get("max_shift", 0.75), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps, sigmas, mu=mu + ) + self._num_timesteps = len(timesteps) + + # 7. Prepare for denoising loop + + batch.prompt_embeds = [prompt_embeds] + batch.negative_prompt_embeds = [negative_prompt_embeds] + batch.latents = latents + batch.timesteps = timesteps + batch.num_inference_steps = num_inference_steps + batch.sigmas = sigmas.tolist() # Convert numpy array to list for validation + batch.generator = generator + batch.raw_latent_shape = latents.shape + + batch.prior_token_id = prior_token_id + batch.prior_token_drop_cond = torch.full_like( + prior_token_id, False, dtype=torch.bool + ) + batch.prior_token_drop_uncond = torch.full_like( + prior_token_id, True, dtype=torch.bool + ) + batch.target_size = target_size + batch.crop_coords = crops_coords_top_left + + batch.kv_caches = kv_caches + + batch.height = height + batch.width = width + + return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py b/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py new file mode 100644 index 000000000..9e74a197a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py @@ -0,0 +1,58 @@ +from sglang.multimodal_gen.runtime.models.model_stages.glm_image import ( + GlmImageBeforeDenoisingStage, +) +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 ( + DecodingStage, + DenoisingStage, +) +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 GlmImagePipeline(LoRAPipeline, ComposedPipelineBase): + pipeline_name = "GlmImagePipeline" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "vision_language_encoder", + "processor", + "transformer", + "scheduler", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + self.add_stage( + stage_name="GlmImageBeforeDenoisingStage", + stage=GlmImageBeforeDenoisingStage( + vae=self.get_module("vae"), + text_encoder=self.get_module("text_encoder"), + tokenizer=self.get_module("tokenizer"), + processor=self.get_module("processor"), + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vision_language_encoder=self.get_module("vision_language_encoder"), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = [GlmImagePipeline] 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 4ea1c1503..123bc9fb2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -224,6 +224,8 @@ class DecodingStage(PipelineStage): else: trajectory_decoded = None + frames = server_args.pipeline_config.post_decoding(frames, server_args) + # Update batch with decoded image output_batch = OutputBatch( output=frames,