diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 7de0125d5..f065dd0dd 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -51,11 +51,11 @@ jobs: with: filters: | main_package: - - added|modified|deleted|renamed: "python/**" - - added|modified|deleted|renamed: "!python/sglang/multimodal_gen/**" - - added|modified|deleted|renamed: "scripts/ci/**" - - added|modified|deleted|renamed: "test/**" - - added|modified|deleted|renamed: ".github/workflows/pr-test.yml" + - "python/sglang/!(multimodal_gen)/**" + - "python/*.toml" + - "scripts/ci/**" + - "test/**" + - ".github/workflows/pr-test.yml" sgl_kernel: - "sgl-kernel/**" multimodal_gen: diff --git a/.github/workflows/vllm-dependency-test.yml b/.github/workflows/vllm-dependency-test.yml index 64fdc5cb2..a3b87bb5d 100644 --- a/.github/workflows/vllm-dependency-test.yml +++ b/.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" diff --git a/python/sglang/multimodal_gen/configs/pipelines/wan.py b/python/sglang/multimodal_gen/configs/pipelines/wan.py index af6a697c2..d98e5fe86 100644 --- a/python/sglang/multimodal_gen/configs/pipelines/wan.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py index 698e3cd9a..c0a589038 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py +++ b/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" diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 8a746fa08..2f3caa12f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py index 8ed4a0557..17aae037a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py +++ b/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 - z, mask2, seq_len = None, None, None + # TI2V specific preparations - BEFORE SP sharding + 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) - _, mask2 = masks_like([latent_model_input], zero=True) + # z: [B, C, 1, H, W] + 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, - "mask2": mask2, + "z": z_sp, # Use SP-sharded version + # 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 - if ( - server_args.pipeline_config.ti2v_task - and batch.pil_image is not None - ): - timestep = torch.stack([t_device]).to(get_local_torch_device()) - temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() - 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]) + timestep = self.expand_timestep_before_forward( + batch, + server_args, + t_device, + target_dtype, + seq_len, + reserved_frames_mask, + ) 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 - and batch.pil_image is not None - ): - latents = latents.squeeze(0) - latents = (1.0 - mask2[0]) * z + mask2[0] * latents + + latents = self.post_forward_for_ti2v_task( + batch, server_args, reserved_frames_mask, latents, z + ) # 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 ), diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index bb357b91f..d553069aa 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/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"], diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py b/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py index ea2d35304..b0e3e2a52 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py b/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py index 79d043f52..092725bdf 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py +++ b/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_cfg_parallel": 191.7 * 1.05, - "test_usp": 15.0, - "test_mixed": 15.0, + "test_usp": 530.5 * 1.05, + } + + +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_cfg_parallel": 191.7 * 1.05, - "test_usp": 387.6 * 1.05, - "test_mixed": 15.0, + "test_usp": 82.3 * 1.05, } +# 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() diff --git a/python/sglang/multimodal_gen/test/test_offline_api.py b/python/sglang/multimodal_gen/test/test_offline_api.py index 2e9ea67a3..9c45e2710 100644 --- a/python/sglang/multimodal_gen/test/test_offline_api.py +++ b/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 = [] diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 37f7418d8..b4d60e268 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/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" - self.__class__.results.append( - {"name": name, "key": test_key, "duration": duration_str, "status": status} - ) + duration = float(duration) if succeed else None + self.results.append(TestResult(name, test_key, duration, succeed)) 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", diff --git a/python/sglang/multimodal_gen/utils.py b/python/sglang/multimodal_gen/utils.py index 655af2c1e..79bf32f8e 100644 --- a/python/sglang/multimodal_gen/utils.py +++ b/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: