diffusion: fix wan-2.2-TI2V and support sp (#12926)

.github/workflows/pr-test.yml

@@ -51,11 +51,11 @@ jobs:
         with:
           filters: |
             main_package:
-              - added|modified|deleted|renamed: "python/**"
+              - "python/sglang/!(multimodal_gen)/**"
-              - added|modified|deleted|renamed: "!python/sglang/multimodal_gen/**"
+              - "python/*.toml"
-              - added|modified|deleted|renamed: "scripts/ci/**"
+              - "scripts/ci/**"
-              - added|modified|deleted|renamed: "test/**"
+              - "test/**"
-              - added|modified|deleted|renamed: ".github/workflows/pr-test.yml"
+              - ".github/workflows/pr-test.yml"
             sgl_kernel:
               - "sgl-kernel/**"
             multimodal_gen:

.github/workflows/vllm-dependency-test.yml

@@ -5,6 +5,7 @@ on:
     branches: [ main ]
     paths:
       - "python/**"
+      - "!python/sglang/multimodal_gen/**"
       - "scripts/ci/**"
       - "test/**"
       - ".github/workflows/vllm-dependency-test.yml"
@@ -12,6 +13,7 @@ on:
     branches: [ main ]
     paths:
       - "python/**"
+      - "!python/sglang/multimodal_gen/**"
       - "scripts/ci/**"
       - "test/**"
       - ".github/workflows/vllm-dependency-test.yml"

python/sglang/multimodal_gen/configs/pipelines/wan.py

@@ -140,8 +140,7 @@ class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
         vae_stride = self.vae_stride
         oh = batch.height
         ow = batch.width
-        shape = (z_dim, F, oh // vae_stride[1], ow // vae_stride[2])
+        shape = (batch_size, z_dim, F, oh // vae_stride[1], ow // vae_stride[2])
-
         return shape
 
     def __post_init__(self) -> None:

python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py

@@ -406,6 +406,9 @@ class NDRotaryEmbedding(torch.nn.Module):
         start_frame: int = 0,
         device: torch.device | str | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Handles sp internally
+        """
         # Caching wrapper: use grid parameters directly as the key.
         # grid_tuple = _to_tuple(grid_size, dim=self.ndim)
         device_str = str(device) if device is not None else "cpu"

python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py

@@ -690,7 +690,7 @@ class WanTransformer3DModel(CachableDiT):
         d = self.hidden_size // self.num_attention_heads
         self.rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
 
-        self.rope = NDRotaryEmbedding(
+        self.rotary_emb = NDRotaryEmbedding(
             rope_dim_list=self.rope_dim_list,
             rope_theta=10000,
             dtype=torch.float32 if current_platform.is_mps() else torch.float64,
@@ -725,7 +725,8 @@ class WanTransformer3DModel(CachableDiT):
         post_patch_height = height // p_h
         post_patch_width = width // p_w
 
-        freqs_cos, freqs_sin = self.rope.forward_from_grid(
+        # The rotary embedding layer correctly handles SP offsets internally.
+        freqs_cos, freqs_sin = self.rotary_emb.forward_from_grid(
             (
                 post_patch_num_frames * self.sp_size,
                 post_patch_height,
@@ -746,6 +747,7 @@ class WanTransformer3DModel(CachableDiT):
 
         # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
         if timestep.dim() == 2:
+            # ti2v
             ts_seq_len = timestep.shape[1]
             timestep = timestep.flatten()  # batch_size * seq_len
         else:

python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py

@@ -217,9 +217,6 @@ class DenoisingStage(PipelineStage):
             target_dtype != torch.float32
         ) and not server_args.disable_autocast
 
-        # Handle sequence parallelism if enabled
-        self._preprocess_sp_latents(batch)
-
         # Get timesteps and calculate warmup steps
         timesteps = batch.timesteps
         if timesteps is None:
@@ -263,8 +260,14 @@ class DenoisingStage(PipelineStage):
         else:
             boundary_timestep = None
 
-        # TI2V specific preparations
+        # TI2V specific preparations - BEFORE SP sharding
-        z, mask2, seq_len = None, None, None
+        z, z_sp, reserved_frames_masks, reserved_frames_mask_sp, seq_len = (
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
         # FIXME: should probably move to latent preparation stage, to handle with offload
         if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
             # Wan2.2 TI2V directly replaces the first frame of the latent with
@@ -285,11 +288,27 @@ class DenoisingStage(PipelineStage):
                 z = z * self.vae.scaling_factor.to(z.device, z.dtype)
             else:
                 z = z * self.vae.scaling_factor
-            latent_model_input = latents.to(target_dtype).squeeze(0)
+            # z: [B, C, 1, H, W]
-            _, mask2 = masks_like([latent_model_input], zero=True)
+            latent_model_input = latents.to(target_dtype)
+            # Keep as [B, C, T, H, W] for proper broadcasting
+            assert latent_model_input.ndim == 5
 
-            latents = (1.0 - mask2[0]) * z + mask2[0] * latent_model_input
+            # Create mask with proper shape [B, C, T, H, W]
+            latent_for_mask = latent_model_input.squeeze(0)  # [C, T, H, W]
+            _, reserved_frames_masks = masks_like([latent_for_mask], zero=True)
+            reserved_frames_mask = reserved_frames_masks[0].unsqueeze(
+                0
+            )  # [1, C, T, H, W]
+
+            # replace GLOBAL first frame with image - proper broadcasting
+            # z: [B, C, 1, H, W], reserved_frames_mask: [1, C, T, H, W]
+            # Both will broadcast correctly
+            latents = (
+                1.0 - reserved_frames_mask
+            ) * z + reserved_frames_mask * latent_model_input
+            assert latents.ndim == 5
             latents = latents.to(get_local_torch_device())
+            batch.latents = latents
 
             F = batch.num_frames
             temporal_scale = (
@@ -309,6 +328,74 @@ class DenoisingStage(PipelineStage):
                 int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size()
             )
 
+        # Handle sequence parallelism AFTER TI2V processing
+        self._preprocess_sp_latents(batch)
+        latents = batch.latents
+
+        # Shard z and reserved_frames_mask for TI2V if SP is enabled
+        if (
+            server_args.pipeline_config.ti2v_task
+            and batch.pil_image is not None
+            and get_sp_world_size() > 1
+        ):
+            sp_world_size = get_sp_world_size()
+            rank_in_sp_group = get_sp_parallel_rank()
+
+            if getattr(batch, "did_sp_shard_latents", False):
+                # Shard z (image latent) along time dimension
+                # z shape: [1, C, 1, H, W] - only first frame
+                # Only rank 0 has the first frame after sharding
+                if z.shape[2] == 1:
+                    # z is single frame, only rank 0 needs it
+                    if rank_in_sp_group == 0:
+                        z_sp = z
+                    else:
+                        # Other ranks don't have the first frame
+                        z_sp = None
+                else:
+                    # Should not happen for TI2V
+                    z_sp = z
+
+                # Shard reserved_frames_mask along time dimension to match sharded latents
+                # reserved_frames_mask is a list from masks_like, extract reserved_frames_mask[0] first
+                # reserved_frames_mask[0] shape: [C, T, H, W]
+                # All ranks need their portion of reserved_frames_mask for timestep calculation
+                if reserved_frames_masks is not None:
+                    reserved_frames_mask = reserved_frames_masks[
+                        0
+                    ]  # Extract tensor from list
+                    time_dim = reserved_frames_mask.shape[1]  # [C, T, H, W]
+                    if time_dim > 0 and time_dim % sp_world_size == 0:
+                        reserved_frames_mask_sp_tensor = rearrange(
+                            reserved_frames_mask,
+                            "c (n t) h w -> c n t h w",
+                            n=sp_world_size,
+                        ).contiguous()
+                        reserved_frames_mask_sp_tensor = reserved_frames_mask_sp_tensor[
+                            :, rank_in_sp_group, :, :, :
+                        ]
+                        reserved_frames_mask_sp = (
+                            reserved_frames_mask_sp_tensor  # Store as tensor, not list
+                        )
+                    else:
+                        reserved_frames_mask_sp = reserved_frames_mask
+                else:
+                    reserved_frames_mask_sp = None
+            else:
+                # SP not enabled or latents not sharded
+                z_sp = z
+                reserved_frames_mask_sp = (
+                    reserved_frames_masks[0]
+                    if reserved_frames_masks is not None
+                    else None
+                )  # Extract tensor
+        else:
+            # TI2V not enabled or SP not enabled
+            z_sp = z
+            reserved_frames_mask_sp = (
+                reserved_frames_masks[0] if reserved_frames_masks is not None else None
+            )  # Extract tensor
+
         guidance = self.get_or_build_guidance(
             # TODO: replace with raw_latent_shape?
             latents.shape[0],
@@ -370,8 +457,9 @@ class DenoisingStage(PipelineStage):
             "prompt_embeds": prompt_embeds,
             "neg_prompt_embeds": neg_prompt_embeds,
             "boundary_timestep": boundary_timestep,
-            "z": z,
+            "z": z_sp,  # Use SP-sharded version
-            "mask2": mask2,
+            # ndim == 5
+            "reserved_frames_mask": reserved_frames_mask_sp,  # Use SP-sharded version
             "seq_len": seq_len,
             "guidance": guidance,
         }
@@ -582,6 +670,75 @@ class DenoisingStage(PipelineStage):
         assert current_model is not None, "The model for the current step is not set."
         return current_model, current_guidance_scale
 
+    def expand_timestep_before_forward(
+        self,
+        batch: Req,
+        server_args: ServerArgs,
+        t_device,
+        target_dtype,
+        seq_len,
+        reserved_frames_mask,
+    ):
+        bsz = batch.raw_latent_shape[0]
+        # expand timestep
+        if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+            # Explicitly cast t_device to the target float type at the beginning.
+            # This ensures any precision-based rounding (e.g., float32(999.0) -> bfloat16(1000.0))
+            # is applied consistently *before* it's used by any rank.
+            t_device_rounded = t_device.to(target_dtype)
+
+            local_seq_len = seq_len
+            if get_sp_world_size() > 1 and getattr(
+                batch, "did_sp_shard_latents", False
+            ):
+                local_seq_len = seq_len // get_sp_world_size()
+
+            if get_sp_parallel_rank() == 0 and reserved_frames_mask is not None:
+                # Rank 0 has the first frame, create a special timestep tensor
+                # NOTE: The spatial downsampling in the next line is suspicious but kept
+                # to match original model's potential training configuration.
+                temp_ts = (
+                    reserved_frames_mask[0][:, ::2, ::2] * t_device_rounded
+                ).flatten()
+
+                # Pad to full local sequence length
+                temp_ts = torch.cat(
+                    [
+                        temp_ts,
+                        temp_ts.new_ones(local_seq_len - temp_ts.size(0))
+                        * t_device_rounded,
+                    ]
+                )
+                timestep = temp_ts.unsqueeze(0).repeat(bsz, 1)
+            else:
+                # Other ranks get a uniform timestep tensor of the correct shape [B, local_seq_len]
+                timestep = t_device.repeat(bsz, local_seq_len)
+        else:
+            timestep = t_device.repeat(bsz)
+        return timestep
+
+    def post_forward_for_ti2v_task(
+        self, batch: Req, server_args: ServerArgs, reserved_frames_mask, latents, z
+    ):
+        """
+        For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep
+        """
+        if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+            # Apply TI2V mask blending with SP-aware z and reserved_frames_mask.
+            # This ensures the first frame is always the condition image after each step.
+            # This is only applied on rank 0, where z is not None.
+            if z is not None and reserved_frames_mask is not None:
+                # z: [1, C, 1, H, W]
+                # latents: [1, C, T_local, H, W]
+                # reserved_frames_mask: [C, T_local, H, W]
+                # Unsqueeze mask to [1, C, T_local, H, W] for broadcasting.
+                # z will broadcast along the time dimension.
+                latents = (
+                    1.0 - reserved_frames_mask.unsqueeze(0)
+                ) * z + reserved_frames_mask.unsqueeze(0) * latents
+
+        return latents
+
     @torch.no_grad()
     def forward(
         self,
@@ -613,7 +770,7 @@ class DenoisingStage(PipelineStage):
         latents = prepared_vars["latents"]
         boundary_timestep = prepared_vars["boundary_timestep"]
         z = prepared_vars["z"]
-        mask2 = prepared_vars["mask2"]
+        reserved_frames_mask = prepared_vars["reserved_frames_mask"]
         seq_len = prepared_vars["seq_len"]
         guidance = prepared_vars["guidance"]
 
@@ -663,23 +820,14 @@ class DenoisingStage(PipelineStage):
                             [latent_model_input, batch.image_latent], dim=1
                         ).to(target_dtype)
 
-                    # expand timestep
+                    timestep = self.expand_timestep_before_forward(
-                    if (
+                        batch,
-                        server_args.pipeline_config.ti2v_task
+                        server_args,
-                        and batch.pil_image is not None
+                        t_device,
-                    ):
+                        target_dtype,
-                        timestep = torch.stack([t_device]).to(get_local_torch_device())
+                        seq_len,
-                        temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
+                        reserved_frames_mask,
-                        temp_ts = torch.cat(
+                    )
-                            [
-                                temp_ts,
-                                temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep,
-                            ]
-                        )
-                        timestep = temp_ts.unsqueeze(0)
-                        t_expand = timestep.repeat(latent_model_input.shape[0], 1)
-                    else:
-                        t_expand = t_device.repeat(latent_model_input.shape[0])
 
                     latent_model_input = self.scheduler.scale_model_input(
                         latent_model_input, t_device
@@ -690,7 +838,7 @@ class DenoisingStage(PipelineStage):
                     noise_pred = self._predict_noise_with_cfg(
                         current_model,
                         latent_model_input,
-                        t_expand,
+                        timestep,
                         batch,
                         i,
                         attn_metadata,
@@ -714,12 +862,10 @@ class DenoisingStage(PipelineStage):
                         **extra_step_kwargs,
                         return_dict=False,
                     )[0]
-                    if (
+
-                        server_args.pipeline_config.ti2v_task
+                    latents = self.post_forward_for_ti2v_task(
-                        and batch.pil_image is not None
+                        batch, server_args, reserved_frames_mask, latents, z
-                    ):
+                    )
-                        latents = latents.squeeze(0)
-                        latents = (1.0 - mask2[0]) * z + mask2[0] * latents
 
                     # save trajectory latents if needed
                     if batch.return_trajectory_latents:
@@ -876,7 +1022,7 @@ class DenoisingStage(PipelineStage):
         self,
         current_model,
         latent_model_input,
-        t_expand,
+        timestep,
         prompt_embeds,
         target_dtype,
         guidance: torch.Tensor,
@@ -885,7 +1031,7 @@ class DenoisingStage(PipelineStage):
         return current_model(
             hidden_states=latent_model_input,
             encoder_hidden_states=prompt_embeds,
-            timestep=t_expand,
+            timestep=timestep,
             guidance=guidance,
             **kwargs,
         )
@@ -894,7 +1040,7 @@ class DenoisingStage(PipelineStage):
         self,
         current_model: torch.nn.Module,
         latent_model_input: torch.Tensor,
-        t_expand,
+        timestep,
         batch,
         timestep_index: int,
         attn_metadata,
@@ -913,7 +1059,7 @@ class DenoisingStage(PipelineStage):
         Args:
             current_model: The transformer model to use for the current step.
             latent_model_input: The input latents for the model.
-            t_expand: The expanded timestep tensor.
+            timestep: The expanded timestep tensor.
             batch: The current batch information.
             timestep_index: The current timestep index.
             attn_metadata: Attention metadata for custom backends.
@@ -940,7 +1086,7 @@ class DenoisingStage(PipelineStage):
                 noise_pred_cond = self._predict_noise(
                     current_model=current_model,
                     latent_model_input=latent_model_input,
-                    t_expand=t_expand,
+                    timestep=timestep,
                     prompt_embeds=server_args.pipeline_config.get_pos_prompt_embeds(
                         batch
                     ),
@@ -968,7 +1114,7 @@ class DenoisingStage(PipelineStage):
                 noise_pred_uncond = self._predict_noise(
                     current_model=current_model,
                     latent_model_input=latent_model_input,
-                    t_expand=t_expand,
+                    timestep=timestep,
                     prompt_embeds=server_args.pipeline_config.get_neg_prompt_embeds(
                         batch
                     ),

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -370,7 +370,7 @@ def maybe_download_model(
         logger.info(
             "Downloading model snapshot from HF Hub for %s...", model_name_or_path
         )
-        with get_lock(model_name_or_path):
+        with get_lock(model_name_or_path).acquire(poll_interval=2):
             local_path = snapshot_download(
                 repo_id=model_name_or_path,
                 ignore_patterns=["*.onnx", "*.msgpack"],

python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py

@@ -20,6 +20,10 @@ class TestFastWan2_1_T2V(TestGenerateBase):
         "test_mixed": 15.0,
     }
 
+    # disabled for vsa
+    def test_usp(self):
+        pass
+
 
 class TestFastWan2_2_T2V(TestGenerateBase):
     model_path = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
@@ -38,7 +42,7 @@ class TestWan2_1_T2V(TestGenerateBase):
     extra_args = []
     data_type: DataType = DataType.VIDEO
     thresholds = {
-        "test_single_gpu": 76.0,
+        "test_single_gpu": 76.0 * 1.05,
         "test_cfg_parallel": 46.5 * 1.05,
         "test_usp": 22.5,
         "test_mixed": 26.5,

python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py

@@ -1,5 +1,3 @@
-# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
-
 import unittest
 
 from sglang.multimodal_gen.configs.sample.base import DataType
@@ -18,7 +16,8 @@ class TestGenerateTI2VBase(TestGenerateBase):
             "sglang",
             "generate",
             f'--prompt="Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline\'s intricate details and the refreshing atmosphere of the seaside."',
-            "--image-path=https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true",
+            "--image-path",
+            "https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true",
             "--save-output",
             "--log-level=debug",
             f"--output-path={cls.output_path}",
@@ -36,12 +35,15 @@ class TestGenerateTI2VBase(TestGenerateBase):
 
 class TestWan2_1_I2V_14B_480P(TestGenerateTI2VBase):
     model_path = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-    extra_args = ["--attention-backend=video_sparse_attn"]
     thresholds = {
-        "test_single_gpu": 13.0,
+        "test_usp": 530.5 * 1.05,
-        "test_cfg_parallel": 191.7 * 1.05,
+    }
-        "test_usp": 15.0,
+
-        "test_mixed": 15.0,
+
+class TestWan2_1_I2V_14B_720P(TestGenerateTI2VBase):
+    model_path = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
+    thresholds = {
+        "test_usp": 530.5 * 1.05,
     }
 
 
@@ -50,13 +52,19 @@ class TestWan2_2_TI2V_5B(TestGenerateTI2VBase):
     # FIXME: doesn't work with vsa at the moment
     # extra_args = ["--attention-backend=video_sparse_attn"]
     thresholds = {
-        "test_single_gpu": 13.0,
+        "test_usp": 82.3 * 1.05,
-        "test_cfg_parallel": 191.7 * 1.05,
-        "test_usp": 387.6 * 1.05,
-        "test_mixed": 15.0,
     }
 
 
+# OOM
+# class TestWan2_2_I2V_A14B(TestGenerateTI2VBase):
+#     model_path = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+#     # FIXME: doesn't work with vsa at the moment
+#     thresholds = {
+#         "test_usp": 66.3 * 1.05,
+#     }
+
+
 if __name__ == "__main__":
     del TestGenerateTI2VBase, TestGenerateBase
     unittest.main()

python/sglang/multimodal_gen/test/test_offline_api.py

@@ -19,7 +19,7 @@ class TestGeneratorAPIBase(unittest.TestCase):
     server_kwargs = {}
 
     # sampling
-    output_path: str = "outputs"
+    output_path: str = "test_outputs"
 
     results = []
 

python/sglang/multimodal_gen/test/test_utils.py

@@ -1,4 +1,5 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+import dataclasses
 import os
 import shlex
 import socket
@@ -6,6 +7,7 @@ import subprocess
 import sys
 import time
 import unittest
+from typing import Optional
 
 from PIL import Image
 
@@ -15,7 +17,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 
-def run_command(command):
+def run_command(command) -> Optional[float]:
     """Runs a command and returns the execution time and status."""
     print(f"Running command: {' '.join(command)}")
 
@@ -75,6 +77,18 @@ def check_image_size(ut, image, width, height):
     ut.assertEqual(image.size, (width, height))
 
 
+@dataclasses.dataclass
+class TestResult:
+    name: str
+    key: str
+    duration: Optional[float]
+    succeed: bool
+
+    @property
+    def duration_str(self):
+        return f"{self.duration:.4f}" if self.duration else "NA"
+
+
 class TestCLIBase(unittest.TestCase):
     model_path: str = None
     extra_args = []
@@ -84,7 +98,7 @@ class TestCLIBase(unittest.TestCase):
 
     width: int = 720
     height: int = 720
-    output_path: str = "outputs"
+    output_path: str = "test_outputs"
 
     base_command = [
         "sglang",
@@ -105,7 +119,7 @@ class TestCLIBase(unittest.TestCase):
     def setUpClass(cls):
         cls.results = []
 
-    def _run_command(self, name, model_path: str, test_key: str = "", args=[]):
+    def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]):
         command = (
             self.base_command
             + [f"--model-path={model_path}"]
@@ -115,11 +129,10 @@ class TestCLIBase(unittest.TestCase):
         )
         duration = run_command(command)
         status = "Success" if duration else "Failed"
+        succeed = duration is not None
 
-        duration_str = f"{duration:.4f}s" if duration else "NA"
+        duration = float(duration) if succeed else None
-        self.__class__.results.append(
+        self.results.append(TestResult(name, test_key, duration, succeed))
-            {"name": name, "key": test_key, "duration": duration_str, "status": status}
-        )
 
         return name, duration, status
 
@@ -133,7 +146,7 @@ class TestGenerateBase(TestCLIBase):
 
     width: int = 720
     height: int = 720
-    output_path: str = "outputs"
+    output_path: str = "test_outputs"
     image_path: str | None = None
     prompt: str | None = "A curious raccoon"
 
@@ -150,7 +163,7 @@ class TestGenerateBase(TestCLIBase):
         f"--output-path={output_path}",
     ]
 
-    results = []
+    results: list[TestResult] = []
 
     @classmethod
     def setUpClass(cls):
@@ -167,24 +180,28 @@ class TestGenerateBase(TestCLIBase):
             test_key: order for order, test_key in enumerate(test_keys)
         }
 
-        ordered_results: list[dict] = [{}] * len(test_keys)
+        ordered_results: list[TestResult] = [None] * len(test_keys)
-
         for result in cls.results:
-            order = test_key_to_order[result["key"]]
+            order = test_key_to_order[result.key]
             ordered_results[order] = result
 
         for result in ordered_results:
             if not result:
                 continue
             status = (
-                result["status"] and result["duration"] <= cls.thresholds[result["key"]]
+                "Succeed"
+                if (
+                    result.succeed
+                    and float(result.duration) <= float(cls.thresholds[result.key])
+                )
+                else "Failed"
             )
-            print(f"| {result['name']:<30} | {result['duration']:<8} | {status:<7} |")
+            print(f"| {result.name:<30} | {result.duration_str:<8} | {status:<7} |")
         print()
-        durations = [result["duration"] for result in cls.results]
+        durations = [result.duration_str for result in cls.results]
         print(" | ".join([""] + durations + [""]))
 
-    def _run_test(self, name, args, model_path: str, test_key: str):
+    def _run_test(self, name: str, args, model_path: str, test_key: str):
         time_threshold = self.thresholds[test_key]
         name, duration, status = self._run_command(
             name, args=args, model_path=model_path, test_key=test_key
@@ -220,7 +237,7 @@ class TestGenerateBase(TestCLIBase):
     def test_single_gpu(self):
         """single gpu"""
         self._run_test(
-            name=f"{self.model_name()}, single gpu",
+            name=f"{self.model_name()}_single gpu",
             args=None,
             model_path=self.model_path,
             test_key="test_single_gpu",
@@ -231,7 +248,7 @@ class TestGenerateBase(TestCLIBase):
         if self.data_type == DataType.IMAGE:
             return
         self._run_test(
-            name=f"{self.model_name()}, cfg parallel",
+            name=f"{self.model_name()}_cfg parallel",
             args="--num-gpus 2 --enable-cfg-parallel",
             model_path=self.model_path,
             test_key="test_cfg_parallel",
@@ -242,7 +259,7 @@ class TestGenerateBase(TestCLIBase):
         if self.data_type == DataType.IMAGE:
             return
         self._run_test(
-            name=f"{self.model_name()}, usp",
+            name=f"{self.model_name()}_usp",
             args="--num-gpus 4 --ulysses-degree=2 --ring-degree=2",
             model_path=self.model_path,
             test_key="test_usp",
@@ -253,7 +270,7 @@ class TestGenerateBase(TestCLIBase):
         if self.data_type == DataType.IMAGE:
             return
         self._run_test(
-            name=f"{self.model_name()}, mixed",
+            name=f"{self.model_name()}_mixed",
             args="--num-gpus 4 --ulysses-degree=2 --ring-degree=1 --enable-cfg-parallel",
             model_path=self.model_path,
             test_key="test_mixed",

python/sglang/multimodal_gen/utils.py

@@ -698,12 +698,38 @@ def is_vmoba_available() -> bool:
 
 # adapted from: https://github.com/Wan-Video/Wan2.2/blob/main/wan/utils/utils.py
 def masks_like(
-    tensor, zero=False, generator=None, p=0.2
+    tensors, zero=False, generator=None, p=0.2
 ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
-    assert isinstance(tensor, list)
+    """
-    out1 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensor]
+    Generate binary masks for Text-to-Image-to-Video (TI2V) tasks.
 
-    out2 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensor]
+    Creates masks to control which frames should be preserved vs replaced.
+    Primarily used to fix the first frame to the input image while generating other frames.
+
+    Args:
+        tensors: List of tensors with shape [C, T, H, W]
+        zero: If True, set first frame (dim 1, index 0) to zero. Default: False
+        generator: Optional random generator for stochastic masking
+        p: Probability of applying special noise when generator is provided. Default: 0.2
+
+    Returns:
+        Tuple of two lists of tensors:
+        - When zero=False: Both lists contain all-ones tensors
+        - When zero=True (no generator): First frame set to 0, others to 1
+        - When zero=True (with generator): First frame set to small random values with probability p
+
+    Example:
+        >>> latent = torch.randn(48, 69, 96, 160)  # [C, T, H, W]
+        >>> _, mask = masks_like([latent], zero=True)
+        >>> # mask[0][:, 0] == 0 (first frame)
+        >>> # mask[0][:, 1:] == 1 (other frames)
+        >>> blended = (1.0 - mask[0]) * image + mask[0] * latent
+        >>> # Result: first frame = image, other frames = latent
+    """
+    assert isinstance(tensors, list)
+    out1 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensors]
+
+    out2 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensors]
 
     if zero:
         if generator is not None: