From 9530b766301c1e456e59935ef64be2ef828600e7 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 2 Dec 2025 18:59:40 +0800 Subject: [PATCH] [diffusion] refactor: simplify DmdDenoisingStage (#14269) --- .../runtime/entrypoints/openai/image_api.py | 7 +- .../runtime/entrypoints/openai/video_api.py | 5 +- .../pipelines_core/stages/denoising.py | 276 +++++++++--------- .../pipelines_core/stages/denoising_dmd.py | 196 ++++--------- .../pipelines_core/stages/input_validation.py | 24 +- .../sglang/multimodal_gen/test/slack_utils.py | 9 +- 6 files changed, 224 insertions(+), 293 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 3a0411cd6..7ac86c2c4 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -54,7 +54,10 @@ def _build_sampling_params_from_request( background: Optional[str], image_path: Optional[str] = None, ) -> SamplingParams: - width, height = _parse_size(size) + if size is None: + width, height = None, None + else: + width, height = _parse_size(size) ext = _choose_ext(output_format, background) server_args = get_global_server_args() # Build user params @@ -149,7 +152,7 @@ async def edits( model: Optional[str] = Form(None), n: Optional[int] = Form(1), response_format: Optional[str] = Form(None), - size: Optional[str] = Form("1024x1024"), + size: Optional[str] = Form(None), output_format: Optional[str] = Form(None), background: Optional[str] = Form("auto"), user: Optional[str] = Form(None), diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index a1e20a8dd..c346fc528 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -47,7 +47,10 @@ router = APIRouter(prefix="/v1/videos", tags=["videos"]) def _build_sampling_params_from_request( request_id: str, request: VideoGenerationsRequest ) -> SamplingParams: - width, height = _parse_size(request.size or "720x1280") + if request.size is None: + width, height = None, None + else: + width, height = _parse_size(request.size) seconds = request.seconds if request.seconds is not None else 4 # Prefer user-provided fps/num_frames from request; fallback to defaults fps_default = 24 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index e74c755f1..aae4af503 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -20,6 +20,7 @@ from einops import rearrange from tqdm.auto import tqdm from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode +from sglang.multimodal_gen.configs.pipeline_configs.wan import Wan2_2_TI2V_5B_Config from sglang.multimodal_gen.runtime.distributed import ( cfg_model_parallel_all_reduce, get_local_torch_device, @@ -184,6 +185,118 @@ class DenoisingStage(PipelineStage): # return StageParallelismType.CFG_PARALLEL if get_global_server_args().enable_cfg_parallel else StageParallelismType.REPLICATED return StageParallelismType.REPLICATED + def _preprocess_latents_for_ti2v( + self, latents, target_dtype, batch, server_args: ServerArgs + ): + # FIXME: should probably move to latent preparation stage, to handle with offload + # Wan2.2 TI2V directly replaces the first frame of the latent with + # the image latent instead of appending along the channel dim + assert batch.image_latent is None, "TI2V task should not have image latents" + assert self.vae is not None, "VAE is not provided for TI2V task" + self.vae = self.vae.to(batch.condition_image.device) + z = self.vae.encode(batch.condition_image).mean.float() + if self.vae.device != "cpu" and server_args.vae_cpu_offload: + self.vae = self.vae.to("cpu") + if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None: + if isinstance(self.vae.shift_factor, torch.Tensor): + z -= self.vae.shift_factor.to(z.device, z.dtype) + else: + z -= self.vae.shift_factor + + if isinstance(self.vae.scaling_factor, torch.Tensor): + z = z * self.vae.scaling_factor.to(z.device, z.dtype) + else: + z = z * self.vae.scaling_factor + # 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 + + # 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 = ( + server_args.pipeline_config.vae_config.arch_config.scale_factor_temporal + ) + spatial_scale = ( + server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial + ) + patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size + seq_len = ( + ((F - 1) // temporal_scale + 1) + * (batch.height // spatial_scale) + * (batch.width // spatial_scale) + // (patch_size[1] * patch_size[2]) + ) + seq_len = int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size() + return seq_len, z, reserved_frames_masks + + def _postprocess_latents_for_ti2v(self, z, reserved_frames_masks, batch): + rank_in_sp_group = get_sp_parallel_rank() + sp_world_size = get_sp_world_size() + + 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 + + return reserved_frames_mask_sp, z_sp + def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): """ Prepare all necessary invariant variables for the denoising loop. @@ -264,144 +377,38 @@ class DenoisingStage(PipelineStage): else: boundary_timestep = 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 ( + # specifically for Wan2_2_TI2V_5B_Config, not applicable for FastWan2_2_TI2V_5B_Config + should_preprocess_for_wan_ti2v = ( server_args.pipeline_config.task_type == ModelTaskType.TI2V and batch.condition_image is not None - ): - # Wan2.2 TI2V directly replaces the first frame of the latent with - # the image latent instead of appending along the channel dim - assert batch.image_latent is None, "TI2V task should not have image latents" - assert self.vae is not None, "VAE is not provided for TI2V task" - self.vae = self.vae.to(batch.condition_image.device) - z = self.vae.encode(batch.condition_image).mean.float() - if self.vae.device != "cpu" and server_args.vae_cpu_offload: - self.vae = self.vae.to("cpu") - if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None: - if isinstance(self.vae.shift_factor, torch.Tensor): - z -= self.vae.shift_factor.to(z.device, z.dtype) - else: - z -= self.vae.shift_factor + and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config + ) - if isinstance(self.vae.scaling_factor, torch.Tensor): - z = z * self.vae.scaling_factor.to(z.device, z.dtype) - else: - z = z * self.vae.scaling_factor - # 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 - - # 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 = ( - server_args.pipeline_config.vae_config.arch_config.scale_factor_temporal + # TI2V specific preparations - before SP sharding + if should_preprocess_for_wan_ti2v: + seq_len, z, reserved_frames_masks = self._preprocess_latents_for_ti2v( + latents, target_dtype, batch, server_args ) - spatial_scale = ( - server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial - ) - patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size - seq_len = ( - ((F - 1) // temporal_scale + 1) - * (batch.height // spatial_scale) - * (batch.width // spatial_scale) - // (patch_size[1] * patch_size[2]) - ) - seq_len = ( - int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size() + else: + seq_len, z, reserved_frames_masks = ( + None, + None, + None, ) - # Handle sequence parallelism AFTER TI2V processing + # Handle sequence parallelism after TI2V processing self._preprocess_sp_latents(batch, server_args) latents = batch.latents # Shard z and reserved_frames_mask for TI2V if SP is enabled - if ( - server_args.pipeline_config.task_type == ModelTaskType.TI2V - and batch.condition_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 + if should_preprocess_for_wan_ti2v: + reserved_frames_mask_sp, z_sp = self._postprocess_latents_for_ti2v( + z, reserved_frames_masks, batch + ) else: - # TI2V not enabled or SP not enabled - z_sp = z - reserved_frames_mask_sp = ( + reserved_frames_mask_sp, z_sp = ( reserved_frames_masks[0] if reserved_frames_masks is not None else None - ) # Extract tensor + ), z guidance = self.get_or_build_guidance( # TODO: replace with raw_latent_shape? @@ -682,15 +689,18 @@ class DenoisingStage(PipelineStage): server_args: ServerArgs, t_device, target_dtype, - seq_len, + seq_len: int | None, reserved_frames_mask, ): bsz = batch.raw_latent_shape[0] - # expand timestep - if ( + should_preprocess_for_wan_ti2v = ( server_args.pipeline_config.task_type == ModelTaskType.TI2V and batch.condition_image is not None - ): + and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config + ) + + # expand timestep + if should_preprocess_for_wan_ti2v: # 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. @@ -732,10 +742,12 @@ class DenoisingStage(PipelineStage): """ For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep """ - if ( + should_preprocess_for_wan_ti2v = ( server_args.pipeline_config.task_type == ModelTaskType.TI2V and batch.condition_image is not None - ): + and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config + ) + if should_preprocess_for_wan_ti2v: # 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. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 1af44795d..1ee183b0a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -3,20 +3,8 @@ import time import torch -from einops import rearrange -from sglang.multimodal_gen.runtime.distributed import ( - get_local_torch_device, - get_sp_parallel_rank, - get_sp_world_size, - sequence_model_parallel_all_gather, -) -from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( - SlidingTileAttentionBackend, -) -from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import ( - VideoSparseAttentionBackend, -) +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, @@ -24,10 +12,6 @@ from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler from sglang.multimodal_gen.runtime.models.utils import pred_noise_to_pred_video from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage -from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( - st_attn_available, - vsa_available, -) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler @@ -36,7 +20,6 @@ from sglang.multimodal_gen.utils import dict_to_3d_list logger = init_logger(__name__) -# TODO: use base methods of DenoisingStage class DmdDenoisingStage(DenoisingStage): """ Denoising stage for DMD. @@ -46,6 +29,29 @@ class DmdDenoisingStage(DenoisingStage): super().__init__(transformer, scheduler) self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0) + def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs): + # 1. to shard latents (B, C, T, H, W) along dim 2 + super()._preprocess_sp_latents(batch, server_args) + + # 2. DMD expects (B, T, C, H, W) for the main latents in the loop + if batch.latents is not None: + batch.latents = batch.latents.permute(0, 2, 1, 3, 4) + + # Note: batch.image_latent is kept as (B, C, T, H, W) here + + def _postprocess_sp_latents( + self, + batch: Req, + latents: torch.Tensor, + trajectory_tensor: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # 1. convert back from DMD's (B, T, C, H, W) to standard (B, C, T, H, W) + # this is because base gather_latents_for_sp expects dim=2 for T + latents = latents.permute(0, 2, 1, 3, 4) + + # 2. use base method to gather + return super()._postprocess_sp_latents(batch, latents, trajectory_tensor) + def forward( self, batch: Req, @@ -53,38 +59,25 @@ class DmdDenoisingStage(DenoisingStage): ) -> Req: """ Run the denoising loop. - - Args: - batch: The current batch information. - server_args: The inference arguments. - - Returns: - The batch with denoised latents. """ - # Setup precision and autocast settings - # TODO(will): make the precision configurable for inference - # target_dtype = PRECISION_TO_TYPE[server_args.precision] - target_dtype = torch.bfloat16 - autocast_enabled = ( - target_dtype != torch.float32 - ) and not server_args.disable_autocast + prepared_vars = self._prepare_denoising_loop(batch, server_args) - # Get timesteps and calculate warmup steps - timesteps = batch.timesteps + target_dtype = prepared_vars["target_dtype"] + autocast_enabled = prepared_vars["autocast_enabled"] + num_warmup_steps = prepared_vars["num_warmup_steps"] + latents = prepared_vars["latents"] + video_raw_latent_shape = latents.shape - # TODO(will): remove this once we add input/output validation for stages - if timesteps is None: - raise ValueError("Timesteps must be provided") - num_inference_steps = batch.num_inference_steps - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + timesteps = torch.tensor( + server_args.pipeline_config.dmd_denoising_steps, + dtype=torch.long, + device=get_local_torch_device(), + ) - # Prepare image latents and embeddings for I2V generation + # prepare image_kwargs image_embeds = batch.image_embeds if len(image_embeds) > 0: - assert torch.isnan(image_embeds[0]).sum() == 0 - image_embeds = [ - image_embed.to(target_dtype) for image_embed in image_embeds - ] + image_embeds = [img.to(target_dtype) for img in image_embeds] image_kwargs = self.prepare_extra_func_kwargs( self.transformer.forward, @@ -94,56 +87,12 @@ class DmdDenoisingStage(DenoisingStage): }, ) - pos_cond_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - "encoder_hidden_states_2": batch.clip_embedding_pos, - "encoder_attention_mask": batch.prompt_attention_mask, - }, - ) + pos_cond_kwargs = prepared_vars["pos_cond_kwargs"] + prompt_embeds = prepared_vars["prompt_embeds"] - # Prepare STA parameters - if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: - self.prepare_sta_param(batch, server_args) - - # Get latents and embeddings - assert batch.latents is not None, "latents must be provided" - latents = batch.latents - latents = latents.permute(0, 2, 1, 3, 4) - - video_raw_latent_shape = latents.shape - prompt_embeds = batch.prompt_embeds - assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan" - timesteps = torch.tensor( - server_args.pipeline_config.dmd_denoising_steps, - dtype=torch.long, - device=get_local_torch_device(), - ) - - # Handle sequence parallelism if enabled - sp_world_size, rank_in_sp_group = ( - get_sp_world_size(), - get_sp_parallel_rank(), - ) - sp_group = sp_world_size > 1 - if sp_group: - latents = rearrange( - latents, "b (n t) c h w -> b n t c h w", n=sp_world_size - ).contiguous() - latents = latents[:, rank_in_sp_group, :, :, :, :] - if batch.image_latent is not None: - image_latent = rearrange( - batch.image_latent, - "b c (n t) h w -> b c n t h w", - n=sp_world_size, - ).contiguous() - - image_latent = image_latent[:, :, rank_in_sp_group, :, :, :] - batch.image_latent = image_latent - - # Run denoising loop denoising_loop_start_time = time.time() self.start_profile(batch=batch) + with self.progress_bar(total=len(timesteps)) as progress_bar: for i, t in enumerate(timesteps): # Skip if interrupted @@ -171,16 +120,11 @@ class DmdDenoisingStage(DenoisingStage): # Prepare inputs for transformer t_expand = t.repeat(latent_model_input.shape[0]) - guidance_expand = ( - torch.tensor( - [server_args.pipeline_config.embedded_cfg_scale] - * latent_model_input.shape[0], - dtype=torch.float32, - device=get_local_torch_device(), - ).to(target_dtype) - * 1000.0 - if server_args.pipeline_config.embedded_cfg_scale is not None - else None + + guidance_expand = self.get_or_build_guidance( + latent_model_input.shape[0], + target_dtype, + get_local_torch_device(), ) # Predict noise residual @@ -189,41 +133,13 @@ class DmdDenoisingStage(DenoisingStage): dtype=target_dtype, enabled=autocast_enabled, ): - if ( - vsa_available - and self.attn_backend == VideoSparseAttentionBackend - ): - self.attn_metadata_builder_cls = ( - self.attn_backend.get_builder_cls() - ) - - if self.attn_metadata_builder_cls is not None: - self.attn_metadata_builder = ( - self.attn_metadata_builder_cls() - ) - # TODO(will): clean this up - attn_metadata = self.attn_metadata_builder.build( # type: ignore - current_timestep=i, # type: ignore - raw_latent_shape=batch.raw_latent_shape[2:5], # type: ignore - patch_size=server_args.pipeline_config.dit_config.patch_size, # type: ignore - STA_param=batch.STA_param, # type: ignore - VSA_sparsity=server_args.VSA_sparsity, # type: ignore - device=get_local_torch_device(), # type: ignore - ) # type: ignore - assert ( - attn_metadata is not None - ), "attn_metadata cannot be None" - else: - attn_metadata = None - else: - attn_metadata = None + attn_metadata = self._build_attn_metadata(i, batch, server_args) batch.is_cfg_negative = False with set_forward_context( current_timestep=i, attn_metadata=attn_metadata, forward_batch=batch, - # server_args=server_args ): # Run transformer pred_noise = self.transformer( @@ -251,13 +167,6 @@ class DmdDenoisingStage(DenoisingStage): dtype=pred_video.dtype, generator=batch.generator[0], ).to(self.device) - if sp_group: - noise = rearrange( - noise, - "b (n t) c h w -> b n t c h w", - n=sp_world_size, - ).contiguous() - noise = noise[:, rank_in_sp_group, :, :, :, :] latents = self.scheduler.add_noise( pred_video.flatten(0, 1), noise.flatten(0, 1), @@ -284,11 +193,12 @@ class DmdDenoisingStage(DenoisingStage): (denoising_loop_end_time - denoising_loop_start_time) / len(timesteps), ) - # Gather results if using sequence parallelism - if sp_group: - latents = sequence_model_parallel_all_gather(latents, dim=1) - latents = latents.permute(0, 2, 1, 3, 4) - # Update batch with final latents - batch.latents = latents + self._post_denoising_loop( + batch=batch, + latents=latents, + trajectory_latents=[], + trajectory_timesteps=[], + server_args=server_args, + ) return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index b2bbb0676..8ad061d95 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -93,8 +93,8 @@ class InputValidationStage(PipelineStage): # adjust output image size calculated_width, calculated_height = calculated_size - width = calculated_width if batch.width_not_provided else batch.width - height = calculated_height if batch.height_not_provided else batch.height + width = batch.width or calculated_width + height = batch.height or calculated_height multiple_of = ( server_args.pipeline_config.vae_config.get_vae_scale_factor() * 2 ) @@ -182,16 +182,6 @@ class InputValidationStage(PipelineStage): "`negative_prompt_embeds` must be provided" ) - # Validate height and width - if batch.height is None or batch.width is None: - raise ValueError( - "Height and width must be provided. Please set `height` and `width`." - ) - if batch.height % 8 != 0 or batch.width % 8 != 0: - raise ValueError( - f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}." - ) - # Validate number of inference steps if batch.num_inference_steps <= 0: raise ValueError( @@ -234,8 +224,7 @@ class InputValidationStage(PipelineStage): lambda _: V.string_or_list_strings(batch.prompt) or V.list_not_empty(batch.prompt_embeds), ) - result.add_check("height", batch.height, V.positive_int) - result.add_check("width", batch.width, V.positive_int) + result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) @@ -249,6 +238,13 @@ class InputValidationStage(PipelineStage): def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify input validation stage outputs.""" result = VerificationResult() + result.add_check("height", batch.height, V.positive_int) + result.add_check("width", batch.width, V.positive_int) + # Validate height and width + if batch.height % 8 != 0 or batch.width % 8 != 0: + raise ValueError( + f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}." + ) result.add_check("seeds", batch.seeds, V.list_not_empty) result.add_check("generator", batch.generator, V.generator_or_list_generators) return result diff --git a/python/sglang/multimodal_gen/test/slack_utils.py b/python/sglang/multimodal_gen/test/slack_utils.py index ce25381a8..7237c891a 100644 --- a/python/sglang/multimodal_gen/test/slack_utils.py +++ b/python/sglang/multimodal_gen/test/slack_utils.py @@ -9,6 +9,8 @@ from datetime import datetime from urllib.parse import urlparse from urllib.request import urlopen +from sglang.multimodal_gen.runtime.utils.perf_logger import get_git_commit_hash + logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -47,7 +49,12 @@ except Exception as e: def _get_status_message(run_id, current_case_id, thread_messages=None): date_str = datetime.now().strftime("%d/%m") - base_header = f"*🧵 for nightly test of {date_str}*\n*GitHub Run ID:* {run_id}\n*Total Tasks:* {len(ALL_CASES)}" + base_header = f""""*🧵 for nightly test of {date_str}* +*Git Revision:* {get_git_commit_hash()} +*GitHub Run ID:* {run_id} +*Total Tasks:* {len(ALL_CASES)} + +""" if not ALL_CASES: return base_header