From 841eb29d3d97725eb91ec2c31ff4e46728a893e7 Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Fri, 28 Nov 2025 21:48:31 +0800 Subject: [PATCH] [diffusion] model: support z-image (#14067) --- .../configs/models/dits/zimage.py | 40 + .../configs/pipeline_configs/__init__.py | 2 + .../configs/pipeline_configs/flux.py | 4 +- .../configs/pipeline_configs/zimage.py | 74 ++ .../multimodal_gen/configs/sample/zimage.py | 32 + python/sglang/multimodal_gen/registry.py | 11 + .../runtime/models/dits/zimage.py | 730 ++++++++++++++++++ .../runtime/models/encoders/stepllm.py | 2 +- .../runtime/pipelines/zimage_pipeline.py | 116 +++ .../pipelines_core/stages/causal_denoising.py | 2 +- .../pipelines_core/stages/conditioning.py | 2 +- .../pipelines_core/stages/denoising.py | 2 +- .../pipelines_core/stages/input_validation.py | 4 +- .../test/server/perf_baselines.json | 25 + .../test/server/testcase_configs.py | 13 + 15 files changed, 1051 insertions(+), 8 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/models/dits/zimage.py create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py create mode 100644 python/sglang/multimodal_gen/configs/sample/zimage.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/zimage.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/zimage.py b/python/sglang/multimodal_gen/configs/models/dits/zimage.py new file mode 100644 index 000000000..7d3576fcd --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/zimage.py @@ -0,0 +1,40 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field +from typing import Tuple + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class ZImageArchConfig(DiTArchConfig): + all_patch_size: Tuple[int, ...] = (2,) + all_f_patch_size: Tuple[int, ...] = (1,) + in_channels: int = 16 + out_channels: int | None = None + dim: int = 3840 + num_layers: int = 30 + n_refiner_layers: int = 2 + num_attention_heads: int = 30 + n_kv_heads: int = 30 + norm_eps: float = 1e-5 + qk_norm: bool = True + cap_feat_dim: int = 2560 + rope_theta: float = 256.0 + t_scale: float = 1000.0 + axes_dims: Tuple[int, int, int] = (32, 48, 48) + axes_lens: Tuple[int, int, int] = (1024, 512, 512) + + def __post_init__(self): + super().__post_init__() + self.out_channels = self.out_channels or self.in_channels + self.num_channels_latents = self.in_channels + self.hidden_size = self.dim + + +@dataclass +class ZImageDitConfig(DiTConfig): + arch_config: ZImageArchConfig = field(default_factory=ZImageArchConfig) + + prefix: str = "zimage" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index 370ec38e9..250c0d7ab 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -17,6 +17,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import ( WanT2V480PConfig, WanT2V720PConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig __all__ = [ "HunyuanConfig", @@ -30,4 +31,5 @@ __all__ = [ "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig", + "ZImagePipelineConfig", ] diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index b8a2601c3..825e64d00 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -321,7 +321,7 @@ def _prepare_image_ids( return image_latent_ids -def flux_2_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: +def flux2_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: hidden_states_layers: list[int] = [10, 20, 30] out = torch.stack([outputs.hidden_states[k] for k in hidden_states_layers], dim=1) @@ -412,7 +412,7 @@ class Flux2PipelineConfig(FluxPipelineConfig): ) postprocess_text_funcs: tuple[Callable[[str], str], ...] = field( - default_factory=lambda: (flux_2_postprocess_text,) + default_factory=lambda: (flux2_postprocess_text,) ) vae_config: VAEConfig = field(default_factory=Flux2VAEConfig) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py new file mode 100644 index 000000000..b107243f9 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py @@ -0,0 +1,74 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +from dataclasses import dataclass, field +from typing import Callable + +import torch + +from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig +from sglang.multimodal_gen.configs.models.encoders import ( + BaseEncoderOutput, + TextEncoderConfig, +) +from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ImagePipelineConfig, + ModelTaskType, +) + + +def zimage_preprocess_text(prompt: str): + messages = [ + {"role": "user", "content": prompt}, + ] + return messages + + +def zimage_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: + device = outputs.hidden_states[-2].device + prompt_mask = _text_inputs.attention_mask.to(device).bool() + return outputs.hidden_states[-2][0][prompt_mask[0]] + + +class TransformersModelConfig(EncoderConfig): + tokenizer_kwargs: dict = field(default_factory=lambda: {}) + + +@dataclass +class ZImagePipelineConfig(ImagePipelineConfig): + + should_use_guidance: bool = False + task_type: ModelTaskType = ModelTaskType.T2I + + dit_config: DiTConfig = field(default_factory=ZImageDitConfig) + vae_config: VAEConfig = field(default_factory=FluxVAEConfig) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (TextEncoderConfig(),) + ) + + preprocess_text_funcs: tuple[Callable, ...] = field( + default_factory=lambda: (zimage_preprocess_text,) + ) + + postprocess_text_funcs: tuple[Callable, ...] = field( + default_factory=lambda: (zimage_postprocess_text,) + ) + + def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict: + # flatten to 1-d list + inputs = tokenizer.apply_chat_template( + prompts, + tokenize=True, + add_generation_prompt=True, + enable_thinking=True, + padding="max_length", + max_length=512, # TODO (yhyang201): set max length according to config + truncation=True, + return_tensors="pt", + return_dict=True, + ) + return inputs + + def post_denoising_loop(self, latents, batch): + bs, channels, num_frames, height, width = latents.shape + return latents.view(bs, channels, height, width) diff --git a/python/sglang/multimodal_gen/configs/sample/zimage.py b/python/sglang/multimodal_gen/configs/sample/zimage.py new file mode 100644 index 000000000..2772167e3 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/zimage.py @@ -0,0 +1,32 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.sample.base import SamplingParams +from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams + + +@dataclass +class ZImageSamplingParams(SamplingParams): + num_inference_steps: int = 9 + + num_frames: int = 1 + height: int = 720 + width: int = 1280 + fps: int = 24 + + guidance_scale: float = 0.0 + + teacache_params: TeaCacheParams = field( + default_factory=lambda: TeaCacheParams( + teacache_thresh=0.15, + coefficients=[ + 7.33226126e02, + -4.01131952e02, + 6.75869174e01, + -3.14987800e00, + 9.61237896e-02, + ], + ) + ) diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index bf74b8c1e..97e6c7c61 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -24,6 +24,7 @@ from sglang.multimodal_gen.configs.pipeline_configs import ( WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig, + ZImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig @@ -56,6 +57,7 @@ from sglang.multimodal_gen.configs.sample.wan import ( WanT2V_1_3B_SamplingParams, WanT2V_14B_SamplingParams, ) +from sglang.multimodal_gen.configs.sample.zimage import ZImageSamplingParams from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) @@ -404,6 +406,15 @@ def _register_configs(): ], model_detectors=[lambda id: "flux.2" in id.lower()], ) + register_configs( + model_name="Z-image", + sampling_param_cls=ZImageSamplingParams, + pipeline_config_cls=ZImagePipelineConfig, + model_paths=[ + "Tongyi-MAI/Z-Image-Turbo", + ], + model_detectors=[lambda id: "z-image" in id.lower()], + ) # Qwen-Image register_configs( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py new file mode 100644 index 000000000..ab9af9114 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -0,0 +1,730 @@ +import math +from typing import Any, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence + +from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +ADALN_EMBED_DIM = 256 +SEQ_MULTI_OF = 32 + + +class SelectFirstElement(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return x[0] + + +class TimestepEmbedder(nn.Module): + def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): + super().__init__() + if mid_size is None: + mid_size = out_size + + self.mlp = nn.ModuleList( + [ + ReplicatedLinear(frequency_embedding_size, mid_size, bias=True), + nn.SiLU(), + ReplicatedLinear(mid_size, out_size, bias=True), + ] + ) + + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + with torch.amp.autocast("cuda", enabled=False): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) + / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to( + self.mlp[0].weight.dtype + ) + t_emb, _ = self.mlp[0](t_freq) + t_emb = self.mlp[1](t_emb) + t_emb, _ = self.mlp[2](t_emb) + return t_emb + + +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False) + self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False) + self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False) + + def _forward_silu_gating(self, x1, x3): + return F.silu(x1) * x3 + + def forward(self, x): + x1, _ = self.w1(x) + x3, _ = self.w3(x) + out, _ = self.w2(self._forward_silu_gating(x1, x3)) + return out + + +class ZImageAttention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + num_kv_heads: int, + qk_norm: bool = True, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = dim // num_heads + self.qk_norm = qk_norm + + self.to_q = ReplicatedLinear(dim, dim, bias=False) + self.to_k = ReplicatedLinear(dim, self.head_dim * num_kv_heads, bias=False) + self.to_v = ReplicatedLinear(dim, self.head_dim * num_kv_heads, bias=False) + + if self.qk_norm: + self.norm_q = RMSNorm(self.head_dim, eps=eps) + self.norm_k = RMSNorm(self.head_dim, eps=eps) + else: + self.norm_q = None + self.norm_k = None + + self.to_out = nn.ModuleList([ReplicatedLinear(dim, dim, bias=False)]) + + self.attn = USPAttention( + num_heads=num_heads, + head_size=self.head_dim, + num_kv_heads=num_kv_heads, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends={ + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + }, + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ): + q, _ = self.to_q(hidden_states) + k, _ = self.to_k(hidden_states) + v, _ = self.to_v(hidden_states) + + q = q.view(*q.shape[:-1], self.num_heads, self.head_dim) + k = k.view(*k.shape[:-1], self.num_kv_heads, self.head_dim) + v = v.view(*v.shape[:-1], self.num_kv_heads, self.head_dim) + + if self.norm_q is not None: + q = self.norm_q(q) + if self.norm_k is not None: + k = self.norm_k(k) + + # Apply RoPE + def apply_rotary_emb( + x_in: torch.Tensor, freqs_cis: torch.Tensor + ) -> torch.Tensor: + with torch.amp.autocast("cuda", enabled=False): + x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x * freqs_cis).flatten(3) + return x_out.type_as(x_in) # todo + + if freqs_cis is not None: + q = apply_rotary_emb(q, freqs_cis) + k = apply_rotary_emb(k, freqs_cis) + + hidden_states = self.attn(q, k, v) + hidden_states = hidden_states.flatten(2) + + hidden_states, _ = self.to_out[0](hidden_states) + + return hidden_states + + +class ZImageTransformerBlock(nn.Module): + def __init__( + self, + layer_id: int, + dim: int, + n_heads: int, + n_kv_heads: int, + norm_eps: float, + qk_norm: bool, + modulation=True, + ): + super().__init__() + self.dim = dim + self.head_dim = dim // n_heads + self.layer_id = layer_id + self.modulation = modulation + + self.attention = ZImageAttention( + dim=dim, + num_heads=n_heads, + num_kv_heads=n_kv_heads, + qk_norm=qk_norm, + eps=1e-5, + ) + + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) + + self.attention_norm1 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) + + self.attention_norm2 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + if modulation: + self.adaLN_modulation = nn.Sequential( + ReplicatedLinear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True) + ) + + def forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor, + freqs_cis: Tuple[torch.Tensor, torch.Tensor], + adaln_input: Optional[torch.Tensor] = None, + ): + if self.modulation: + assert adaln_input is not None + scale_msa_gate, _ = self.adaLN_modulation(adaln_input) + scale_msa, gate_msa, scale_mlp, gate_mlp = scale_msa_gate.unsqueeze( + 1 + ).chunk(4, dim=2) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + + # Attention block + attn_out = self.attention( + self.attention_norm1(x) * scale_msa, + freqs_cis=freqs_cis, + ) + x = x + gate_msa * self.attention_norm2(attn_out) + + # FFN block + x = x + gate_mlp * self.ffn_norm2( + self.feed_forward( + self.ffn_norm1(x) * scale_mlp, + ) + ) + else: + # Attention block + attn_out = self.attention( + self.attention_norm1(x), + freqs_cis=freqs_cis, + ) + x = x + self.attention_norm2(attn_out) + + # FFN block + x = x + self.ffn_norm2( + self.feed_forward( + self.ffn_norm1(x), + ) + ) + + return x + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = ReplicatedLinear(hidden_size, out_channels, bias=True) + + self.act = nn.SiLU() + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + ReplicatedLinear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), + ) + + def forward(self, x, c): + scale, _ = self.adaLN_modulation(c) + scale = 1.0 + scale + x = self.norm_final(x) * scale.unsqueeze(1) + x, _ = self.linear(x) + return x + + +class RopeEmbedder: + def __init__( + self, + theta: float = 256.0, + axes_dims: List[int] = (16, 56, 56), + axes_lens: List[int] = (64, 128, 128), + ): + self.theta = theta + self.axes_dims = axes_dims + self.axes_lens = axes_lens + assert len(axes_dims) == len( + axes_lens + ), "axes_dims and axes_lens must have the same length" + self.freqs_cis = None + + @staticmethod + def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0): + with torch.device("cpu"): + freqs_cis = [] + for i, (d, e) in enumerate(zip(dim, end)): + freqs = 1.0 / ( + theta + ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d) + ) + timestep = torch.arange(e, device=freqs.device, dtype=torch.float64) + freqs = torch.outer(timestep, freqs).float() + freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to( + torch.complex64 + ) # complex64 + freqs_cis.append(freqs_cis_i) + + return freqs_cis + + def __call__(self, ids: torch.Tensor): + assert ids.ndim == 2 + assert ids.shape[-1] == len(self.axes_dims) + device = ids.device + + if self.freqs_cis is None: + self.freqs_cis = self.precompute_freqs_cis( + self.axes_dims, self.axes_lens, theta=self.theta + ) + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + else: + # Ensure freqs_cis are on the same device as ids + if self.freqs_cis[0].device != device: + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + + result = [] + for i in range(len(self.axes_dims)): + index = ids[:, i] + result.append(self.freqs_cis[i][index]) + return torch.cat(result, dim=-1) + + +class ZImageTransformer2DModel(CachableDiT): + _supports_gradient_checkpointing = True + _no_split_modules = ["ZImageTransformerBlock"] + + def __init__( + self, + config: ZImageDitConfig, + hf_config: dict[str, Any], + ) -> None: + 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.all_patch_size = arch_config.all_patch_size + self.all_f_patch_size = arch_config.all_f_patch_size + self.dim = arch_config.dim + self.n_heads = arch_config.num_attention_heads + + self.rope_theta = arch_config.rope_theta + self.t_scale = arch_config.t_scale + self.gradient_checkpointing = False + + assert len(self.all_patch_size) == len(self.all_f_patch_size) + + all_x_embedder = {} + all_final_layer = {} + for patch_idx, (patch_size, f_patch_size) in enumerate( + zip(self.all_patch_size, self.all_f_patch_size) + ): + x_embedder = ReplicatedLinear( + f_patch_size * patch_size * patch_size * self.in_channels, + self.dim, + bias=True, + ) + all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder + + final_layer = FinalLayer( + self.dim, patch_size * patch_size * f_patch_size * self.out_channels + ) + all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer + + self.all_x_embedder = nn.ModuleDict(all_x_embedder) + self.all_final_layer = nn.ModuleDict(all_final_layer) + + self.noise_refiner = nn.ModuleList( + [ + ZImageTransformerBlock( + 1000 + layer_id, + self.dim, + self.n_heads, + arch_config.n_kv_heads, + arch_config.norm_eps, + arch_config.qk_norm, + modulation=True, + ) + for layer_id in range(arch_config.n_refiner_layers) + ] + ) + self.context_refiner = nn.ModuleList( + [ + ZImageTransformerBlock( + layer_id, + self.dim, + self.n_heads, + arch_config.n_kv_heads, + arch_config.norm_eps, + arch_config.qk_norm, + modulation=False, + ) + for layer_id in range(arch_config.n_refiner_layers) + ] + ) + self.t_embedder = TimestepEmbedder( + min(self.dim, ADALN_EMBED_DIM), mid_size=1024 + ) + + self.cap_embedder = nn.Sequential( + RMSNorm(arch_config.cap_feat_dim, eps=arch_config.norm_eps), + ReplicatedLinear(arch_config.cap_feat_dim, self.dim, bias=True), + ) + + self.x_pad_token = nn.Parameter(torch.empty((1, self.dim))) + self.cap_pad_token = nn.Parameter(torch.empty((1, self.dim))) + + self.layers = nn.ModuleList( + [ + ZImageTransformerBlock( + layer_id, + self.dim, + self.n_heads, + arch_config.n_kv_heads, + arch_config.norm_eps, + arch_config.qk_norm, + ) + for layer_id in range(arch_config.num_layers) + ] + ) + head_dim = self.dim // self.n_heads + assert head_dim == sum(arch_config.axes_dims) + self.axes_dims = arch_config.axes_dims + self.axes_lens = arch_config.axes_lens + + self.rope_embedder = RopeEmbedder( + theta=self.rope_theta, axes_dims=self.axes_dims, axes_lens=self.axes_lens + ) + + def unpatchify( + self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size + ) -> List[torch.Tensor]: + pH = pW = patch_size + pF = f_patch_size + bsz = len(x) + assert len(size) == bsz + for i in range(bsz): + F, H, W = size[i] + ori_len = (F // pF) * (H // pH) * (W // pW) + # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)" + x[i] = ( + x[i][:ori_len] + .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels) + .permute(6, 0, 3, 1, 4, 2, 5) + .reshape(self.out_channels, F, H, W) + ) + return x + + @staticmethod + def create_coordinate_grid(size, start=None, device=None): + if start is None: + start = (0 for _ in size) + + axes = [ + torch.arange(x0, x0 + span, dtype=torch.int32, device=device) + for x0, span in zip(start, size) + ] + grids = torch.meshgrid(axes, indexing="ij") + return torch.stack(grids, dim=-1) + + def patchify_and_embed( + self, + all_image: List[torch.Tensor], + all_cap_feats: List[torch.Tensor], + patch_size: int, + f_patch_size: int, + ): + pH = pW = patch_size + pF = f_patch_size + device = all_image[0].device + + all_image_out = [] + all_image_size = [] + all_image_pos_ids = [] + all_image_pad_mask = [] + all_cap_pos_ids = [] + all_cap_pad_mask = [] + all_cap_feats_out = [] + + for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)): + ### Process Caption + cap_ori_len = len(cap_feat) + cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF + # padded position ids + cap_padded_pos_ids = self.create_coordinate_grid( + size=(cap_ori_len + cap_padding_len, 1, 1), + start=(1, 0, 0), + device=device, + ).flatten(0, 2) + all_cap_pos_ids.append(cap_padded_pos_ids) + # pad mask + all_cap_pad_mask.append( + torch.cat( + [ + torch.zeros((cap_ori_len,), dtype=torch.bool, device=device), + torch.ones((cap_padding_len,), dtype=torch.bool, device=device), + ], + dim=0, + ) + ) + # padded feature + cap_padded_feat = torch.cat( + [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], + dim=0, + ) + all_cap_feats_out.append(cap_padded_feat) + + ### Process Image + C, F, H, W = image.size() + all_image_size.append((F, H, W)) + F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW + + image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) + # "c f pf h ph w pw -> (f h w) (pf ph pw c)" + image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape( + F_tokens * H_tokens * W_tokens, pF * pH * pW * C + ) + + image_ori_len = len(image) + image_padding_len = (-image_ori_len) % SEQ_MULTI_OF + + image_ori_pos_ids = self.create_coordinate_grid( + size=(F_tokens, H_tokens, W_tokens), + start=(cap_ori_len + cap_padding_len + 1, 0, 0), + device=device, + ).flatten(0, 2) + image_padding_pos_ids = ( + self.create_coordinate_grid( + size=(1, 1, 1), + start=(0, 0, 0), + device=device, + ) + .flatten(0, 2) + .repeat(image_padding_len, 1) + ) + image_padded_pos_ids = torch.cat( + [image_ori_pos_ids, image_padding_pos_ids], dim=0 + ) + all_image_pos_ids.append(image_padded_pos_ids) + # pad mask + all_image_pad_mask.append( + torch.cat( + [ + torch.zeros((image_ori_len,), dtype=torch.bool, device=device), + torch.ones( + (image_padding_len,), dtype=torch.bool, device=device + ), + ], + dim=0, + ) + ) + # padded feature + image_padded_feat = torch.cat( + [image, image[-1:].repeat(image_padding_len, 1)], dim=0 + ) + all_image_out.append(image_padded_feat) + + return ( + all_image_out, + all_cap_feats_out, + all_image_size, + all_image_pos_ids, + all_cap_pos_ids, + all_image_pad_mask, + all_cap_pad_mask, + ) + + def forward( + self, + hidden_states: List[torch.Tensor], + timestep, + encoder_hidden_states: List[torch.Tensor], + guidance=0, + patch_size=2, + f_patch_size=1, + **kwargs, + ): + assert patch_size in self.all_patch_size + assert f_patch_size in self.all_f_patch_size + + x = hidden_states + cap_feats = encoder_hidden_states + timestep = 1000.0 - timestep + t = timestep + bsz = len(x) + device = x[0].device + # t = t * self.t_scale + t = self.t_embedder(t) + + adaln_input = t.type_as(x) + ( + x, + cap_feats, + x_size, + x_pos_ids, + cap_pos_ids, + x_inner_pad_mask, + cap_inner_pad_mask, + ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) + # x embed & refine + x_item_seqlens = [len(_) for _ in x] + assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) + x_max_item_seqlen = max(x_item_seqlens) + + x = torch.cat(x, dim=0) + x, _ = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) + x[torch.cat(x_inner_pad_mask)] = self.x_pad_token + x = list(x.split(x_item_seqlens, dim=0)) + x_freqs_cis = list( + self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) + ) + + # RoPE returns (cos, sin) now + # x_pos_ids_cat = torch.cat(x_pos_ids, dim=0) + # x_cos, x_sin = self.rope_embedder(x_pos_ids_cat) + + # x_cos_list = list(x_cos.split(x_item_seqlens, dim=0)) + # x_sin_list = list(x_sin.split(x_item_seqlens, dim=0)) + + x = pad_sequence(x, batch_first=True, padding_value=0.0) + x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) + # x_cos = pad_sequence(x_cos_list, batch_first=True, padding_value=0.0) + # x_sin = pad_sequence(x_sin_list, batch_first=True, padding_value=0.0) + # B, T, D_half = x_cos.shape # D_half = 64 + + # x_cos_triton = x_cos.reshape(B * T, D_half).contiguous() # [B*T, 64] + # x_sin_triton = x_sin.reshape(B * T, D_half).contiguous() # [B*T, 64] + + x_attn_mask = torch.zeros( + (bsz, x_max_item_seqlen), dtype=torch.bool, device=device + ) + for i, seq_len in enumerate(x_item_seqlens): + x_attn_mask[i, :seq_len] = 1 + + # Refiner logic + for layer in self.noise_refiner: + x = layer(x, x_attn_mask, x_freqs_cis, adaln_input) + + # cap embed & refine + cap_item_seqlens = [len(_) for _ in cap_feats] + assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) + cap_max_item_seqlen = max(cap_item_seqlens) + + cap_feats = torch.cat(cap_feats, dim=0) + + # cap_embedder is Sequential with ReplicatedLinear. + # We need to handle this if ReplicatedLinear returns tuple. + # In __init__, cap_embedder = Sequential(RMSNorm, ReplicatedLinear). + # RMSNorm returns Tensor. ReplicatedLinear returns (Tensor, Gathered). + # Sequential returns (Tensor, Gathered). + # So we need to unpack. + cap_feats_out = self.cap_embedder(cap_feats) + if isinstance(cap_feats_out, tuple): + cap_feats = cap_feats_out[0] + else: + cap_feats = cap_feats_out + + cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token + cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) + + cap_freqs_cis = list( + self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split( + cap_item_seqlens, dim=0 + ) + ) + + cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0) + cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) + + cap_attn_mask = torch.zeros( + (bsz, cap_max_item_seqlen), dtype=torch.bool, device=device + ) + for i, seq_len in enumerate(cap_item_seqlens): + cap_attn_mask[i, :seq_len] = 1 + + for layer in self.context_refiner: + cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis) + + # unified + unified = [] + unified_freqs_cis = [] + + for i in range(bsz): + x_len = x_item_seqlens[i] + cap_len = cap_item_seqlens[i] + unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) + unified_freqs_cis.append( + torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]]) + ) + + unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] + assert unified_item_seqlens == [len(_) for _ in unified] + unified_max_item_seqlen = max(unified_item_seqlens) + + unified = pad_sequence(unified, batch_first=True, padding_value=0.0) + unified_freqs_cis = pad_sequence( + unified_freqs_cis, batch_first=True, padding_value=0.0 + ) + + unified_attn_mask = torch.zeros( + (bsz, unified_max_item_seqlen), dtype=torch.bool, device=device + ) + for i, seq_len in enumerate(unified_item_seqlens): + unified_attn_mask[i, :seq_len] = 1 + + for layer in self.layers: + unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input) + + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( + unified, adaln_input + ) + unified = list(unified.unbind(dim=0)) + x = self.unpatchify(unified, x_size, patch_size, f_patch_size) + + return -x[0] + + +EntryClass = ZImageTransformer2DModel diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py b/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py index 18f10046c..dae653eac 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py @@ -21,7 +21,7 @@ import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange -from transformers.modeling_utils import PretrainedConfig, PreTrainedModel +from transformers import PretrainedConfig, PreTrainedModel from sglang.multimodal_gen.runtime.models.dits.stepvideo import StepVideoRMSNorm diff --git a/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py new file mode 100644 index 000000000..f8fd441de --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py @@ -0,0 +1,116 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +# SPDX-License-Identifier: Apache-2.0 + + +from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline, Req +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + ConditioningStage, + DecodingStage, + DenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) +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, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +def prepare_mu(batch: Req, server_args: ServerArgs): + height = batch.height + width = batch.width + vae_scale_factor = server_args.pipeline_config.vae_config.vae_scale_factor + image_seq_len = ((int(height) // vae_scale_factor) // 2) * ( + (int(width) // vae_scale_factor) // 2 + ) + mu = calculate_shift( + image_seq_len, + # hard code, since scheduler_config is not in PipelineConfig now + 256, + 4096, + 0.5, + 1.15, + ) + return "mu", mu + + +class ZImagePipeline(LoRAPipeline, ComposedPipelineBase): + pipeline_name = "ZImagePipeline" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage_primary", + stage=TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + ], + tokenizers=[ + self.get_module("tokenizer"), + ], + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"), + prepare_extra_set_timesteps_kwargs=[prepare_mu], + ), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + 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 = ZImagePipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index a00ef3f8d..d06dc142b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py @@ -492,7 +492,7 @@ class CausalDMDDenoisingStage(DenoisingStage): result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) - result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) + result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) result.add_check("eta", batch.eta, V.non_negative_float) result.add_check("generator", batch.generator, V.generator_or_list_generators) result.add_check( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py index 284dac49e..f0cb4569a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py @@ -91,7 +91,7 @@ class ConditioningStage(PipelineStage): batch.do_classifier_free_guidance, V.bool_value, ) - result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) + result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) result.add_check( "negative_prompt_embeds", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 6f498efd1..8febe38dc 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -1353,7 +1353,7 @@ class DenoisingStage(PipelineStage): result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) - result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) + result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) result.add_check("eta", batch.eta, V.non_negative_float) result.add_check("generator", batch.generator, V.generator_or_list_generators) result.add_check( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index 3661a1890..892d34aed 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -194,7 +194,7 @@ class InputValidationStage(PipelineStage): ) # Validate guidance scale if using classifier-free guidance - if batch.do_classifier_free_guidance and batch.guidance_scale <= 0: + if batch.do_classifier_free_guidance and batch.guidance_scale < 0: raise ValueError( f"Guidance scale must be positive, but got {batch.guidance_scale}" ) @@ -318,7 +318,7 @@ class InputValidationStage(PipelineStage): result.add_check( "guidance_scale", batch.guidance_scale, - lambda x: not batch.do_classifier_free_guidance or V.positive_float(x), + lambda x: not batch.do_classifier_free_guidance or V.non_negative_float(x), ) return result diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 88d3872bf..5ba0dd56e 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -361,6 +361,31 @@ "expected_avg_denoise_ms": 167.89, "expected_median_denoise_ms": 169.67 }, + "zimage_image_t2i": { + "stages_ms": { + "InputValidationStage": 0.03, + "TextEncodingStage": 104.21, + "ConditioningStage": 0.01, + "TimestepPreparationStage": 1.33, + "LatentPreparationStage": 1.13, + "DenoisingStage": 850.85, + "DecodingStage": 289.32 + }, + "denoise_step_ms": { + "0": 101.56, + "1": 28.26, + "2": 101.74, + "3": 101.68, + "4": 102.19, + "5": 102.05, + "6": 102.03, + "7": 102.28, + "8": 105.54 + }, + "expected_e2e_ms": 1248.41, + "expected_avg_denoise_ms": 94.15, + "expected_median_denoise_ms": 102.03 + }, "qwen_image_edit_ti2i": { "notes": "single uploaded reference image, Qwen/Qwen-Image-Edit", "expected_e2e_ms": 138500.0, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index df705db93..3d7b96b54 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -255,6 +255,19 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [ output_size="1024x1024", ), ), + DiffusionTestCase( + "zimage_image_t2i", + DiffusionServerArgs( + model_path="Tongyi-MAI/Z-Image-Turbo", + modality="image", + warmup_text=1, + warmup_edit=0, + ), + DiffusionSamplingParams( + prompt="Doraemon is eating dorayaki.", + output_size="1024x1024", + ), + ), # === Text and Image to Image (TI2I) === # TODO: Timeout with Torch2.9. Add back when it can pass CI # DiffusionTestCase(