minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] refactor: streamline denoising stages (#22633)

Browse Source
This commit is contained in:
Mick
2026-04-13 13:34:37 +08:00
committed by GitHub
parent 7d2c11970c
commit d524f110ac
9 changed files with 1785 additions and 2306 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+25 -9
View File
@@ -151,7 +151,6 @@ def maybe_unpad_latents(latents, batch):
return latents
# config for a single pipeline
@dataclass
class PipelineConfig:
"""The base configuration class for a generation pipeline."""
@@ -360,10 +359,28 @@ class PipelineConfig:
def preprocess_decoding(self, latents, server_args=None, vae=None):
return latents
@staticmethod
def _gather_sp_tensor(tensor: torch.Tensor, *, dim: int) -> torch.Tensor:
"""All-gather an SP-sharded tensor along the specified logical dimension."""
return sequence_model_parallel_all_gather(tensor.contiguous(), dim=dim)
@staticmethod
def _trim_sp_gather_padding(
tensor: torch.Tensor, *, orig_len: int | None, dim: int
) -> torch.Tensor:
"""Trim padding introduced before SP sharding back to the original length."""
if orig_len is None:
return tensor
orig_len = int(orig_len)
if orig_len <= 0 or tensor.shape[dim] <= orig_len:
return tensor
slices = [slice(None)] * tensor.ndim
slices[dim] = slice(orig_len)
return tensor[tuple(slices)]
def gather_latents_for_sp(self, latents, batch=None):
# For video latents [B, C, T_local, H, W], gather along time dim=2
latents = sequence_model_parallel_all_gather(latents, dim=2)
return latents
return self._gather_sp_tensor(latents, dim=2)
def can_shard_audio_latents_for_sp(self, audio_latents) -> bool:
"""Return whether this pipeline uses packed audio latents that can be SP-sharded."""
@@ -405,9 +422,9 @@ class PipelineConfig:
noise_pred = self.gather_latents_for_sp(noise_pred)
raw_latent_shape = getattr(batch, "raw_latent_shape", None)
if raw_latent_shape is not None and noise_pred.dim() == 3:
orig_s = raw_latent_shape[1]
if noise_pred.shape[1] > orig_s:
noise_pred = noise_pred[:, :orig_s, :]
noise_pred = self._trim_sp_gather_padding(
noise_pred, orig_len=raw_latent_shape[1], dim=1
)
return noise_pred
def preprocess_vae_image(self, batch, vae_image_processor):
@@ -850,8 +867,7 @@ class ImagePipelineConfig(PipelineConfig):
def gather_latents_for_sp(self, latents, batch=None):
# For image latents [B, S_local, D], gather along sequence dim=1
latents = sequence_model_parallel_all_gather(latents, dim=1)
return latents
return self._gather_sp_tensor(latents, dim=1)
def _unpad_and_unpack_latents(self, latents, batch):
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
@@ -908,7 +924,7 @@ class SpatialImagePipelineConfig(ImagePipelineConfig):
if latents.dim() != 4:
return super().gather_latents_for_sp(latents, batch=batch)
# Gather along dim=2 (H') to match shard_latents_for_sp
return sequence_model_parallel_all_gather(latents, dim=2)
return self._gather_sp_tensor(latents, dim=2)
@dataclass
python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py
+10 -8
View File
@@ -19,7 +19,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
from sglang.multimodal_gen.runtime.distributed import (
get_sp_parallel_rank,
get_sp_world_size,
sequence_model_parallel_all_gather,
)
@@ -383,7 +382,7 @@ class LTX2PipelineConfig(PipelineConfig):
if get_sp_world_size() <= 1:
return latents
if isinstance(latents, torch.Tensor) and latents.ndim == 3:
return sequence_model_parallel_all_gather(latents.contiguous(), dim=1)
return self._gather_sp_tensor(latents, dim=1)
return super().gather_latents_for_sp(latents, batch=batch)
def shard_audio_latents_for_sp(self, batch, audio_latents):
@@ -432,18 +431,21 @@ class LTX2PipelineConfig(PipelineConfig):
)
def gather_audio_latents_for_sp(self, audio_latents, batch):
"""Gather packed audio latents after SP and trim any pad-only tail tokens."""
if get_sp_world_size() <= 1:
return audio_latents
if not (isinstance(audio_latents, torch.Tensor) and audio_latents.ndim == 3):
return audio_latents
audio_latents = sequence_model_parallel_all_gather(
audio_latents.contiguous(), dim=1
audio_latents = self._gather_sp_tensor(
audio_latents,
dim=1,
)
return self._trim_sp_gather_padding(
audio_latents,
orig_len=getattr(batch, "sp_audio_orig_num_frames", None),
dim=1,
)
orig_num_frames = int(batch.sp_audio_orig_num_frames)
if orig_num_frames > 0:
audio_latents = audio_latents[:, :orig_num_frames, :]
return audio_latents
def prepare_video_rope_coords_for_sp(
self,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py
+4
View File
@@ -56,6 +56,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation impo
from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av import (
LTX2AVLatentPreparationStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import (
LTX2DenoisingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.text_connector import (
LTX2TextConnectorStage,
)
@@ -80,6 +83,7 @@ __all__ = [
"LTX2AVLatentPreparationStage",
"DenoisingStage",
"DmdDenoisingStage",
"LTX2DenoisingStage",
"LTX2AVDenoisingStage",
"CausalDMDDenoisingStage",
"EncodingStage",
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+300 -332
View File
@@ -10,13 +10,13 @@ import math
import os
import time
import weakref
from collections.abc import Iterable
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field, fields
from functools import lru_cache
from typing import Any
import torch
import torch.nn as nn
from einops import rearrange
from tqdm.auto import tqdm
from sglang.multimodal_gen import envs
@@ -25,9 +25,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux import (
Flux2PipelineConfig,
FluxPipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.wan import (
Wan2_2_TI2V_5B_Config,
)
from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
CacheDitConfig,
@@ -41,7 +38,6 @@ from sglang.multimodal_gen.runtime.distributed import (
cfg_model_parallel_all_reduce,
get_local_torch_device,
get_sp_group,
get_sp_parallel_rank,
get_sp_world_size,
get_tp_group,
get_world_group,
@@ -68,6 +64,13 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
PipelineStage,
StageParallelismType,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v import (
blend_wan_ti2v_latents,
expand_wan_ti2v_timestep,
prepare_wan_ti2v_latents,
prepare_wan_ti2v_sp_inputs,
should_apply_wan_ti2v,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
StageValidators as V,
)
@@ -89,12 +92,60 @@ from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiT
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
from sglang.multimodal_gen.utils import dict_to_3d_list, masks_like
from sglang.multimodal_gen.utils import dict_to_3d_list
from sglang.srt.utils.common import get_compiler_backend
logger = init_logger(__name__)
@dataclass(slots=True)
class DenoisingContext:
"""Loop-scoped state shared across the denoising skeleton and its hooks."""
extra_step_kwargs: dict[str, Any]
target_dtype: torch.dtype
autocast_enabled: bool
timesteps: torch.Tensor
num_inference_steps: int
num_warmup_steps: int
image_kwargs: dict[str, Any]
pos_cond_kwargs: dict[str, Any]
neg_cond_kwargs: dict[str, Any]
latents: torch.Tensor
boundary_timestep: float | None
z: torch.Tensor | None
reserved_frames_mask: torch.Tensor | None
seq_len: int | None
guidance: torch.Tensor
is_warmup: bool
trajectory_timesteps: list[torch.Tensor] = field(default_factory=list)
trajectory_latents: list[torch.Tensor] = field(default_factory=list)
extra: dict[str, Any] = field(default_factory=dict)
def __getitem__(self, key: str) -> Any:
return getattr(self, key)
def get(self, key: str, default: Any = None) -> Any:
return getattr(self, key, default)
def to_kwargs(self) -> dict[str, Any]:
"""Return a shallow field mapping for derived context construction."""
return {item.name: getattr(self, item.name) for item in fields(self)}
@dataclass(slots=True)
class DenoisingStepState:
"""Per-step hot-path state computed once and reused within a denoising step."""
step_index: int
t_host: torch.Tensor
t_device: torch.Tensor
t_int: int
current_model: Any
current_guidance_scale: Any
attn_metadata: Any | None
class DenoisingStage(PipelineStage):
"""
Stage for running the denoising loop in diffusion pipelines.
@@ -431,118 +482,6 @@ 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 _handle_boundary_ratio(
self,
server_args,
@@ -572,7 +511,7 @@ class DenoisingStage(PipelineStage):
Prepare all necessary invariant variables for the denoising loop.
Returns:
A dictionary containing all the prepared variables for the denoising loop.
A context object containing the invariant state for the denoising loop.
"""
assert self.transformer is not None
pipeline = self.pipeline() if self.pipeline else None
@@ -642,17 +581,16 @@ class DenoisingStage(PipelineStage):
assert neg_prompt_embeds is not None
# Removed Tensor truthiness assert to avoid GPU sync
# 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
and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config
)
should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)
# 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
seq_len, z, reserved_frames_masks = prepare_wan_ti2v_latents(
self.vae,
latents,
target_dtype,
batch,
server_args,
)
else:
seq_len, z, reserved_frames_masks = (
@@ -667,7 +605,7 @@ class DenoisingStage(PipelineStage):
# Shard z and reserved_frames_mask for TI2V if SP is enabled
if should_preprocess_for_wan_ti2v:
reserved_frames_mask_sp, z_sp = self._postprocess_latents_for_ti2v(
reserved_frames_mask_sp, z_sp = prepare_wan_ti2v_sp_inputs(
z, reserved_frames_masks, batch
)
else:
@@ -732,26 +670,199 @@ class DenoisingStage(PipelineStage):
else:
neg_cond_kwargs = {}
return {
"extra_step_kwargs": extra_step_kwargs,
"target_dtype": target_dtype,
"autocast_enabled": autocast_enabled,
"timesteps": timesteps,
"num_inference_steps": num_inference_steps,
"num_warmup_steps": num_warmup_steps,
"image_kwargs": image_kwargs,
"pos_cond_kwargs": pos_cond_kwargs,
"neg_cond_kwargs": neg_cond_kwargs,
"latents": latents,
"prompt_embeds": prompt_embeds,
"neg_prompt_embeds": neg_prompt_embeds,
"boundary_timestep": boundary_timestep,
"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,
}
return DenoisingContext(
extra_step_kwargs=extra_step_kwargs,
target_dtype=target_dtype,
autocast_enabled=autocast_enabled,
timesteps=timesteps,
num_inference_steps=num_inference_steps,
num_warmup_steps=num_warmup_steps,
image_kwargs=image_kwargs,
pos_cond_kwargs=pos_cond_kwargs,
neg_cond_kwargs=neg_cond_kwargs,
latents=latents,
boundary_timestep=boundary_timestep,
z=z_sp,
reserved_frames_mask=reserved_frames_mask_sp,
seq_len=seq_len,
guidance=guidance,
is_warmup=batch.is_warmup,
)
def _before_denoising_loop(
self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
) -> None:
"""Prepare scheduler state before entering the shared denoising loop."""
self.scheduler.set_begin_index(0)
def _prepare_step_state(
self,
ctx: DenoisingContext,
batch: Req,
server_args: ServerArgs,
step_index: int,
t_host: torch.Tensor,
timesteps_cpu: torch.Tensor,
) -> DenoisingStepState:
"""Build the per-step state shared by the loop and model-specific hooks."""
t_int = int(t_host.item())
t_device = ctx.timesteps[step_index]
current_model, current_guidance_scale = self._select_and_manage_model(
t_int=t_int,
boundary_timestep=ctx.boundary_timestep,
server_args=server_args,
batch=batch,
)
attn_metadata = self._prepare_step_attn_metadata(
ctx=ctx,
batch=batch,
server_args=server_args,
step_index=step_index,
t_int=t_int,
timesteps_cpu=timesteps_cpu,
)
return DenoisingStepState(
step_index=step_index,
t_host=t_host,
t_device=t_device,
t_int=t_int,
current_model=current_model,
current_guidance_scale=current_guidance_scale,
attn_metadata=attn_metadata,
)
def _prepare_step_attn_metadata(
self,
ctx: DenoisingContext,
batch: Req,
server_args: ServerArgs,
step_index: int,
t_int: int,
timesteps_cpu: torch.Tensor,
) -> Any | None:
"""Build attention metadata for the current denoising step."""
# Keep attention metadata preparation overridable so model-specific stages
# can preserve their original semantics without duplicating step state setup.
return self._build_attn_metadata(
step_index,
batch,
server_args,
timestep_value=t_int,
timesteps=timesteps_cpu,
)
def _get_prompt_embeds_validator(
self, batch: Req
) -> Callable[[Any], bool] | list[Callable[[Any], bool]]:
"""Return the prompt-embedding validator used by verify_input."""
del batch
return V.list_not_empty
def _get_negative_prompt_embeds_validator(
self, batch: Req
) -> Callable[[Any], bool] | list[Callable[[Any], bool]]:
"""Return the negative-prompt validator used by verify_input."""
return lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)
def _run_denoising_step(
self,
ctx: DenoisingContext,
step: DenoisingStepState,
batch: Req,
server_args: ServerArgs,
) -> None:
"""Run one scheduler-backed denoising step in the shared base path.
Model-specific stages should override this instead of the whole loop whenever possible to achieve better performance
"""
# 1. Prepare latent inputs in the model's compute dtype.
latent_model_input = ctx.latents.to(ctx.target_dtype)
if batch.image_latent is not None:
assert (
not server_args.pipeline_config.task_type == ModelTaskType.TI2V
), "image latents should not be provided for TI2V task"
latent_model_input = torch.cat(
[latent_model_input, batch.image_latent], dim=1
).to(ctx.target_dtype)
# 2. Expand the timestep to the shape expected by the current model.
timestep = self.expand_timestep_before_forward(
batch,
server_args,
step.t_device,
ctx.target_dtype,
ctx.seq_len,
ctx.reserved_frames_mask,
)
# 3. Apply scheduler-side input scaling before the model forward.
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, step.t_device
)
# 4. Run the model prediction path, including CFG when enabled.
noise_pred = self._predict_noise_with_cfg(
current_model=step.current_model,
latent_model_input=latent_model_input,
timestep=timestep,
batch=batch,
timestep_index=step.step_index,
attn_metadata=step.attn_metadata,
target_dtype=ctx.target_dtype,
current_guidance_scale=step.current_guidance_scale,
image_kwargs=ctx.image_kwargs,
pos_cond_kwargs=ctx.pos_cond_kwargs,
neg_cond_kwargs=ctx.neg_cond_kwargs,
server_args=server_args,
guidance=ctx.guidance,
latents=ctx.latents,
)
if server_args.comfyui_mode:
batch.noise_pred = noise_pred
# 5. Advance the scheduler state with the predicted noise.
ctx.latents = self.scheduler.step(
model_output=noise_pred,
timestep=step.t_device,
sample=ctx.latents,
**ctx.extra_step_kwargs,
return_dict=False,
)[0]
# 6. Re-apply any model-specific latent constraints after the update.
ctx.latents = self.post_forward_for_ti2v_task(
batch,
server_args,
ctx.reserved_frames_mask,
ctx.latents,
ctx.z,
)
def _record_trajectory(
self,
ctx: DenoisingContext,
step: DenoisingStepState,
batch: Req,
server_args: ServerArgs,
) -> None:
"""Append the current step to the returned latent trajectory, if requested."""
if not batch.return_trajectory_latents:
return
ctx.trajectory_timesteps.append(step.t_host)
ctx.trajectory_latents.append(ctx.latents)
def _finalize_denoising_loop(
self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
) -> None:
"""Finalize the shared loop by handing state to post-denoising processing."""
self._post_denoising_loop(
batch=batch,
latents=ctx.latents,
trajectory_latents=ctx.trajectory_latents,
trajectory_timesteps=ctx.trajectory_timesteps,
server_args=server_args,
is_warmup=ctx.is_warmup,
)
def _post_denoising_loop(
self,
@@ -972,45 +1083,18 @@ class DenoisingStage(PipelineStage):
reserved_frames_mask,
):
bsz = batch.raw_latent_shape[0]
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
)
should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)
# 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.
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)
assert seq_len is not None, "Wan TI2V requires a token sequence length."
timestep = expand_wan_ti2v_timestep(
batch,
t_device,
target_dtype,
seq_len,
reserved_frames_mask,
)
else:
timestep = t_device.repeat(bsz)
return timestep
@@ -1018,27 +1102,10 @@ class DenoisingStage(PipelineStage):
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
"""
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
)
"""Re-apply Wan TI2V first-frame conditioning after each denoising step."""
should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)
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.
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
latents = blend_wan_ti2v_latents(latents, reserved_frames_mask, z)
return latents
@@ -1051,159 +1118,56 @@ class DenoisingStage(PipelineStage):
"""
Run the denoising loop.
"""
# Prepare variables for the denoising loop
prepared_vars = self._prepare_denoising_loop(batch, server_args)
extra_step_kwargs = prepared_vars["extra_step_kwargs"]
target_dtype = prepared_vars["target_dtype"]
autocast_enabled = prepared_vars["autocast_enabled"]
timesteps = prepared_vars["timesteps"]
num_inference_steps = prepared_vars["num_inference_steps"]
num_warmup_steps = prepared_vars["num_warmup_steps"]
image_kwargs = prepared_vars["image_kwargs"]
pos_cond_kwargs = prepared_vars["pos_cond_kwargs"]
neg_cond_kwargs = prepared_vars["neg_cond_kwargs"]
latents = prepared_vars["latents"]
boundary_timestep = prepared_vars["boundary_timestep"]
z = prepared_vars["z"]
reserved_frames_mask = prepared_vars["reserved_frames_mask"]
seq_len = prepared_vars["seq_len"]
guidance = prepared_vars["guidance"]
# Initialize lists for ODE trajectory
trajectory_timesteps: list[torch.Tensor] = []
trajectory_latents: list[torch.Tensor] = []
# Run denoising loop
ctx = self._prepare_denoising_loop(batch, server_args)
denoising_start_time = time.time()
self._before_denoising_loop(ctx, batch, server_args)
# to avoid device-sync caused by timestep comparison
is_warmup = batch.is_warmup
self.scheduler.set_begin_index(0)
timesteps_cpu = timesteps.cpu()
timesteps_cpu = ctx.timesteps.cpu()
num_timesteps = timesteps_cpu.shape[0]
with torch.autocast(
device_type=current_platform.device_type,
dtype=target_dtype,
enabled=autocast_enabled,
dtype=ctx.target_dtype,
enabled=ctx.autocast_enabled,
):
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t_host in enumerate(timesteps_cpu):
with self.progress_bar(total=ctx.num_inference_steps) as progress_bar:
for step_index, t_host in enumerate(timesteps_cpu):
with StageProfiler(
f"denoising_step_{i}",
f"denoising_step_{step_index}",
logger=logger,
metrics=batch.metrics,
perf_dump_path_provided=batch.perf_dump_path is not None,
record_as_step=True,
):
t_int = int(t_host.item())
t_device = timesteps[i]
current_model, current_guidance_scale = (
self._select_and_manage_model(
t_int=t_int,
boundary_timestep=boundary_timestep,
server_args=server_args,
batch=batch,
)
)
# Expand latents for I2V
latent_model_input = latents.to(target_dtype)
if batch.image_latent is not None:
assert (
not server_args.pipeline_config.task_type
== ModelTaskType.TI2V
), "image latents should not be provided for TI2V task"
latent_model_input = torch.cat(
[latent_model_input, batch.image_latent], dim=1
).to(target_dtype)
timestep = self.expand_timestep_before_forward(
step = self._prepare_step_state(
ctx,
batch,
server_args,
t_device,
target_dtype,
seq_len,
reserved_frames_mask,
step_index,
t_host,
timesteps_cpu,
)
self._run_denoising_step(ctx, step, batch, server_args)
self._record_trajectory(ctx, step, batch, server_args)
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t_device
)
# Predict noise residual
attn_metadata = self._build_attn_metadata(
i,
batch,
server_args,
timestep_value=t_int,
timesteps=timesteps_cpu,
)
noise_pred = self._predict_noise_with_cfg(
current_model=current_model,
latent_model_input=latent_model_input,
timestep=timestep,
batch=batch,
timestep_index=i,
attn_metadata=attn_metadata,
target_dtype=target_dtype,
current_guidance_scale=current_guidance_scale,
image_kwargs=image_kwargs,
pos_cond_kwargs=pos_cond_kwargs,
neg_cond_kwargs=neg_cond_kwargs,
server_args=server_args,
guidance=guidance,
latents=latents,
)
# Save noise_pred to batch for external access (e.g., ComfyUI)
if server_args.comfyui_mode:
batch.noise_pred = noise_pred
# Compute the previous noisy sample
latents = self.scheduler.step(
model_output=noise_pred,
timestep=t_device,
sample=latents,
**extra_step_kwargs,
return_dict=False,
)[0]
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:
trajectory_timesteps.append(t_host)
trajectory_latents.append(latents)
# Update progress bar
if i == num_timesteps - 1 or (
(i + 1) > num_warmup_steps
and (i + 1) % self.scheduler.order == 0
if step_index == num_timesteps - 1 or (
(step_index + 1) > ctx.num_warmup_steps
and (step_index + 1) % self.scheduler.order == 0
and progress_bar is not None
):
progress_bar.update()
if not is_warmup:
if not ctx.is_warmup:
self.step_profile()
denoising_end_time = time.time()
if num_timesteps > 0 and not is_warmup:
if num_timesteps > 0 and not ctx.is_warmup:
self.log_info(
"average time per step: %.4f seconds",
(denoising_end_time - denoising_start_time) / len(timesteps),
(denoising_end_time - denoising_start_time) / len(ctx.timesteps),
)
self._post_denoising_loop(
batch=batch,
latents=latents,
trajectory_latents=trajectory_latents,
trajectory_timesteps=trajectory_timesteps,
server_args=server_args,
is_warmup=is_warmup,
)
self._finalize_denoising_loop(ctx, batch, server_args)
return batch
# TODO: this will extends the preparation stage, should let subclass/passed-in variables decide which to prepare
@@ -1846,7 +1810,11 @@ class DenoisingStage(PipelineStage):
result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
# disable temporarily for image-generation models
# result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
result.add_check(
"prompt_embeds",
batch.prompt_embeds,
self._get_prompt_embeds_validator(batch),
)
result.add_check("image_embeds", batch.image_embeds, V.is_list)
# result.add_check(
# "image_latent", batch.image_latent, V.none_or_tensor_with_dims(5)
@@ -1865,7 +1833,7 @@ class DenoisingStage(PipelineStage):
result.add_check(
"negative_prompt_embeds",
batch.negative_prompt_embeds,
lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x),
self._get_negative_prompt_embeds_validator(batch),
)
return result
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py
+14 -1957
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py
+1240
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Model-specific helpers and stages for diffusion pipeline components."""
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py
+174
View File
@@ -0,0 +1,174 @@
"""WAN TI2V-specific helpers shared by the generic denoising stage."""
import math
import torch
from einops import rearrange
from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
from sglang.multimodal_gen.configs.pipeline_configs.wan import (
Wan2_2_TI2V_5B_Config,
)
from sglang.multimodal_gen.runtime.distributed import (
get_local_torch_device,
get_sp_parallel_rank,
get_sp_world_size,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.utils import masks_like
def should_apply_wan_ti2v(batch: Req, server_args: ServerArgs) -> bool:
"""Return whether the request should use the Wan2.2 TI2V latent path."""
return bool(
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
)
def prepare_wan_ti2v_latents(
vae: object,
latents: torch.Tensor,
target_dtype: torch.dtype,
batch: Req,
server_args: ServerArgs,
) -> tuple[int, torch.Tensor, list[torch.Tensor]]:
"""Encode the conditioning image and splice it into Wan TI2V latents."""
# 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 vae is not None, "VAE is not provided for TI2V task"
vae = vae.to(batch.condition_image.device)
z = vae.encode(batch.condition_image).mean.float()
if getattr(vae, "device", None) != "cpu" and server_args.vae_cpu_offload:
vae = vae.to("cpu")
if hasattr(vae, "shift_factor") and vae.shift_factor is not None:
if isinstance(vae.shift_factor, torch.Tensor):
z -= vae.shift_factor.to(z.device, z.dtype)
else:
z -= vae.shift_factor
if isinstance(vae.scaling_factor, torch.Tensor):
z = z * vae.scaling_factor.to(z.device, z.dtype)
else:
z = z * vae.scaling_factor
latent_model_input = latents.to(target_dtype)
assert latent_model_input.ndim == 5
latent_for_mask = latent_model_input.squeeze(0)
_, reserved_frames_masks = masks_like([latent_for_mask], zero=True)
reserved_frames_mask = reserved_frames_masks[0].unsqueeze(0)
latents = (
1.0 - reserved_frames_mask
) * z + reserved_frames_mask * latent_model_input
assert latents.ndim == 5
batch.latents = latents.to(get_local_torch_device())
num_frames = 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 = (
((num_frames - 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 prepare_wan_ti2v_sp_inputs(
z: torch.Tensor | None,
reserved_frames_masks: list[torch.Tensor] | None,
batch: Req,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
"""Shard Wan TI2V image-conditioning state to match SP-sharded video latents."""
rank_in_sp_group = get_sp_parallel_rank()
sp_world_size = get_sp_world_size()
if getattr(batch, "did_sp_shard_latents", False):
if z is not None and z.shape[2] == 1:
z_sp = z if rank_in_sp_group == 0 else None
else:
z_sp = z
if reserved_frames_masks is not None:
reserved_frames_mask = reserved_frames_masks[0]
time_dim = reserved_frames_mask.shape[1]
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 = reserved_frames_mask_sp_tensor[
:, rank_in_sp_group, :, :, :
]
else:
reserved_frames_mask_sp = reserved_frames_mask
else:
reserved_frames_mask_sp = None
else:
z_sp = z
reserved_frames_mask_sp = (
reserved_frames_masks[0] if reserved_frames_masks is not None else None
)
return reserved_frames_mask_sp, z_sp
def expand_wan_ti2v_timestep(
batch: Req,
t_device: torch.Tensor,
target_dtype: torch.dtype,
seq_len: int,
reserved_frames_mask: torch.Tensor | None,
) -> torch.Tensor:
"""Expand the timestep tensor for Wan TI2V's first-frame masking semantics."""
batch_size = batch.raw_latent_shape[0]
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:
temp_ts = (reserved_frames_mask[0][:, ::2, ::2] * t_device_rounded).flatten()
temp_ts = torch.cat(
[
temp_ts,
temp_ts.new_ones(local_seq_len - temp_ts.size(0)) * t_device_rounded,
]
)
return temp_ts.unsqueeze(0).repeat(batch_size, 1)
return t_device.repeat(batch_size, local_seq_len)
def blend_wan_ti2v_latents(
latents: torch.Tensor,
reserved_frames_mask: torch.Tensor | None,
z: torch.Tensor | None,
) -> torch.Tensor:
"""Restore Wan TI2V's conditioned first frame after each denoising step."""
if z is None or reserved_frames_mask is None:
return latents
return (
1.0 - reserved_frames_mask.unsqueeze(0)
) * z + reserved_frames_mask.unsqueeze(0) * latents