[diffusion] model: support z-image (#14067)

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"

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",
 ]

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)
 

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)

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,
+            ],
+        )
+    )

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(

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

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
 

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

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(

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

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(

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
 

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,

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(