chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] model: GLM-Image (#16894)

Browse Source 
Co-authored-by: jianyingzhu <53300651@qq.com>
This commit is contained in:
Yuhao Yang
2026-01-14 02:02:03 +08:00
committed by GitHub
parent 2a7b67adff
commit a0b4ba9032
13 changed files with 1951 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
python/sglang/multimodal_gen/configs/models/dits/glmimage.py Normal file
View File

@@ -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"

60
python/sglang/multimodal_gen/configs/models/vaes/glmimage.py Normal file
View File

@@ -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

3
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
View File

@@ -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 {}

86
python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py Normal file
View File

@@ -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")

12
python/sglang/multimodal_gen/configs/sample/glmimage.py Normal file
View File

@@ -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

10
python/sglang/multimodal_gen/registry.py
View File

@@ -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()

1
python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py
View File

@@ -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")

31
python/sglang/multimodal_gen/runtime/loader/component_loader.py
View File

@@ -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.

820
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py Normal file
View File

@@ -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

9
python/sglang/multimodal_gen/runtime/models/encoders/t5.py
View File

@@ -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)

823
python/sglang/multimodal_gen/runtime/models/model_stages/glm_image.py Normal file
View File

@@ -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 <sop>H W<eop> shape specification
Returns:
Tuple of (expanded_prompt, token_h, token_w, prev_token_h, prev_token_w)
"""
match = re.search(r"<sop>(\d+)\s+(\d+)<eop>", prompt)
if match is None:
raise ValueError(
f"Prompt must contain shape info in format '<sop>H W<eop>', 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"<sop>{token_h} {token_w}<eop>"
new_shape = (
f"<sop>{token_h} {token_w}<eop><sop>{prev_token_h} {prev_token_w}<eop>"
)
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., "description<sop>36 24<eop>")
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

58
python/sglang/multimodal_gen/runtime/pipelines/glm_image.py Normal file
View File

@@ -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]

2
python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
View File

@@ -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,