[diffusion] refactor: simplify DmdDenoisingStage (#14269)
This commit is contained in:
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user