[diffusion] refactor: simplify DmdDenoisingStage (#14269)

This commit is contained in:
Mick
2025-12-02 18:59:40 +08:00
committed by GitHub
parent 3067b3f050
commit 9530b76630
6 changed files with 224 additions and 293 deletions

View File

@@ -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),

View File

@@ -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

View File

@@ -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.

View File

@@ -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

View File

@@ -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

View File

@@ -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