From d524f110ac677d8bb138bff0e2f306e0ba7b62be Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 13 Apr 2026 13:34:37 +0800 Subject: [PATCH] [diffusion] refactor: streamline denoising stages (#22633) --- .../configs/pipeline_configs/base.py | 34 +- .../configs/pipeline_configs/ltx_2.py | 18 +- .../runtime/pipelines_core/stages/__init__.py | 4 + .../pipelines_core/stages/denoising.py | 632 +++--- .../pipelines_core/stages/denoising_av.py | 1971 +---------------- .../pipelines_core/stages/ltx_2_denoising.py | 1240 +++++++++++ .../stages/model_specific_stages/__init__.py | 1 + .../stages/model_specific_stages/wan_ti2v.py | 174 ++ .../utils/diffusion/comparison_configs.json | 17 + 9 files changed, 1785 insertions(+), 2306 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index b149b9cfb..6406af765 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py index 126cf1a54..00b5dc80c 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py index ba6fc5f46..208791f2e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py @@ -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", 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 7f385c967..7b868be9f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index a688d8f7e..e735960bd 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py @@ -1,54 +1,26 @@ import copy -import json -import math -import os -import time -from io import BytesIO -import av -import numpy as np -import PIL.Image import torch -from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution -from diffusers.models.modeling_outputs import AutoencoderKLOutput from diffusers.utils.torch_utils import randn_tensor -from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( is_ltx23_native_variant, ) -from sglang.multimodal_gen.runtime.distributed import get_sp_world_size -from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context -from sglang.multimodal_gen.runtime.models.vaes.ltx_2_3_condition_encoder import ( - LTX23VideoConditionEncoder, -) -from sglang.multimodal_gen.runtime.models.vision_utils import ( - load_image, - normalize, - numpy_to_pt, - pil_to_numpy, -) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req -from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage -from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( - StageValidators as V, +from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( + LTX2DenoisingStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( - VerificationResult, -) -from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin 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.utils import PRECISION_TO_TYPE logger = init_logger(__name__) -class LTX2AVDenoisingStage(DenoisingStage): +class LTX2AVDenoisingStage(LTX2DenoisingStage): """ - LTX-2 specific denoising stage that handles joint video and audio generation. + Thin AV layer that adds audio trajectory gathering and final unpacking on top of + the LTX-2 denoising semantics. """ def __init__(self, transformer, scheduler, vae=None, audio_vae=None, **kwargs): @@ -56,1862 +28,6 @@ class LTX2AVDenoisingStage(DenoisingStage): transformer=transformer, scheduler=scheduler, vae=vae, **kwargs ) self.audio_vae = audio_vae - self._condition_image_encoder = None - self._condition_image_encoder_dir = None - - @staticmethod - def _get_video_latent_num_frames_for_model( - batch: Req, server_args: ServerArgs, latents: torch.Tensor - ) -> int: - """Return the latent-frame length the DiT model should see. - - - If video latents were time-sharded for SP and are packed as token latents - ([B, S, D]), the model only sees the local shard and must use the local - latent-frame count (stored on the batch during SP sharding). - - Otherwise, fall back to the global latent-frame count inferred from the - requested output frames and the VAE temporal compression ratio. - """ - did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) - is_token_latents = isinstance(latents, torch.Tensor) and latents.ndim == 3 - - if did_sp_shard and is_token_latents: - if not hasattr(batch, "sp_video_latent_num_frames"): - raise ValueError( - "SP-sharded LTX2 token latents require `batch.sp_video_latent_num_frames` " - "to be set by `LTX2PipelineConfig.shard_latents_for_sp()`." - ) - return int(batch.sp_video_latent_num_frames) - - pc = server_args.pipeline_config - return int( - (batch.num_frames - 1) - // int(pc.vae_config.arch_config.temporal_compression_ratio) - + 1 - ) - - @staticmethod - def _truncate_sp_padded_token_latents( - batch: Req, latents: torch.Tensor - ) -> torch.Tensor: - """Remove token padding introduced by SP time-sharding (if applicable).""" - did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) - if not did_sp_shard or not ( - isinstance(latents, torch.Tensor) and latents.ndim == 3 - ): - return latents - - raw_shape = getattr(batch, "raw_latent_shape", None) - if not (isinstance(raw_shape, tuple) and len(raw_shape) == 3): - return latents - - orig_s = int(raw_shape[1]) - cur_s = int(latents.shape[1]) - if cur_s == orig_s: - return latents - if cur_s < orig_s: - raise ValueError( - f"Unexpected gathered token-latents seq_len {cur_s} < original seq_len {orig_s}." - ) - return latents[:, :orig_s, :].contiguous() - - def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None: - """Disable cache-dit for TI2V-style requests (image-conditioned), to avoid stale activations. - - NOTE: base denoising stage calls this hook with (num_inference_steps, batch). - """ - if getattr(self, "_disable_cache_dit_for_request", False): - return - return super()._maybe_enable_cache_dit(num_inference_steps, batch) - - def _get_ltx2_stage1_guider_params( - self, batch: Req, server_args: ServerArgs, stage: str - ) -> dict[str, object] | None: - if stage != "stage1": - return None - return batch.extra.get("ltx2_stage1_guider_params") - - @staticmethod - def _ltx2_should_skip_step(step_index: int, skip_step: int) -> bool: - if skip_step == 0: - return False - return step_index % (skip_step + 1) != 0 - - @staticmethod - def _ltx2_apply_rescale( - cond: torch.Tensor, pred: torch.Tensor, rescale_scale: float - ) -> torch.Tensor: - if rescale_scale == 0.0: - return pred - factor = cond.std() / pred.std() - factor = rescale_scale * factor + (1.0 - rescale_scale) - return pred * factor - - @staticmethod - def _prepare_ltx2_ti2v_clean_state( - latents: torch.Tensor, - image_latent: torch.Tensor, - num_img_tokens: int, - zero_clean_latent: bool, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - latents = latents.clone() - conditioned = image_latent[:, :num_img_tokens, :].to( - device=latents.device, dtype=latents.dtype - ) - latents[:, :num_img_tokens, :] = conditioned - denoise_mask = torch.ones( - (latents.shape[0], latents.shape[1], 1), - device=latents.device, - dtype=torch.float32, - ) - denoise_mask[:, :num_img_tokens, :] = 0.0 - if zero_clean_latent: - clean_latent = torch.zeros_like(latents) - else: - clean_latent = latents.detach().clone() - clean_latent[:, :num_img_tokens, :] = conditioned - return latents, denoise_mask, clean_latent - - @staticmethod - def _ltx2_velocity_to_x0( - sample: torch.Tensor, - velocity: torch.Tensor, - sigma: float | torch.Tensor, - ) -> torch.Tensor: - if isinstance(sigma, torch.Tensor): - sigma = sigma.to(device=sample.device, dtype=torch.float32) - while sigma.ndim < sample.ndim: - sigma = sigma.unsqueeze(-1) - return (sample.float() - sigma * velocity.float()).to(sample.dtype) - return (sample.float() - float(sigma) * velocity.float()).to(sample.dtype) - - @staticmethod - def _repeat_batch_dim(tensor: torch.Tensor, target_batch_size: int) -> torch.Tensor: - """Repeat along batch dim while preserving any tokenwise timestep layout.""" - if tensor.shape[0] == int(target_batch_size): - return tensor - if tensor.shape[0] <= 0 or int(target_batch_size) % int(tensor.shape[0]) != 0: - raise ValueError( - f"Cannot repeat tensor with batch={tensor.shape[0]} to target_batch_size={target_batch_size}" - ) - repeat_factor = int(target_batch_size) // int(tensor.shape[0]) - return tensor.repeat(repeat_factor, *([1] * (tensor.ndim - 1))) - - @staticmethod - def _build_ltx2_sp_padding_mask( - batch: Req, - *, - seq_len: int, - batch_size: int, - key: str, - device: torch.device, - ) -> torch.Tensor | None: - valid = getattr(batch, key, None) - if valid is None: - return None - valid = int(valid) - if valid <= 0 or valid >= int(seq_len): - return None - mask = torch.ones( - (batch_size, int(seq_len)), device=device, dtype=torch.float32 - ) - mask[:, valid:] = 0.0 - return mask - - @staticmethod - def _get_ltx_prompt_attention_mask( - batch: Req, - *, - is_ltx23_variant: bool, - negative: bool = False, - ) -> torch.Tensor | None: - if is_ltx23_variant: - return None - return ( - batch.negative_attention_mask if negative else batch.prompt_attention_mask - ) - - @classmethod - def _should_use_ltx23_legacy_one_stage( - cls, - server_args: ServerArgs, - pipeline_name: str | None, - ) -> bool: - if not is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ): - return False - if server_args.pipeline_class_name == "LTX2TwoStagePipeline": - return False - return pipeline_name != "LTX2TwoStagePipeline" - - @classmethod - def _should_shard_ltx23_legacy_one_stage_audio_latents( - cls, - batch: Req, - server_args: ServerArgs, - ) -> bool: - return bool( - get_sp_world_size() > 1 - and is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ) - and cls._should_use_ltx23_legacy_one_stage(server_args, None) - and server_args.pipeline_config.can_shard_audio_latents_for_sp( - batch.audio_latents - ) - ) - - @classmethod - def _ltx2_calculate_guided_x0( - cls, - *, - cond: torch.Tensor, - uncond_text: torch.Tensor | float, - uncond_perturbed: torch.Tensor | float, - uncond_modality: torch.Tensor | float, - cfg_scale: float, - stg_scale: float, - rescale_scale: float, - modality_scale: float, - ) -> torch.Tensor: - pred = ( - cond - + (cfg_scale - 1.0) * (cond - uncond_text) - + stg_scale * (cond - uncond_perturbed) - + (modality_scale - 1.0) * (cond - uncond_modality) - ) - return cls._ltx2_apply_rescale(cond, pred, rescale_scale) - - @staticmethod - def _resize_center_crop( - img: PIL.Image.Image, *, width: int, height: int - ) -> PIL.Image.Image: - return img.resize((width, height), resample=PIL.Image.Resampling.BILINEAR) - - @staticmethod - def _apply_video_codec_compression( - img_array: np.ndarray, crf: int = 33 - ) -> np.ndarray: - """Encode as a single H.264 frame and decode back to simulate compression artifacts.""" - if crf == 0: - return img_array - height, width = img_array.shape[0] // 2 * 2, img_array.shape[1] // 2 * 2 - img_array = img_array[:height, :width] - buffer = BytesIO() - container = av.open(buffer, mode="w", format="mp4") - stream = container.add_stream( - "libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"} - ) - stream.height, stream.width = height, width - frame = av.VideoFrame.from_ndarray(img_array, format="rgb24").reformat( - format="yuv420p" - ) - container.mux(stream.encode(frame)) - container.mux(stream.encode()) - container.close() - buffer.seek(0) - container = av.open(buffer) - decoded = next(container.decode(container.streams.video[0])) - container.close() - return decoded.to_ndarray(format="rgb24") - - @staticmethod - def _resize_center_crop_tensor( - img: PIL.Image.Image, - *, - width: int, - height: int, - device: torch.device, - dtype: torch.dtype, - apply_codec_compression: bool = True, - codec_crf: int = 33, - ) -> torch.Tensor: - """Resize, center-crop, and normalize to [1, C, 1, H, W] tensor in [-1, 1].""" - img_array = np.array(img).astype(np.uint8)[..., :3] - if apply_codec_compression: - img_array = LTX2AVDenoisingStage._apply_video_codec_compression( - img_array, crf=codec_crf - ) - tensor = ( - torch.from_numpy(img_array.astype(np.float32)) - .permute(2, 0, 1) - .unsqueeze(0) - .to(device=device) - ) - src_h, src_w = tensor.shape[2], tensor.shape[3] - scale = max(height / src_h, width / src_w) - new_h, new_w = math.ceil(src_h * scale), math.ceil(src_w * scale) - tensor = torch.nn.functional.interpolate( - tensor, size=(new_h, new_w), mode="bilinear", align_corners=False - ) - top, left = (new_h - height) // 2, (new_w - width) // 2 - tensor = tensor[:, :, top : top + height, left : left + width] - return ((tensor / 127.5 - 1.0).to(dtype=dtype)).unsqueeze(2) - - @staticmethod - def _pil_to_normed_tensor(img: PIL.Image.Image) -> torch.Tensor: - # PIL -> numpy [0,1] -> torch [B,C,H,W], then [-1,1] - arr = pil_to_numpy(img) - t = numpy_to_pt(arr) - return normalize(t) - - @staticmethod - def _should_apply_ltx2_ti2v(batch: Req) -> bool: - """True if we have an image-latent token prefix to condition with. - - SP note: when token latents are time-sharded, only the rank that owns the - *global* first latent frame should apply TI2V conditioning (rank with start_frame==0). - """ - if ( - batch.image_latent is None - or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0 - ): - return False - did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) - if not did_sp_shard: - return True - return int(getattr(batch, "sp_video_start_frame", 0)) == 0 - - @staticmethod - def _should_replicate_ltx23_audio_for_sp( - batch: Req, - server_args: ServerArgs, - *, - is_ltx23_variant: bool, - ) -> bool: - return False - - def _get_condition_image_encoder( - self, - server_args: ServerArgs, - *, - device: torch.device, - dtype: torch.dtype, - ) -> LTX23VideoConditionEncoder | None: - arch_config = server_args.pipeline_config.vae_config.arch_config - encoder_subdir = str(getattr(arch_config, "condition_encoder_subdir", "")) - if not encoder_subdir: - return None - - vae_model_path = server_args.model_paths["vae"] - encoder_dir = os.path.join(vae_model_path, encoder_subdir) - config_path = os.path.join(encoder_dir, "config.json") - weights_path = os.path.join(encoder_dir, "model.safetensors") - if not os.path.exists(config_path) or not os.path.exists(weights_path): - raise ValueError( - f"LTX-2 condition encoder files not found under {encoder_dir}" - ) - - cached_dir = self._condition_image_encoder_dir - encoder = self._condition_image_encoder - if encoder is None or cached_dir != encoder_dir: - with open(config_path, encoding="utf-8") as f: - config = json.load(f) - encoder = LTX23VideoConditionEncoder(config) - encoder.load_state_dict(safetensors_load_file(weights_path), strict=True) - self._condition_image_encoder = encoder - self._condition_image_encoder_dir = encoder_dir - - encoder = encoder.to(device=device, dtype=dtype) - return encoder - - def _prepare_ltx2_image_latent(self, batch: Req, server_args: ServerArgs) -> None: - """Encode `batch.image_path` into packed token latents for LTX-2 TI2V.""" - if ( - batch.image_latent is not None - and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 - ): - return - batch.ltx2_num_image_tokens = 0 - batch.image_latent = None - - if batch.image_path is None: - return - if batch.width is None or batch.height is None: - raise ValueError("width/height must be provided for LTX-2 TI2V.") - if self.vae is None: - raise ValueError("VAE must be provided for LTX-2 TI2V.") - - image_path = ( - batch.image_path[0] - if isinstance(batch.image_path, list) - else batch.image_path - ) - - img = load_image(image_path) - img_array = np.array(img).astype(np.uint8)[..., :3] - img_array = self._apply_video_codec_compression(img_array, crf=33) - conditioned_img = PIL.Image.fromarray(img_array) - batch.condition_image = self._resize_center_crop( - conditioned_img, width=int(batch.width), height=int(batch.height) - ) - - latents_device = ( - batch.latents.device - if isinstance(batch.latents, torch.Tensor) - else torch.device("cpu") - ) - encode_dtype = batch.latents.dtype - original_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] - vae_autocast_enabled = ( - original_dtype != torch.float32 - ) and not server_args.disable_autocast - condition_image_encoder = self._get_condition_image_encoder( - server_args, device=latents_device, dtype=encode_dtype - ) - if condition_image_encoder is None: - self.vae = self.vae.to(device=latents_device, dtype=encode_dtype) - - video_condition = self._resize_center_crop_tensor( - conditioned_img, - width=int(batch.width), - height=int(batch.height), - device=latents_device, - dtype=encode_dtype, - apply_codec_compression=False, - ) - - with torch.autocast( - device_type=current_platform.device_type, - dtype=original_dtype, - enabled=vae_autocast_enabled, - ): - try: - if ( - condition_image_encoder is None - and server_args.pipeline_config.vae_tiling - ): - self.vae.enable_tiling() - except Exception: - pass - if not vae_autocast_enabled: - video_condition = video_condition.to(encode_dtype) - - if condition_image_encoder is not None: - latent = condition_image_encoder(video_condition) - else: - latent_dist: DiagonalGaussianDistribution = self.vae.encode( - video_condition - ) - if isinstance(latent_dist, AutoencoderKLOutput): - latent_dist = latent_dist.latent_dist - - if condition_image_encoder is None: - mode = server_args.pipeline_config.vae_config.encode_sample_mode() - if mode == "argmax": - latent = latent_dist.mode() - elif mode == "sample": - if batch.generator is None: - raise ValueError("Generator must be provided for VAE sampling.") - latent = latent_dist.sample(batch.generator) - else: - raise ValueError(f"Unsupported encode_sample_mode: {mode}") - - # Per-channel normalization: normalized = (x - mean) / std - mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latent) - std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latent) - latent = (latent - mean) / std - else: - latent = latent.to(dtype=encode_dtype) - - packed = server_args.pipeline_config.maybe_pack_latents( - latent, latent.shape[0], batch - ) - if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): - raise ValueError("Expected packed image latents [B, S0, D].") - - # Fail-fast token count: must match one latent frame's tokens. - vae_sf = int(server_args.pipeline_config.vae_scale_factor) - patch = int(server_args.pipeline_config.patch_size) - latent_h = int(batch.height) // vae_sf - latent_w = int(batch.width) // vae_sf - expected_tokens = (latent_h // patch) * (latent_w // patch) - if int(packed.shape[1]) != int(expected_tokens): - raise ValueError( - "LTX-2 conditioning token count mismatch: " - f"{int(packed.shape[1])=} {int(expected_tokens)=}." - ) - - batch.image_latent = packed - batch.ltx2_num_image_tokens = int(packed.shape[1]) - - if batch.debug: - logger.info( - "LTX2 TI2V conditioning prepared: %d tokens (shape=%s) for %sx%s", - batch.ltx2_num_image_tokens, - tuple(batch.image_latent.shape), - batch.width, - batch.height, - ) - - if condition_image_encoder is None: - self.vae.to(original_dtype) - if server_args.vae_cpu_offload: - self.vae = self.vae.to("cpu") - if condition_image_encoder is not None: - self._condition_image_encoder = condition_image_encoder.to("cpu") - - @torch.no_grad() - def _forward_ltx23_legacy_one_stage( - self, - batch: Req, - server_args: ServerArgs, - prepared_vars: dict[str, object], - ) -> Req: - 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"] - latents = prepared_vars["latents"] - boundary_timestep = prepared_vars["boundary_timestep"] - z = prepared_vars["z"] - reserved_frames_mask = prepared_vars["reserved_frames_mask"] - stage = "stage1" - audio_latents = batch.audio_latents - audio_scheduler = copy.deepcopy(self.scheduler) - batch.ltx23_audio_replicated_for_sp = False - batch.did_sp_shard_audio_latents = False - - latent_num_frames_for_model = self._get_video_latent_num_frames_for_model( - batch=batch, server_args=server_args, latents=latents - ) - latent_height = ( - batch.height - // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio - ) - latent_width = ( - batch.width - // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio - ) - - trajectory_timesteps: list[torch.Tensor] = [] - trajectory_latents: list[torch.Tensor] = [] - trajectory_audio_latents: list[torch.Tensor] = [] - - denoising_start_time = time.time() - - is_warmup = batch.is_warmup - self.scheduler.set_begin_index(0) - audio_scheduler.set_begin_index(0) - timesteps_cpu = timesteps.cpu() - num_timesteps = timesteps_cpu.shape[0] - - do_ti2v = self._should_apply_ltx2_ti2v(batch) - num_img_tokens = int(getattr(batch, "ltx2_num_image_tokens", 0)) - denoise_mask = None - clean_latent = None - if do_ti2v: - if not (isinstance(latents, torch.Tensor) and latents.ndim == 3): - raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") - latents, denoise_mask, clean_latent = self._prepare_ltx2_ti2v_clean_state( - latents=latents, - image_latent=batch.image_latent, - num_img_tokens=num_img_tokens, - zero_clean_latent=True, - ) - - with torch.autocast( - device_type=current_platform.device_type, - dtype=target_dtype, - enabled=autocast_enabled, - ): - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t_host in enumerate(timesteps_cpu): - with StageProfiler( - f"denoising_step_{i}", - 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, - ) - - attn_metadata = self._build_attn_metadata(i, batch, server_args) - - sigmas = getattr(self.scheduler, "sigmas", None) - if sigmas is None or not isinstance(sigmas, torch.Tensor): - raise ValueError( - "Expected scheduler.sigmas to be a tensor for LTX-2." - ) - sigma = sigmas[i].to(device=latents.device, dtype=torch.float32) - sigma_next = sigmas[i + 1].to( - device=latents.device, dtype=torch.float32 - ) - dt = sigma_next - sigma - - latent_model_input = latents.to(target_dtype) - audio_latent_model_input = audio_latents.to(target_dtype) - stage1_guider_params = self._get_ltx2_stage1_guider_params( - batch, server_args, stage - ) - latent_num_frames = latent_num_frames_for_model - - if audio_latent_model_input.ndim == 3: - audio_num_frames_latent = int( - audio_latent_model_input.shape[1] - ) - elif audio_latent_model_input.ndim == 4: - audio_num_frames_latent = int( - audio_latent_model_input.shape[2] - ) - else: - raise ValueError( - f"Unexpected audio latents rank: {audio_latent_model_input.ndim}, shape={tuple(audio_latent_model_input.shape)}" - ) - - video_coords = None - audio_coords = None - - timestep = t_device.expand(int(latent_model_input.shape[0])) - if do_ti2v and denoise_mask is not None: - timestep_video = timestep.unsqueeze( - -1 - ) * denoise_mask.squeeze(-1) - else: - timestep_video = timestep - timestep_audio = timestep - - use_official_cfg_path = stage1_guider_params is None - if use_official_cfg_path: - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = batch.prompt_attention_mask - if batch.do_classifier_free_guidance: - latent_model_input = torch.cat( - [latent_model_input] * 2, dim=0 - ) - audio_latent_model_input = torch.cat( - [audio_latent_model_input] * 2, dim=0 - ) - encoder_hidden_states = torch.cat( - [ - batch.negative_prompt_embeds[0], - encoder_hidden_states, - ], - dim=0, - ) - audio_encoder_hidden_states = torch.cat( - [ - batch.negative_audio_prompt_embeds[0], - audio_encoder_hidden_states, - ], - dim=0, - ) - encoder_attention_mask = torch.cat( - [ - batch.negative_attention_mask, - encoder_attention_mask, - ], - dim=0, - ) - cfg_batch_size = int(latent_model_input.shape[0]) - timestep_video = self._repeat_batch_dim( - timestep_video, cfg_batch_size - ) - timestep_audio = self._repeat_batch_dim( - timestep_audio, cfg_batch_size - ) - - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - model_video, model_audio = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - ) - - model_video = model_video.float() - model_audio = model_audio.float() - if batch.do_classifier_free_guidance: - ( - model_video_uncond, - model_video_text, - ) = model_video.chunk(2) - ( - model_audio_uncond, - model_audio_text, - ) = model_audio.chunk(2) - model_video = model_video_uncond + ( - batch.guidance_scale - * (model_video_text - model_video_uncond) - ) - model_audio = model_audio_uncond + ( - batch.guidance_scale - * (model_audio_text - model_audio_uncond) - ) - v_pos = model_video - a_v_pos = model_audio - - latents = self.scheduler.step( - v_pos, t_device, latents, return_dict=False - )[0] - audio_latents = audio_scheduler.step( - a_v_pos, t_device, audio_latents, return_dict=False - )[0] - latents = self.post_forward_for_ti2v_task( - batch, server_args, reserved_frames_mask, latents, z - ) - - if batch.return_trajectory_latents: - trajectory_timesteps.append(t_host) - trajectory_latents.append(latents) - if audio_latents is not None: - trajectory_audio_latents.append(audio_latents) - - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 - and progress_bar is not None - ): - progress_bar.update() - - if not is_warmup: - self.step_profile() - continue - - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = batch.prompt_attention_mask - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_pos, a_v_pos = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - ) - - if ( - stage1_guider_params is not None - or batch.do_classifier_free_guidance - ): - neg_encoder_hidden_states = ( - batch.negative_prompt_embeds[0] - ) - neg_audio_encoder_hidden_states = ( - batch.negative_audio_prompt_embeds[0] - ) - neg_encoder_attention_mask = ( - batch.negative_attention_mask - ) - - v_neg, a_v_neg = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=neg_encoder_hidden_states, - audio_encoder_hidden_states=neg_audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=neg_encoder_attention_mask, - audio_encoder_attention_mask=neg_encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - ) - else: - v_neg = None - a_v_neg = None - - v_pos = v_pos.float() - a_v_pos = a_v_pos.float() - if v_neg is not None: - v_neg = v_neg.float() - if a_v_neg is not None: - a_v_neg = a_v_neg.float() - - sigma_val = float(sigma.item()) - video_sigma_for_x0: float | torch.Tensor = sigma_val - if do_ti2v and denoise_mask is not None: - video_sigma_for_x0 = sigma.to( - device=latents.device, dtype=torch.float32 - ) * denoise_mask.squeeze(-1) - denoised_video = self._ltx2_velocity_to_x0( - latents, v_pos, video_sigma_for_x0 - ) - denoised_audio = self._ltx2_velocity_to_x0( - audio_latents, a_v_pos, sigma_val - ) - denoised_video_neg = None - denoised_audio_neg = None - denoised_video_perturbed = None - denoised_audio_perturbed = None - denoised_video_modality = None - denoised_audio_modality = None - - if ( - ( - stage1_guider_params is not None - or batch.do_classifier_free_guidance - ) - and v_neg is not None - and a_v_neg is not None - ): - denoised_video_neg = self._ltx2_velocity_to_x0( - latents, v_neg, video_sigma_for_x0 - ) - denoised_audio_neg = self._ltx2_velocity_to_x0( - audio_latents, a_v_neg, sigma_val - ) - if stage1_guider_params is not None: - video_skip = self._ltx2_should_skip_step( - i, int(stage1_guider_params["video_skip_step"]) - ) - audio_skip = self._ltx2_should_skip_step( - i, int(stage1_guider_params["audio_skip_step"]) - ) - - need_perturbed = ( - float(stage1_guider_params["video_stg_scale"]) != 0.0 - or float(stage1_guider_params["audio_stg_scale"]) != 0.0 - ) - if need_perturbed: - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_ptb, a_v_ptb = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - skip_video_self_attn_blocks=tuple( - stage1_guider_params["video_stg_blocks"] - ), - skip_audio_self_attn_blocks=tuple( - stage1_guider_params["audio_stg_blocks"] - ), - ) - denoised_video_perturbed = self._ltx2_velocity_to_x0( - latents, v_ptb.float(), video_sigma_for_x0 - ) - denoised_audio_perturbed = self._ltx2_velocity_to_x0( - audio_latents, a_v_ptb.float(), sigma_val - ) - - need_modality = ( - float(stage1_guider_params["video_modality_scale"]) - != 1.0 - or float(stage1_guider_params["audio_modality_scale"]) - != 1.0 - ) - if need_modality: - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_mod, a_v_mod = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - disable_a2v_cross_attn=True, - disable_v2a_cross_attn=True, - ) - denoised_video_modality = self._ltx2_velocity_to_x0( - latents, v_mod.float(), video_sigma_for_x0 - ) - denoised_audio_modality = self._ltx2_velocity_to_x0( - audio_latents, a_v_mod.float(), sigma_val - ) - - if not video_skip: - denoised_video = self._ltx2_calculate_guided_x0( - cond=denoised_video, - uncond_text=( - denoised_video_neg - if denoised_video_neg is not None - else denoised_video - ), - uncond_perturbed=( - denoised_video_perturbed - if denoised_video_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_video_modality - if denoised_video_modality is not None - else 0.0 - ), - cfg_scale=float( - stage1_guider_params["video_cfg_scale"] - ), - stg_scale=float( - stage1_guider_params["video_stg_scale"] - ), - rescale_scale=float( - stage1_guider_params["video_rescale_scale"] - ), - modality_scale=float( - stage1_guider_params["video_modality_scale"] - ), - ) - if not audio_skip: - denoised_audio = self._ltx2_calculate_guided_x0( - cond=denoised_audio, - uncond_text=( - denoised_audio_neg - if denoised_audio_neg is not None - else denoised_audio - ), - uncond_perturbed=( - denoised_audio_perturbed - if denoised_audio_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_audio_modality - if denoised_audio_modality is not None - else 0.0 - ), - cfg_scale=float( - stage1_guider_params["audio_cfg_scale"] - ), - stg_scale=float( - stage1_guider_params["audio_stg_scale"] - ), - rescale_scale=float( - stage1_guider_params["audio_rescale_scale"] - ), - modality_scale=float( - stage1_guider_params["audio_modality_scale"] - ), - ) - elif ( - batch.do_classifier_free_guidance - and denoised_video_neg is not None - and denoised_audio_neg is not None - ): - denoised_video = denoised_video + ( - batch.guidance_scale - 1.0 - ) * (denoised_video - denoised_video_neg) - denoised_audio = denoised_audio + ( - batch.guidance_scale - 1.0 - ) * (denoised_audio - denoised_audio_neg) - - if ( - do_ti2v - and denoise_mask is not None - and clean_latent is not None - ): - denoised_video = ( - denoised_video * denoise_mask - + clean_latent.float() * (1.0 - denoise_mask) - ) - if sigma_val == 0.0: - v_video = torch.zeros_like(denoised_video) - v_audio = torch.zeros_like(denoised_audio) - else: - v_video = ( - (latents.float() - denoised_video.float()) / sigma_val - ).to(latents.dtype) - v_audio = ( - (audio_latents.float() - denoised_audio.float()) - / sigma_val - ).to(audio_latents.dtype) - - latents = (latents.float() + v_video.float() * dt).to( - dtype=latents.dtype - ) - audio_latents = ( - audio_latents.float() + v_audio.float() * dt - ).to(dtype=audio_latents.dtype) - - latents = self.post_forward_for_ti2v_task( - batch, server_args, reserved_frames_mask, latents, z - ) - - if batch.return_trajectory_latents: - trajectory_timesteps.append(t_host) - trajectory_latents.append(latents) - if audio_latents is not None: - trajectory_audio_latents.append(audio_latents) - - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 - and progress_bar is not None - ): - progress_bar.update() - - if not is_warmup: - self.step_profile() - - denoising_end_time = time.time() - - if num_timesteps > 0 and not is_warmup: - self.log_info( - "average time per step: %.4f seconds", - (denoising_end_time - denoising_start_time) / len(timesteps), - ) - - batch.audio_latents = audio_latents - self._post_denoising_loop( - batch=batch, - latents=latents, - trajectory_latents=trajectory_latents, - trajectory_timesteps=trajectory_timesteps, - trajectory_audio_latents=trajectory_audio_latents, - server_args=server_args, - is_warmup=is_warmup, - ) - - return batch - - @torch.no_grad() - def forward(self, batch: Req, server_args: ServerArgs) -> Req: - """ - Run the denoising loop. - - Args: - batch: The current batch information. - server_args: The inference arguments. - - Returns: - The batch with denoised latents. - """ - # Disable cache-dit for image-conditioned requests (TI2V-style) for correctness/debuggability. - self._disable_cache_dit_for_request = batch.image_path is not None - - # Prepare variables for the denoising loop - prepared_vars = self._prepare_denoising_loop(batch, server_args) - 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"] - latents = prepared_vars["latents"] - boundary_timestep = prepared_vars["boundary_timestep"] - z = prepared_vars["z"] - reserved_frames_mask = prepared_vars["reserved_frames_mask"] - is_ltx23_variant = is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ) - phase = batch.extra.get("ltx2_phase") - pipeline = self.pipeline() if self.pipeline else None - pipeline_name = pipeline.pipeline_name if pipeline is not None else None - use_ltx23_legacy_one_stage = self._should_use_ltx23_legacy_one_stage( - server_args, pipeline_name - ) - stage = ( - phase - if phase is not None - else ("stage1" if use_ltx23_legacy_one_stage else "one_stage") - ) - audio_latents = batch.audio_latents - audio_scheduler = copy.deepcopy(self.scheduler) - - self._prepare_ltx2_image_latent(batch, server_args) - if use_ltx23_legacy_one_stage: - return self._forward_ltx23_legacy_one_stage( - batch, server_args, prepared_vars - ) - do_ti2v = self._should_apply_ltx2_ti2v(batch) - replicate_audio_for_sp = self._should_replicate_ltx23_audio_for_sp( - batch, - server_args, - is_ltx23_variant=is_ltx23_variant, - ) - batch.ltx23_audio_replicated_for_sp = bool(replicate_audio_for_sp) - - if ( - is_ltx23_variant - and get_sp_world_size() > 1 - and server_args.pipeline_config.can_shard_audio_latents_for_sp( - batch.audio_latents - ) - and not replicate_audio_for_sp - and not use_ltx23_legacy_one_stage - ): - ( - batch.audio_latents, - batch.did_sp_shard_audio_latents, - ) = server_args.pipeline_config.shard_audio_latents_for_sp( - batch, batch.audio_latents - ) - audio_latents = batch.audio_latents - else: - batch.did_sp_shard_audio_latents = False - - # For LTX-2 packed token latents, SP sharding happens on the time dimension - # (frames). The model must see local latent frames (RoPE offset is applied - # inside the model using SP rank). - latent_num_frames_for_model = self._get_video_latent_num_frames_for_model( - batch=batch, server_args=server_args, latents=latents - ) - latent_height = ( - batch.height - // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio - ) - latent_width = ( - batch.width - // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio - ) - - # Initialize lists for ODE trajectory - trajectory_timesteps: list[torch.Tensor] = [] - trajectory_latents: list[torch.Tensor] = [] - trajectory_audio_latents: list[torch.Tensor] = [] - - # Run denoising loop - denoising_start_time = time.time() - - # to avoid device-sync caused by timestep comparison - is_warmup = batch.is_warmup - self.scheduler.set_begin_index(0) - audio_scheduler.set_begin_index(0) - timesteps_cpu = timesteps.cpu() - num_timesteps = timesteps_cpu.shape[0] - - num_img_tokens = int(getattr(batch, "ltx2_num_image_tokens", 0)) - denoise_mask = None - clean_latent = None - if do_ti2v: - if not (isinstance(latents, torch.Tensor) and latents.ndim == 3): - raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") - use_zero_clean_latent = is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ) - latents, denoise_mask, clean_latent = self._prepare_ltx2_ti2v_clean_state( - latents=latents, - image_latent=batch.image_latent, - num_img_tokens=num_img_tokens, - zero_clean_latent=use_zero_clean_latent, - ) - with torch.autocast( - device_type=current_platform.device_type, - dtype=target_dtype, - enabled=autocast_enabled, - ): - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t_host in enumerate(timesteps_cpu): - with StageProfiler( - f"denoising_step_{i}", - 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, - ) - - # Predict noise residual - attn_metadata = self._build_attn_metadata(i, batch, server_args) - - # === LTX-2 sigma-space Euler step (flow matching) === - # Use scheduler-generated sigmas (includes terminal sigma=0). - sigmas = getattr(self.scheduler, "sigmas", None) - if sigmas is None or not isinstance(sigmas, torch.Tensor): - raise ValueError( - "Expected scheduler.sigmas to be a tensor for LTX-2." - ) - sigma = sigmas[i].to(device=latents.device, dtype=torch.float32) - sigma_next = sigmas[i + 1].to( - device=latents.device, dtype=torch.float32 - ) - dt = sigma_next - sigma - - latent_model_input = latents.to(target_dtype) - audio_latent_model_input = audio_latents.to(target_dtype) - stage1_guider_params = self._get_ltx2_stage1_guider_params( - batch, server_args, stage - ) - latent_num_frames = latent_num_frames_for_model - - # Audio latent dims - if audio_latent_model_input.ndim == 3: - audio_num_frames_latent = int( - audio_latent_model_input.shape[1] - ) - elif audio_latent_model_input.ndim == 4: - audio_num_frames_latent = int( - audio_latent_model_input.shape[2] - ) - else: - raise ValueError( - f"Unexpected audio latents rank: {audio_latent_model_input.ndim}, shape={tuple(audio_latent_model_input.shape)}" - ) - - video_coords = None - audio_coords = None - if not use_ltx23_legacy_one_stage: - video_coords = server_args.pipeline_config.prepare_video_rope_coords_for_sp( - current_model, - batch, - latent_model_input, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - ) - audio_coords = server_args.pipeline_config.prepare_audio_rope_coords_for_sp( - current_model, - batch, - audio_latent_model_input, - num_frames=audio_num_frames_latent, - ) - - batch_size = int(latent_model_input.shape[0]) - video_num_tokens = int(latent_model_input.shape[1]) - is_ltx23_variant = is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ) - timestep = t_device.expand(batch_size) - if do_ti2v and denoise_mask is not None: - timestep_video = timestep.unsqueeze( - -1 - ) * denoise_mask.squeeze(-1) - elif is_ltx23_variant and not use_ltx23_legacy_one_stage: - timestep_video = timestep.view(batch_size, 1).expand( - batch_size, video_num_tokens - ) - else: - timestep_video = timestep - - if ( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - and audio_latent_model_input.ndim == 3 - ): - audio_num_tokens = int(audio_latent_model_input.shape[1]) - timestep_audio = timestep.view(batch_size, 1).expand( - batch_size, audio_num_tokens - ) - else: - timestep_audio = timestep - prompt_timestep_video = None - prompt_timestep_audio = None - if is_ltx23_variant and not use_ltx23_legacy_one_stage: - timestep_scale_multiplier = float( - getattr( - current_model, "timestep_scale_multiplier", 1000 - ) - ) - prompt_timestep_video = ( - sigma.to( - device=latent_model_input.device, - dtype=torch.float32, - ) - * timestep_scale_multiplier - ).expand(batch_size) - prompt_timestep_audio = ( - sigma.to( - device=audio_latent_model_input.device, - dtype=torch.float32, - ) - * timestep_scale_multiplier - ).expand(batch_size) - - if use_ltx23_legacy_one_stage: - video_self_attention_mask = None - audio_self_attention_mask = None - a2v_cross_attention_mask = None - v2a_cross_attention_mask = None - else: - video_self_attention_mask = ( - self._build_ltx2_sp_padding_mask( - batch, - seq_len=video_num_tokens, - batch_size=batch_size, - key="sp_video_valid_token_count", - device=latent_model_input.device, - ) - ) - audio_self_attention_mask = ( - self._build_ltx2_sp_padding_mask( - batch, - seq_len=audio_num_frames_latent, - batch_size=batch_size, - key="sp_audio_valid_token_count", - device=audio_latent_model_input.device, - ) - ) - a2v_cross_attention_mask = audio_self_attention_mask - v2a_cross_attention_mask = video_self_attention_mask - - def build_model_kwargs( - *, - encoder_hidden_states: torch.Tensor, - audio_encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor | None, - skip_video_self_attn_blocks: tuple[int, ...] | None = None, - skip_audio_self_attn_blocks: tuple[int, ...] | None = None, - disable_a2v_cross_attn: bool = False, - disable_v2a_cross_attn: bool = False, - ) -> dict[str, object]: - kwargs: dict[str, object] = { - "hidden_states": latent_model_input, - "audio_hidden_states": audio_latent_model_input, - "encoder_hidden_states": encoder_hidden_states, - "audio_encoder_hidden_states": audio_encoder_hidden_states, - "timestep": timestep_video, - "audio_timestep": timestep_audio, - "encoder_attention_mask": encoder_attention_mask, - "audio_encoder_attention_mask": encoder_attention_mask, - "num_frames": latent_num_frames, - "height": latent_height, - "width": latent_width, - "fps": batch.fps, - "audio_num_frames": audio_num_frames_latent, - "video_coords": video_coords, - "audio_coords": audio_coords, - "return_latents": False, - "return_dict": False, - } - if not use_ltx23_legacy_one_stage: - kwargs.update( - { - "prompt_timestep": prompt_timestep_video, - "audio_prompt_timestep": prompt_timestep_audio, - "video_self_attention_mask": video_self_attention_mask, - "audio_self_attention_mask": audio_self_attention_mask, - "a2v_cross_attention_mask": a2v_cross_attention_mask, - "v2a_cross_attention_mask": v2a_cross_attention_mask, - "audio_replicated_for_sp": replicate_audio_for_sp, - "legacy_ltx23_one_stage_semantics": False, - } - ) - if skip_video_self_attn_blocks is not None: - kwargs["skip_video_self_attn_blocks"] = ( - skip_video_self_attn_blocks - ) - if skip_audio_self_attn_blocks is not None: - kwargs["skip_audio_self_attn_blocks"] = ( - skip_audio_self_attn_blocks - ) - if disable_a2v_cross_attn: - kwargs["disable_a2v_cross_attn"] = True - if disable_v2a_cross_attn: - kwargs["disable_v2a_cross_attn"] = True - return kwargs - - use_official_cfg_path = stage1_guider_params is None - if use_official_cfg_path: - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = ( - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - ), - ) - ) - if batch.do_classifier_free_guidance: - latent_model_input = torch.cat( - [latent_model_input] * 2, dim=0 - ) - audio_latent_model_input = torch.cat( - [audio_latent_model_input] * 2, dim=0 - ) - encoder_hidden_states = torch.cat( - [ - batch.negative_prompt_embeds[0], - encoder_hidden_states, - ], - dim=0, - ) - audio_encoder_hidden_states = torch.cat( - [ - batch.negative_audio_prompt_embeds[0], - audio_encoder_hidden_states, - ], - dim=0, - ) - if encoder_attention_mask is not None: - encoder_attention_mask = torch.cat( - [ - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - ), - negative=True, - ), - encoder_attention_mask, - ], - dim=0, - ) - cfg_batch_size = int(latent_model_input.shape[0]) - timestep_video = self._repeat_batch_dim( - timestep_video, cfg_batch_size - ) - timestep_audio = self._repeat_batch_dim( - timestep_audio, cfg_batch_size - ) - if prompt_timestep_video is not None: - prompt_timestep_video = self._repeat_batch_dim( - prompt_timestep_video, cfg_batch_size - ) - if prompt_timestep_audio is not None: - prompt_timestep_audio = self._repeat_batch_dim( - prompt_timestep_audio, cfg_batch_size - ) - if video_self_attention_mask is not None: - video_self_attention_mask = self._repeat_batch_dim( - video_self_attention_mask, cfg_batch_size - ) - if audio_self_attention_mask is not None: - audio_self_attention_mask = self._repeat_batch_dim( - audio_self_attention_mask, cfg_batch_size - ) - if a2v_cross_attention_mask is not None: - a2v_cross_attention_mask = self._repeat_batch_dim( - a2v_cross_attention_mask, cfg_batch_size - ) - if v2a_cross_attention_mask is not None: - v2a_cross_attention_mask = self._repeat_batch_dim( - v2a_cross_attention_mask, cfg_batch_size - ) - - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - model_video, model_audio = current_model( - **build_model_kwargs( - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - ) - ) - - model_video = model_video.float() - model_audio = model_audio.float() - if batch.do_classifier_free_guidance: - ( - model_video_uncond, - model_video_text, - ) = model_video.chunk(2) - ( - model_audio_uncond, - model_audio_text, - ) = model_audio.chunk(2) - model_video = model_video_uncond + ( - batch.guidance_scale - * (model_video_text - model_video_uncond) - ) - model_audio = model_audio_uncond + ( - batch.guidance_scale - * (model_audio_text - model_audio_uncond) - ) - v_pos = model_video - a_v_pos = model_audio - v_neg = None - a_v_neg = None - - latents = self.scheduler.step( - v_pos, t_device, latents, return_dict=False - )[0] - audio_latents = audio_scheduler.step( - a_v_pos, t_device, audio_latents, return_dict=False - )[0] - latents = self.post_forward_for_ti2v_task( - batch, server_args, reserved_frames_mask, latents, z - ) - - if batch.return_trajectory_latents: - trajectory_timesteps.append(t_host) - trajectory_latents.append(latents) - if audio_latents is not None: - trajectory_audio_latents.append(audio_latents) - - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 - and progress_bar is not None - ): - progress_bar.update() - - if not is_warmup: - self.step_profile() - continue - else: - # Follow ltx-pipelines structure: separate pos/neg forward passes, - # then apply CFG on denoised (x0) predictions. - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = ( - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - ), - ) - ) - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_pos, a_v_pos = current_model( - **build_model_kwargs( - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - ) - ) - - if ( - stage1_guider_params is not None - or batch.do_classifier_free_guidance - ): - neg_encoder_hidden_states = ( - batch.negative_prompt_embeds[0] - ) - neg_audio_encoder_hidden_states = ( - batch.negative_audio_prompt_embeds[0] - ) - neg_encoder_attention_mask = ( - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - ), - negative=True, - ) - ) - - v_neg, a_v_neg = current_model( - **build_model_kwargs( - encoder_hidden_states=neg_encoder_hidden_states, - audio_encoder_hidden_states=neg_audio_encoder_hidden_states, - encoder_attention_mask=neg_encoder_attention_mask, - ) - ) - else: - v_neg = None - a_v_neg = None - - v_pos = v_pos.float() - a_v_pos = a_v_pos.float() - if v_neg is not None: - v_neg = v_neg.float() - if a_v_neg is not None: - a_v_neg = a_v_neg.float() - - sigma_val = float(sigma.item()) - video_sigma_for_x0: float | torch.Tensor = sigma_val - if do_ti2v and denoise_mask is not None: - video_sigma_for_x0 = sigma.to( - device=latents.device, dtype=torch.float32 - ) * denoise_mask.squeeze(-1) - denoised_video = self._ltx2_velocity_to_x0( - latents, v_pos, video_sigma_for_x0 - ) - denoised_audio = self._ltx2_velocity_to_x0( - audio_latents, a_v_pos, sigma_val - ) - denoised_video_cond = denoised_video - denoised_audio_cond = denoised_audio - denoised_video_neg = None - denoised_audio_neg = None - denoised_video_perturbed = None - denoised_audio_perturbed = None - denoised_video_modality = None - denoised_audio_modality = None - - if ( - ( - stage1_guider_params is not None - or batch.do_classifier_free_guidance - ) - and v_neg is not None - and a_v_neg is not None - ): - denoised_video_neg = self._ltx2_velocity_to_x0( - latents, v_neg, video_sigma_for_x0 - ) - denoised_audio_neg = self._ltx2_velocity_to_x0( - audio_latents, a_v_neg, sigma_val - ) - if stage1_guider_params is not None: - video_skip = self._ltx2_should_skip_step( - i, int(stage1_guider_params["video_skip_step"]) - ) - audio_skip = self._ltx2_should_skip_step( - i, int(stage1_guider_params["audio_skip_step"]) - ) - - need_perturbed = ( - float(stage1_guider_params["video_stg_scale"]) != 0.0 - or float(stage1_guider_params["audio_stg_scale"]) != 0.0 - ) - if need_perturbed: - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_ptb, a_v_ptb = current_model( - **build_model_kwargs( - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - skip_video_self_attn_blocks=tuple( - stage1_guider_params["video_stg_blocks"] - ), - skip_audio_self_attn_blocks=tuple( - stage1_guider_params["audio_stg_blocks"] - ), - ) - ) - denoised_video_perturbed = self._ltx2_velocity_to_x0( - latents, v_ptb.float(), video_sigma_for_x0 - ) - denoised_audio_perturbed = self._ltx2_velocity_to_x0( - audio_latents, a_v_ptb.float(), sigma_val - ) - - need_modality = ( - float(stage1_guider_params["video_modality_scale"]) - != 1.0 - or float(stage1_guider_params["audio_modality_scale"]) - != 1.0 - ) - if need_modality: - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_mod, a_v_mod = current_model( - **build_model_kwargs( - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - disable_a2v_cross_attn=True, - disable_v2a_cross_attn=True, - ) - ) - denoised_video_modality = self._ltx2_velocity_to_x0( - latents, v_mod.float(), video_sigma_for_x0 - ) - denoised_audio_modality = self._ltx2_velocity_to_x0( - audio_latents, a_v_mod.float(), sigma_val - ) - - if not video_skip: - denoised_video = self._ltx2_calculate_guided_x0( - cond=denoised_video, - uncond_text=( - denoised_video_neg - if denoised_video_neg is not None - else denoised_video - ), - uncond_perturbed=( - denoised_video_perturbed - if denoised_video_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_video_modality - if denoised_video_modality is not None - else 0.0 - ), - cfg_scale=float( - stage1_guider_params["video_cfg_scale"] - ), - stg_scale=float( - stage1_guider_params["video_stg_scale"] - ), - rescale_scale=float( - stage1_guider_params["video_rescale_scale"] - ), - modality_scale=float( - stage1_guider_params["video_modality_scale"] - ), - ) - if not audio_skip: - denoised_audio = self._ltx2_calculate_guided_x0( - cond=denoised_audio, - uncond_text=( - denoised_audio_neg - if denoised_audio_neg is not None - else denoised_audio - ), - uncond_perturbed=( - denoised_audio_perturbed - if denoised_audio_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_audio_modality - if denoised_audio_modality is not None - else 0.0 - ), - cfg_scale=float( - stage1_guider_params["audio_cfg_scale"] - ), - stg_scale=float( - stage1_guider_params["audio_stg_scale"] - ), - rescale_scale=float( - stage1_guider_params["audio_rescale_scale"] - ), - modality_scale=float( - stage1_guider_params["audio_modality_scale"] - ), - ) - elif ( - batch.do_classifier_free_guidance - and denoised_video_neg is not None - and denoised_audio_neg is not None - ): - denoised_video = denoised_video + ( - batch.guidance_scale - 1.0 - ) * (denoised_video - denoised_video_neg) - denoised_audio = denoised_audio + ( - batch.guidance_scale - 1.0 - ) * (denoised_audio - denoised_audio_neg) - - # Apply conditioning mask (keep conditioned tokens clean). - if ( - do_ti2v - and denoise_mask is not None - and clean_latent is not None - ): - denoised_video = ( - denoised_video * denoise_mask - + clean_latent.float() * (1.0 - denoise_mask) - ) - # Euler step in sigma space: x_next = x + (sigma_next - sigma) * v, - # where v = (x - x0) / sigma. - if sigma_val == 0.0: - v_video = torch.zeros_like(denoised_video) - v_audio = torch.zeros_like(denoised_audio) - else: - v_video = ( - (latents.float() - denoised_video.float()) / sigma_val - ).to(latents.dtype) - v_audio = ( - (audio_latents.float() - denoised_audio.float()) - / sigma_val - ).to(audio_latents.dtype) - - latents = (latents.float() + v_video.float() * dt).to( - dtype=latents.dtype - ) - audio_latents = ( - audio_latents.float() + v_audio.float() * dt - ).to(dtype=audio_latents.dtype) - - 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) - if audio_latents is not None: - trajectory_audio_latents.append(audio_latents) - - # Update progress bar - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 - and progress_bar is not None - ): - progress_bar.update() - - if not is_warmup: - self.step_profile() - - denoising_end_time = time.time() - - if num_timesteps > 0 and not is_warmup: - self.log_info( - "average time per step: %.4f seconds", - (denoising_end_time - denoising_start_time) / len(timesteps), - ) - - batch.audio_latents = audio_latents - self._post_denoising_loop( - batch=batch, - latents=latents, - trajectory_latents=trajectory_latents, - trajectory_timesteps=trajectory_timesteps, - trajectory_audio_latents=trajectory_audio_latents, - server_args=server_args, - is_warmup=is_warmup, - ) - - return batch def _post_denoising_loop( self, @@ -1925,7 +41,7 @@ class LTX2AVDenoisingStage(DenoisingStage): *args, **kwargs, ): - # 1. Handle Trajectory (Video) - Copy from base + """Finalize AV requests by gathering audio latents and unpacking both streams.""" if trajectory_latents: trajectory_tensor = torch.stack(trajectory_latents, dim=1) trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0) @@ -1936,26 +52,16 @@ class LTX2AVDenoisingStage(DenoisingStage): latents, trajectory_tensor = self._postprocess_sp_latents( batch, latents, trajectory_tensor ) - - # If SP time-sharding padded whole frames worth of tokens, remove padding - # after gather and before unpacking. latents = self._truncate_sp_padded_token_latents(batch, latents) if trajectory_tensor is not None and trajectory_timesteps_tensor is not None: batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu() batch.trajectory_latents = trajectory_tensor.cpu() - # 2. Handle Trajectory (Audio) - LTX-2 specific if trajectory_audio_latents: trajectory_audio_tensor = torch.stack(trajectory_audio_latents, dim=1) - # We don't have SP support for audio latents yet (or needed?) batch.trajectory_audio_latents = trajectory_audio_tensor.cpu() - # 3. Unpack and Denormalize - # Call pipeline_config._unpad_and_unpack_latents - # latents is video latents. - # batch.audio_latents is audio latents. - audio_latents = batch.audio_latents if batch.did_sp_shard_audio_latents and isinstance(audio_latents, torch.Tensor): audio_latents = server_args.pipeline_config.gather_audio_latents_for_sp( @@ -1963,7 +69,6 @@ class LTX2AVDenoisingStage(DenoisingStage): ) batch.audio_latents = audio_latents - # NOTE: self.vae and self.audio_vae should be populated via __init__ or manual setting if self.vae is None or self.audio_vae is None: logger.warning( "VAE or Audio VAE not found in DenoisingStage. Skipping unpack and denormalize." @@ -1971,66 +76,22 @@ class LTX2AVDenoisingStage(DenoisingStage): batch.latents = latents batch.audio_latents = audio_latents else: - ( - latents, - audio_latents, - ) = server_args.pipeline_config._unpad_and_unpack_latents( - latents, audio_latents, batch, self.vae, self.audio_vae + latents, audio_latents = ( + server_args.pipeline_config._unpad_and_unpack_latents( + latents, audio_latents, batch, self.vae, self.audio_vae + ) ) - batch.latents = latents batch.audio_latents = audio_latents - # 4. Cleanup - # TODO: make this a general denoising-stage hook + if isinstance(self.transformer, OffloadableDiTMixin): for manager in self.transformer.layerwise_offload_managers: manager.release_all() - def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: - """Verify denoising stage inputs. - - Note: LTX-2 connector stage converts `prompt_embeds`/`negative_prompt_embeds` - from list-of-tensors to a single tensor (video context) and stores audio - context separately. - """ - - result = VerificationResult() - result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) - - # LTX-2 may carry prompt embeddings as either a tensor (preferred) or legacy list. - result.add_check( - "prompt_embeds", - batch.prompt_embeds, - lambda x: V.is_tensor(x) or V.list_not_empty(x), - ) - - # Keep base expectation: image_embeds is always a list (may be empty). - result.add_check("image_embeds", batch.image_embeds, V.is_list) - - result.add_check( - "num_inference_steps", batch.num_inference_steps, V.positive_int - ) - result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) - result.add_check("eta", batch.eta, V.non_negative_float) - result.add_check("generator", batch.generator, V.generator_or_list_generators) - result.add_check( - "do_classifier_free_guidance", - batch.do_classifier_free_guidance, - V.bool_value, - ) - - # When CFG is enabled, negative prompt embeddings must exist (tensor or legacy list). - result.add_check( - "negative_prompt_embeds", - batch.negative_prompt_embeds, - lambda x: (not batch.do_classifier_free_guidance) - or V.is_tensor(x) - or V.list_not_empty(x), - ) - return result - class LTX2RefinementStage(LTX2AVDenoisingStage): + """Stage-2 refinement wrapper that re-noises distilled LTX latents once.""" + def __init__( self, transformer, scheduler, distilled_sigmas, vae=None, audio_vae=None ): @@ -2085,9 +146,6 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): @staticmethod def _should_reset_stage2_generators(server_args: ServerArgs) -> bool: - # Official LTX-2.3 two-stage refinement continues from the generator state - # after stage 1. Resetting back to the request seed changes the distilled - # noise injection immediately at stage 2 step 0. arch_config = getattr( server_args.pipeline_config.vae_config, "arch_config", None ) @@ -2096,6 +154,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): return "LTX-2.3" not in str(getattr(server_args, "model_path", "")) def forward(self, batch: Req, server_args: ServerArgs) -> Req: + """Run the distilled refinement schedule on top of the shared AV denoiser.""" batch.extra["ltx2_phase"] = "stage2" if self._should_reset_stage2_generators(server_args): self._reset_stage2_generators(batch) @@ -2125,11 +184,9 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): device=batch.audio_latents.device, dtype=torch.float32 ) - # Stage 2 runs at full resolution, so Stage 1 TI2V conditioning is invalid. batch.image_latent = None batch.ltx2_num_image_tokens = 0 - # Use a private scheduler copy to avoid mutating shared state. original_scheduler = self.scheduler original_batch_timesteps = batch.timesteps original_batch_num_inference_steps = batch.num_inference_steps diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py new file mode 100644 index 000000000..a644fc693 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py @@ -0,0 +1,1240 @@ +import copy +import json +import math +import os +from dataclasses import dataclass, field +from io import BytesIO + +import av +import numpy as np +import PIL.Image +import torch +from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from safetensors.torch import load_file as safetensors_load_file + +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + is_ltx23_native_variant, +) +from sglang.multimodal_gen.runtime.distributed import get_sp_world_size +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.models.vaes.ltx_2_3_condition_encoder import ( + LTX23VideoConditionEncoder, +) +from sglang.multimodal_gen.runtime.models.vision_utils import ( + load_image, + normalize, + numpy_to_pt, + pil_to_numpy, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingContext, + DenoisingStage, + DenoisingStepState, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + StageValidators as V, +) +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +@dataclass(slots=True) +class LTX2DenoisingContext(DenoisingContext): + """Loop-scoped denoising state for joint LTX-2 video and audio generation.""" + + audio_latents: torch.Tensor | None = None + audio_scheduler: object | None = None + is_ltx23_variant: bool = False + use_ltx23_legacy_one_stage: bool = False + replicate_audio_for_sp: bool = False + stage: str = "one_stage" + latent_num_frames_for_model: int = 0 + latent_height: int = 0 + latent_width: int = 0 + denoise_mask: torch.Tensor | None = None + clean_latent: torch.Tensor | None = None + trajectory_audio_latents: list[torch.Tensor] = field(default_factory=list) + + +class LTX2DenoisingStage(DenoisingStage): + """ + LTX-2 specific denoising stage that handles joint video and audio generation. + """ + + def __init__(self, transformer, scheduler, vae=None, **kwargs): + super().__init__( + transformer=transformer, scheduler=scheduler, vae=vae, **kwargs + ) + self._condition_image_encoder = None + self._condition_image_encoder_dir = None + + @staticmethod + def _get_video_latent_num_frames_for_model( + batch: Req, server_args: ServerArgs, latents: torch.Tensor + ) -> int: + """Return the latent-frame length the DiT model should see. + + - If video latents were time-sharded for SP and are packed as token latents + ([B, S, D]), the model only sees the local shard and must use the local + latent-frame count (stored on the batch during SP sharding). + - Otherwise, fall back to the global latent-frame count inferred from the + requested output frames and the VAE temporal compression ratio. + """ + did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) + is_token_latents = isinstance(latents, torch.Tensor) and latents.ndim == 3 + + if did_sp_shard and is_token_latents: + if not hasattr(batch, "sp_video_latent_num_frames"): + raise ValueError( + "SP-sharded LTX2 token latents require `batch.sp_video_latent_num_frames` " + "to be set by `LTX2PipelineConfig.shard_latents_for_sp()`." + ) + return int(batch.sp_video_latent_num_frames) + + pc = server_args.pipeline_config + return int( + (batch.num_frames - 1) + // int(pc.vae_config.arch_config.temporal_compression_ratio) + + 1 + ) + + @staticmethod + def _truncate_sp_padded_token_latents( + batch: Req, latents: torch.Tensor + ) -> torch.Tensor: + """Remove token padding introduced by SP time-sharding (if applicable).""" + did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) + if not did_sp_shard or not ( + isinstance(latents, torch.Tensor) and latents.ndim == 3 + ): + return latents + + raw_shape = getattr(batch, "raw_latent_shape", None) + if not (isinstance(raw_shape, tuple) and len(raw_shape) == 3): + return latents + + orig_s = int(raw_shape[1]) + cur_s = int(latents.shape[1]) + if cur_s == orig_s: + return latents + if cur_s < orig_s: + raise ValueError( + f"Unexpected gathered token-latents seq_len {cur_s} < original seq_len {orig_s}." + ) + return latents[:, :orig_s, :].contiguous() + + def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None: + """Disable cache-dit for TI2V-style requests (image-conditioned), to avoid stale activations. + + NOTE: base denoising stage calls this hook with (num_inference_steps, batch). + """ + if getattr(self, "_disable_cache_dit_for_request", False): + return + return super()._maybe_enable_cache_dit(num_inference_steps, batch) + + def _get_ltx2_stage1_guider_params( + self, batch: Req, server_args: ServerArgs, stage: str + ) -> dict[str, object] | None: + if stage != "stage1": + return None + return batch.extra.get("ltx2_stage1_guider_params") + + @staticmethod + def _ltx2_should_skip_step(step_index: int, skip_step: int) -> bool: + if skip_step == 0: + return False + return step_index % (skip_step + 1) != 0 + + @staticmethod + def _ltx2_apply_rescale( + cond: torch.Tensor, pred: torch.Tensor, rescale_scale: float + ) -> torch.Tensor: + if rescale_scale == 0.0: + return pred + factor = cond.std() / pred.std() + factor = rescale_scale * factor + (1.0 - rescale_scale) + return pred * factor + + @staticmethod + def _prepare_ltx2_ti2v_clean_state( + latents: torch.Tensor, + image_latent: torch.Tensor, + num_img_tokens: int, + zero_clean_latent: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + latents = latents.clone() + conditioned = image_latent[:, :num_img_tokens, :].to( + device=latents.device, dtype=latents.dtype + ) + latents[:, :num_img_tokens, :] = conditioned + denoise_mask = torch.ones( + (latents.shape[0], latents.shape[1], 1), + device=latents.device, + dtype=torch.float32, + ) + denoise_mask[:, :num_img_tokens, :] = 0.0 + if zero_clean_latent: + clean_latent = torch.zeros_like(latents) + else: + clean_latent = latents.detach().clone() + clean_latent[:, :num_img_tokens, :] = conditioned + return latents, denoise_mask, clean_latent + + @staticmethod + def _ltx2_velocity_to_x0( + sample: torch.Tensor, + velocity: torch.Tensor, + sigma: float | torch.Tensor, + ) -> torch.Tensor: + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(device=sample.device, dtype=torch.float32) + while sigma.ndim < sample.ndim: + sigma = sigma.unsqueeze(-1) + return (sample.float() - sigma * velocity.float()).to(sample.dtype) + return (sample.float() - float(sigma) * velocity.float()).to(sample.dtype) + + @staticmethod + def _repeat_batch_dim(tensor: torch.Tensor, target_batch_size: int) -> torch.Tensor: + """Repeat along batch dim while preserving any tokenwise timestep layout.""" + if tensor.shape[0] == int(target_batch_size): + return tensor + if tensor.shape[0] <= 0 or int(target_batch_size) % int(tensor.shape[0]) != 0: + raise ValueError( + f"Cannot repeat tensor with batch={tensor.shape[0]} to target_batch_size={target_batch_size}" + ) + repeat_factor = int(target_batch_size) // int(tensor.shape[0]) + return tensor.repeat(repeat_factor, *([1] * (tensor.ndim - 1))) + + @staticmethod + def _build_ltx2_sp_padding_mask( + batch: Req, + *, + seq_len: int, + batch_size: int, + key: str, + device: torch.device, + ) -> torch.Tensor | None: + valid = getattr(batch, key, None) + if valid is None: + return None + valid = int(valid) + if valid <= 0 or valid >= int(seq_len): + return None + mask = torch.ones( + (batch_size, int(seq_len)), device=device, dtype=torch.float32 + ) + mask[:, valid:] = 0.0 + return mask + + @staticmethod + def _get_ltx_prompt_attention_mask( + batch: Req, + *, + is_ltx23_variant: bool, + negative: bool = False, + ) -> torch.Tensor | None: + if is_ltx23_variant: + return None + return ( + batch.negative_attention_mask if negative else batch.prompt_attention_mask + ) + + @classmethod + def _should_use_ltx23_legacy_one_stage( + cls, + server_args: ServerArgs, + pipeline_name: str | None, + ) -> bool: + if not is_ltx23_native_variant( + server_args.pipeline_config.vae_config.arch_config + ): + return False + if server_args.pipeline_class_name == "LTX2TwoStagePipeline": + return False + return pipeline_name != "LTX2TwoStagePipeline" + + @classmethod + def _should_shard_ltx23_legacy_one_stage_audio_latents( + cls, + batch: Req, + server_args: ServerArgs, + ) -> bool: + return bool( + get_sp_world_size() > 1 + and is_ltx23_native_variant( + server_args.pipeline_config.vae_config.arch_config + ) + and cls._should_use_ltx23_legacy_one_stage(server_args, None) + and server_args.pipeline_config.can_shard_audio_latents_for_sp( + batch.audio_latents + ) + ) + + @classmethod + def _ltx2_calculate_guided_x0( + cls, + *, + cond: torch.Tensor, + uncond_text: torch.Tensor | float, + uncond_perturbed: torch.Tensor | float, + uncond_modality: torch.Tensor | float, + cfg_scale: float, + stg_scale: float, + rescale_scale: float, + modality_scale: float, + ) -> torch.Tensor: + pred = ( + cond + + (cfg_scale - 1.0) * (cond - uncond_text) + + stg_scale * (cond - uncond_perturbed) + + (modality_scale - 1.0) * (cond - uncond_modality) + ) + return cls._ltx2_apply_rescale(cond, pred, rescale_scale) + + @staticmethod + def _resize_center_crop( + img: PIL.Image.Image, *, width: int, height: int + ) -> PIL.Image.Image: + return img.resize((width, height), resample=PIL.Image.Resampling.BILINEAR) + + @staticmethod + def _apply_video_codec_compression( + img_array: np.ndarray, crf: int = 33 + ) -> np.ndarray: + """Encode as a single H.264 frame and decode back to simulate compression artifacts.""" + if crf == 0: + return img_array + height, width = img_array.shape[0] // 2 * 2, img_array.shape[1] // 2 * 2 + img_array = img_array[:height, :width] + buffer = BytesIO() + container = av.open(buffer, mode="w", format="mp4") + stream = container.add_stream( + "libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"} + ) + stream.height, stream.width = height, width + frame = av.VideoFrame.from_ndarray(img_array, format="rgb24").reformat( + format="yuv420p" + ) + container.mux(stream.encode(frame)) + container.mux(stream.encode()) + container.close() + buffer.seek(0) + container = av.open(buffer) + decoded = next(container.decode(container.streams.video[0])) + container.close() + return decoded.to_ndarray(format="rgb24") + + @staticmethod + def _resize_center_crop_tensor( + img: PIL.Image.Image, + *, + width: int, + height: int, + device: torch.device, + dtype: torch.dtype, + apply_codec_compression: bool = True, + codec_crf: int = 33, + ) -> torch.Tensor: + """Resize, center-crop, and normalize to [1, C, 1, H, W] tensor in [-1, 1].""" + img_array = np.array(img).astype(np.uint8)[..., :3] + if apply_codec_compression: + img_array = LTX2DenoisingStage._apply_video_codec_compression( + img_array, crf=codec_crf + ) + tensor = ( + torch.from_numpy(img_array.astype(np.float32)) + .permute(2, 0, 1) + .unsqueeze(0) + .to(device=device) + ) + src_h, src_w = tensor.shape[2], tensor.shape[3] + scale = max(height / src_h, width / src_w) + new_h, new_w = math.ceil(src_h * scale), math.ceil(src_w * scale) + tensor = torch.nn.functional.interpolate( + tensor, size=(new_h, new_w), mode="bilinear", align_corners=False + ) + top, left = (new_h - height) // 2, (new_w - width) // 2 + tensor = tensor[:, :, top : top + height, left : left + width] + return ((tensor / 127.5 - 1.0).to(dtype=dtype)).unsqueeze(2) + + @staticmethod + def _pil_to_normed_tensor(img: PIL.Image.Image) -> torch.Tensor: + # PIL -> numpy [0,1] -> torch [B,C,H,W], then [-1,1] + arr = pil_to_numpy(img) + t = numpy_to_pt(arr) + return normalize(t) + + @staticmethod + def _should_apply_ltx2_ti2v(batch: Req) -> bool: + """True if we have an image-latent token prefix to condition with. + + SP note: when token latents are time-sharded, only the rank that owns the + *global* first latent frame should apply TI2V conditioning (rank with start_frame==0). + """ + if ( + batch.image_latent is None + or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0 + ): + return False + did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) + if not did_sp_shard: + return True + return int(getattr(batch, "sp_video_start_frame", 0)) == 0 + + @staticmethod + def _should_replicate_ltx23_audio_for_sp( + batch: Req, + server_args: ServerArgs, + *, + is_ltx23_variant: bool, + ) -> bool: + return False + + def _get_condition_image_encoder( + self, + server_args: ServerArgs, + *, + device: torch.device, + dtype: torch.dtype, + ) -> LTX23VideoConditionEncoder | None: + arch_config = server_args.pipeline_config.vae_config.arch_config + encoder_subdir = str(getattr(arch_config, "condition_encoder_subdir", "")) + if not encoder_subdir: + return None + + vae_model_path = server_args.model_paths["vae"] + encoder_dir = os.path.join(vae_model_path, encoder_subdir) + config_path = os.path.join(encoder_dir, "config.json") + weights_path = os.path.join(encoder_dir, "model.safetensors") + if not os.path.exists(config_path) or not os.path.exists(weights_path): + raise ValueError( + f"LTX-2 condition encoder files not found under {encoder_dir}" + ) + + cached_dir = self._condition_image_encoder_dir + encoder = self._condition_image_encoder + if encoder is None or cached_dir != encoder_dir: + with open(config_path, encoding="utf-8") as f: + config = json.load(f) + encoder = LTX23VideoConditionEncoder(config) + encoder.load_state_dict(safetensors_load_file(weights_path), strict=True) + self._condition_image_encoder = encoder + self._condition_image_encoder_dir = encoder_dir + + encoder = encoder.to(device=device, dtype=dtype) + return encoder + + def _prepare_ltx2_image_latent(self, batch: Req, server_args: ServerArgs) -> None: + """Encode `batch.image_path` into packed token latents for LTX-2 TI2V.""" + if ( + batch.image_latent is not None + and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 + ): + return + batch.ltx2_num_image_tokens = 0 + batch.image_latent = None + + if batch.image_path is None: + return + if batch.width is None or batch.height is None: + raise ValueError("width/height must be provided for LTX-2 TI2V.") + if self.vae is None: + raise ValueError("VAE must be provided for LTX-2 TI2V.") + + image_path = ( + batch.image_path[0] + if isinstance(batch.image_path, list) + else batch.image_path + ) + + img = load_image(image_path) + img_array = np.array(img).astype(np.uint8)[..., :3] + img_array = self._apply_video_codec_compression(img_array, crf=33) + conditioned_img = PIL.Image.fromarray(img_array) + batch.condition_image = self._resize_center_crop( + conditioned_img, width=int(batch.width), height=int(batch.height) + ) + + latents_device = ( + batch.latents.device + if isinstance(batch.latents, torch.Tensor) + else torch.device("cpu") + ) + encode_dtype = batch.latents.dtype + original_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + original_dtype != torch.float32 + ) and not server_args.disable_autocast + condition_image_encoder = self._get_condition_image_encoder( + server_args, device=latents_device, dtype=encode_dtype + ) + if condition_image_encoder is None: + self.vae = self.vae.to(device=latents_device, dtype=encode_dtype) + + video_condition = self._resize_center_crop_tensor( + conditioned_img, + width=int(batch.width), + height=int(batch.height), + device=latents_device, + dtype=encode_dtype, + apply_codec_compression=False, + ) + + with torch.autocast( + device_type=current_platform.device_type, + dtype=original_dtype, + enabled=vae_autocast_enabled, + ): + try: + if ( + condition_image_encoder is None + and server_args.pipeline_config.vae_tiling + ): + self.vae.enable_tiling() + except Exception: + pass + if not vae_autocast_enabled: + video_condition = video_condition.to(encode_dtype) + + if condition_image_encoder is not None: + latent = condition_image_encoder(video_condition) + else: + latent_dist: DiagonalGaussianDistribution = self.vae.encode( + video_condition + ) + if isinstance(latent_dist, AutoencoderKLOutput): + latent_dist = latent_dist.latent_dist + + if condition_image_encoder is None: + mode = server_args.pipeline_config.vae_config.encode_sample_mode() + if mode == "argmax": + latent = latent_dist.mode() + elif mode == "sample": + if batch.generator is None: + raise ValueError("Generator must be provided for VAE sampling.") + latent = latent_dist.sample(batch.generator) + else: + raise ValueError(f"Unsupported encode_sample_mode: {mode}") + + # Per-channel normalization: normalized = (x - mean) / std + mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latent) + std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latent) + latent = (latent - mean) / std + else: + latent = latent.to(dtype=encode_dtype) + + packed = server_args.pipeline_config.maybe_pack_latents( + latent, latent.shape[0], batch + ) + if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): + raise ValueError("Expected packed image latents [B, S0, D].") + + # Fail-fast token count: must match one latent frame's tokens. + vae_sf = int(server_args.pipeline_config.vae_scale_factor) + patch = int(server_args.pipeline_config.patch_size) + latent_h = int(batch.height) // vae_sf + latent_w = int(batch.width) // vae_sf + expected_tokens = (latent_h // patch) * (latent_w // patch) + if int(packed.shape[1]) != int(expected_tokens): + raise ValueError( + "LTX-2 conditioning token count mismatch: " + f"{int(packed.shape[1])=} {int(expected_tokens)=}." + ) + + batch.image_latent = packed + batch.ltx2_num_image_tokens = int(packed.shape[1]) + + if batch.debug: + logger.info( + "LTX2 TI2V conditioning prepared: %d tokens (shape=%s) for %sx%s", + batch.ltx2_num_image_tokens, + tuple(batch.image_latent.shape), + batch.width, + batch.height, + ) + + if condition_image_encoder is None: + self.vae.to(original_dtype) + if server_args.vae_cpu_offload: + self.vae = self.vae.to("cpu") + if condition_image_encoder is not None: + self._condition_image_encoder = condition_image_encoder.to("cpu") + + def _prepare_denoising_loop( + self, + batch: Req, + server_args: ServerArgs, + ) -> LTX2DenoisingContext: + """Extend the base context with LTX-2 audio, SP, and TI2V state.""" + self._disable_cache_dit_for_request = batch.image_path is not None + base_ctx = super()._prepare_denoising_loop(batch, server_args) + ctx = LTX2DenoisingContext(**base_ctx.to_kwargs()) + ctx.is_ltx23_variant = is_ltx23_native_variant( + server_args.pipeline_config.vae_config.arch_config + ) + phase = batch.extra.get("ltx2_phase") + pipeline = self.pipeline() if self.pipeline else None + pipeline_name = pipeline.pipeline_name if pipeline is not None else None + ctx.use_ltx23_legacy_one_stage = self._should_use_ltx23_legacy_one_stage( + server_args, pipeline_name + ) + ctx.stage = ( + phase + if phase is not None + else ("stage1" if ctx.use_ltx23_legacy_one_stage else "one_stage") + ) + ctx.audio_latents = batch.audio_latents + # Video and audio keep separate scheduler state throughout the denoising loop. + ctx.audio_scheduler = copy.deepcopy(self.scheduler) + + # Prepare image latents and embeddings for LTX-2 TI2V generation. + self._prepare_ltx2_image_latent(batch, server_args) + do_ti2v = self._should_apply_ltx2_ti2v(batch) + + if ctx.use_ltx23_legacy_one_stage: + batch.ltx23_audio_replicated_for_sp = False + batch.did_sp_shard_audio_latents = False + else: + ctx.replicate_audio_for_sp = self._should_replicate_ltx23_audio_for_sp( + batch, + server_args, + is_ltx23_variant=ctx.is_ltx23_variant, + ) + batch.ltx23_audio_replicated_for_sp = bool(ctx.replicate_audio_for_sp) + if ( + ctx.is_ltx23_variant + and get_sp_world_size() > 1 + and server_args.pipeline_config.can_shard_audio_latents_for_sp( + batch.audio_latents + ) + and not ctx.replicate_audio_for_sp + ): + ( + batch.audio_latents, + batch.did_sp_shard_audio_latents, + ) = server_args.pipeline_config.shard_audio_latents_for_sp( + batch, batch.audio_latents + ) + ctx.audio_latents = batch.audio_latents + else: + batch.did_sp_shard_audio_latents = False + + # For LTX-2 packed token latents, SP sharding happens on the time dimension + # (frames). The model must see local latent frames (RoPE offset is applied + # inside the model using SP rank). + ctx.latent_num_frames_for_model = self._get_video_latent_num_frames_for_model( + batch=batch, server_args=server_args, latents=ctx.latents + ) + ctx.latent_height = ( + batch.height + // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + ) + ctx.latent_width = ( + batch.width + // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + ) + if do_ti2v: + if not (isinstance(ctx.latents, torch.Tensor) and ctx.latents.ndim == 3): + raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") + # Keep conditioned tokens clean and reuse the mask during every step update. + ctx.latents, ctx.denoise_mask, ctx.clean_latent = ( + self._prepare_ltx2_ti2v_clean_state( + latents=ctx.latents, + image_latent=batch.image_latent, + num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)), + zero_clean_latent=ctx.is_ltx23_variant, + ) + ) + return ctx + + def _before_denoising_loop( + self, ctx: LTX2DenoisingContext, batch: Req, server_args: ServerArgs + ) -> None: + """Reset the mirrored audio scheduler before the shared loop begins.""" + super()._before_denoising_loop(ctx, batch, server_args) + if ctx.audio_scheduler is None: + raise ValueError("LTX-2 audio scheduler was not prepared.") + ctx.audio_scheduler.set_begin_index(0) + + def _prepare_step_attn_metadata( + self, + ctx: LTX2DenoisingContext, + batch: Req, + server_args: ServerArgs, + step_index: int, + t_int: int, + timesteps_cpu: torch.Tensor, + ): + """Preserve the legacy LTX-2 attention-metadata contract.""" + # Legacy LTX-2 paths used the plain attention-metadata builder call here. + del ctx, t_int, timesteps_cpu + return self._build_attn_metadata(step_index, batch, server_args) + + def _run_denoising_step( + self, + ctx: LTX2DenoisingContext, + step: DenoisingStepState, + batch: Req, + server_args: ServerArgs, + ) -> None: + """Run one joint video/audio denoising step with LTX-2-specific guidance.""" + if ctx.audio_latents is None: + raise ValueError("LTX-2 requires audio latents for denoising.") + if ctx.audio_scheduler is None: + raise ValueError("LTX-2 audio scheduler was not prepared.") + + # 1. Read the scheduler sigma pair and derive the Euler delta. + sigmas = getattr(self.scheduler, "sigmas", None) + if sigmas is None or not isinstance(sigmas, torch.Tensor): + raise ValueError("Expected scheduler.sigmas to be a tensor for LTX-2.") + sigma = sigmas[step.step_index].to( + device=ctx.latents.device, dtype=torch.float32 + ) + sigma_next = sigmas[step.step_index + 1].to( + device=ctx.latents.device, dtype=torch.float32 + ) + dt = sigma_next - sigma + + # 2. Materialize the current video/audio latent inputs in the compute dtype. + latent_model_input = ctx.latents.to(ctx.target_dtype) + audio_latent_model_input = ctx.audio_latents.to(ctx.target_dtype) + stage1_guider_params = self._get_ltx2_stage1_guider_params( + batch, server_args, ctx.stage + ) + + if audio_latent_model_input.ndim == 3: + audio_num_frames_latent = int(audio_latent_model_input.shape[1]) + elif audio_latent_model_input.ndim == 4: + audio_num_frames_latent = int(audio_latent_model_input.shape[2]) + else: + raise ValueError( + f"Unexpected audio latents rank: {audio_latent_model_input.ndim}, shape={tuple(audio_latent_model_input.shape)}" + ) + + # 3. Prepare any LTX-specific RoPE coordinates and timestep layouts. + video_coords = None + audio_coords = None + if not ctx.use_ltx23_legacy_one_stage: + video_coords = server_args.pipeline_config.prepare_video_rope_coords_for_sp( + step.current_model, + batch, + latent_model_input, + num_frames=ctx.latent_num_frames_for_model, + height=ctx.latent_height, + width=ctx.latent_width, + ) + audio_coords = server_args.pipeline_config.prepare_audio_rope_coords_for_sp( + step.current_model, + batch, + audio_latent_model_input, + num_frames=audio_num_frames_latent, + ) + + batch_size = int(latent_model_input.shape[0]) + timestep = step.t_device.expand(batch_size) + if ctx.denoise_mask is not None: + timestep_video = timestep.unsqueeze(-1) * ctx.denoise_mask.squeeze(-1) + elif ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage: + timestep_video = timestep.view(batch_size, 1).expand( + batch_size, int(latent_model_input.shape[1]) + ) + else: + timestep_video = timestep + + if ( + ctx.is_ltx23_variant + and not ctx.use_ltx23_legacy_one_stage + and audio_latent_model_input.ndim == 3 + ): + timestep_audio = timestep.view(batch_size, 1).expand( + batch_size, int(audio_latent_model_input.shape[1]) + ) + else: + timestep_audio = timestep + + prompt_timestep_video = None + prompt_timestep_audio = None + if ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage: + timestep_scale_multiplier = float( + getattr(step.current_model, "timestep_scale_multiplier", 1000) + ) + prompt_timestep_video = ( + sigma.to(device=latent_model_input.device, dtype=torch.float32) + * timestep_scale_multiplier + ).expand(batch_size) + prompt_timestep_audio = ( + sigma.to(device=audio_latent_model_input.device, dtype=torch.float32) + * timestep_scale_multiplier + ).expand(batch_size) + + # 4. Build attention masks that account for SP padding and replicated audio. + if ctx.use_ltx23_legacy_one_stage: + video_self_attention_mask = None + audio_self_attention_mask = None + a2v_cross_attention_mask = None + v2a_cross_attention_mask = None + else: + video_self_attention_mask = self._build_ltx2_sp_padding_mask( + batch, + seq_len=int(latent_model_input.shape[1]), + batch_size=batch_size, + key="sp_video_valid_token_count", + device=latent_model_input.device, + ) + audio_self_attention_mask = self._build_ltx2_sp_padding_mask( + batch, + seq_len=audio_num_frames_latent, + batch_size=batch_size, + key="sp_audio_valid_token_count", + device=audio_latent_model_input.device, + ) + a2v_cross_attention_mask = audio_self_attention_mask + v2a_cross_attention_mask = video_self_attention_mask + + def build_model_kwargs( + *, + encoder_hidden_states: torch.Tensor, + audio_encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None, + skip_video_self_attn_blocks: tuple[int, ...] | None = None, + skip_audio_self_attn_blocks: tuple[int, ...] | None = None, + disable_a2v_cross_attn: bool = False, + disable_v2a_cross_attn: bool = False, + ) -> dict[str, object]: + kwargs: dict[str, object] = { + "hidden_states": latent_model_input, + "audio_hidden_states": audio_latent_model_input, + "encoder_hidden_states": encoder_hidden_states, + "audio_encoder_hidden_states": audio_encoder_hidden_states, + "timestep": timestep_video, + "audio_timestep": timestep_audio, + "encoder_attention_mask": encoder_attention_mask, + "audio_encoder_attention_mask": encoder_attention_mask, + "num_frames": ctx.latent_num_frames_for_model, + "height": ctx.latent_height, + "width": ctx.latent_width, + "fps": batch.fps, + "audio_num_frames": audio_num_frames_latent, + "video_coords": video_coords, + "audio_coords": audio_coords, + "return_latents": False, + "return_dict": False, + } + if not ctx.use_ltx23_legacy_one_stage: + kwargs.update( + { + "prompt_timestep": prompt_timestep_video, + "audio_prompt_timestep": prompt_timestep_audio, + "video_self_attention_mask": video_self_attention_mask, + "audio_self_attention_mask": audio_self_attention_mask, + "a2v_cross_attention_mask": a2v_cross_attention_mask, + "v2a_cross_attention_mask": v2a_cross_attention_mask, + "audio_replicated_for_sp": ctx.replicate_audio_for_sp, + "legacy_ltx23_one_stage_semantics": False, + } + ) + if skip_video_self_attn_blocks is not None: + kwargs["skip_video_self_attn_blocks"] = skip_video_self_attn_blocks + if skip_audio_self_attn_blocks is not None: + kwargs["skip_audio_self_attn_blocks"] = skip_audio_self_attn_blocks + if disable_a2v_cross_attn: + kwargs["disable_a2v_cross_attn"] = True + if disable_v2a_cross_attn: + kwargs["disable_v2a_cross_attn"] = True + return kwargs + + # 5. Run the branch-specific LTX forward path and apply CFG/guider logic. + prompt_attention_mask = self._get_ltx_prompt_attention_mask( + batch, + is_ltx23_variant=( + ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage + ), + ) + use_official_cfg_path = stage1_guider_params is None + if use_official_cfg_path: + encoder_hidden_states = batch.prompt_embeds[0] + audio_encoder_hidden_states = batch.audio_prompt_embeds[0] + encoder_attention_mask = prompt_attention_mask + if batch.do_classifier_free_guidance: + latent_model_input = torch.cat([latent_model_input] * 2, dim=0) + audio_latent_model_input = torch.cat( + [audio_latent_model_input] * 2, dim=0 + ) + encoder_hidden_states = torch.cat( + [batch.negative_prompt_embeds[0], encoder_hidden_states], dim=0 + ) + audio_encoder_hidden_states = torch.cat( + [ + batch.negative_audio_prompt_embeds[0], + audio_encoder_hidden_states, + ], + dim=0, + ) + if encoder_attention_mask is not None: + encoder_attention_mask = torch.cat( + [ + self._get_ltx_prompt_attention_mask( + batch, + is_ltx23_variant=( + ctx.is_ltx23_variant + and not ctx.use_ltx23_legacy_one_stage + ), + negative=True, + ), + encoder_attention_mask, + ], + dim=0, + ) + cfg_batch_size = int(latent_model_input.shape[0]) + timestep_video = self._repeat_batch_dim(timestep_video, cfg_batch_size) + timestep_audio = self._repeat_batch_dim(timestep_audio, cfg_batch_size) + if prompt_timestep_video is not None: + prompt_timestep_video = self._repeat_batch_dim( + prompt_timestep_video, cfg_batch_size + ) + if prompt_timestep_audio is not None: + prompt_timestep_audio = self._repeat_batch_dim( + prompt_timestep_audio, cfg_batch_size + ) + if video_self_attention_mask is not None: + video_self_attention_mask = self._repeat_batch_dim( + video_self_attention_mask, cfg_batch_size + ) + if audio_self_attention_mask is not None: + audio_self_attention_mask = self._repeat_batch_dim( + audio_self_attention_mask, cfg_batch_size + ) + if a2v_cross_attention_mask is not None: + a2v_cross_attention_mask = self._repeat_batch_dim( + a2v_cross_attention_mask, cfg_batch_size + ) + if v2a_cross_attention_mask is not None: + v2a_cross_attention_mask = self._repeat_batch_dim( + v2a_cross_attention_mask, cfg_batch_size + ) + + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + model_video, model_audio = step.current_model( + **build_model_kwargs( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ) + ) + + model_video = model_video.float() + model_audio = model_audio.float() + if batch.do_classifier_free_guidance: + model_video_uncond, model_video_text = model_video.chunk(2) + model_audio_uncond, model_audio_text = model_audio.chunk(2) + model_video = model_video_uncond + ( + batch.guidance_scale * (model_video_text - model_video_uncond) + ) + model_audio = model_audio_uncond + ( + batch.guidance_scale * (model_audio_text - model_audio_uncond) + ) + + ctx.latents = self.scheduler.step( + model_video, step.t_device, ctx.latents, return_dict=False + )[0] + ctx.audio_latents = ctx.audio_scheduler.step( + model_audio, step.t_device, ctx.audio_latents, return_dict=False + )[0] + ctx.latents = self.post_forward_for_ti2v_task( + batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z + ) + return + + encoder_hidden_states = batch.prompt_embeds[0] + audio_encoder_hidden_states = batch.audio_prompt_embeds[0] + encoder_attention_mask = prompt_attention_mask + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + v_pos, a_v_pos = step.current_model( + **build_model_kwargs( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ) + ) + + if stage1_guider_params is not None or batch.do_classifier_free_guidance: + v_neg, a_v_neg = step.current_model( + **build_model_kwargs( + encoder_hidden_states=batch.negative_prompt_embeds[0], + audio_encoder_hidden_states=batch.negative_audio_prompt_embeds[ + 0 + ], + encoder_attention_mask=self._get_ltx_prompt_attention_mask( + batch, + is_ltx23_variant=( + ctx.is_ltx23_variant + and not ctx.use_ltx23_legacy_one_stage + ), + negative=True, + ), + ) + ) + else: + v_neg = None + a_v_neg = None + + v_pos = v_pos.float() + a_v_pos = a_v_pos.float() + if v_neg is not None: + v_neg = v_neg.float() + if a_v_neg is not None: + a_v_neg = a_v_neg.float() + + sigma_val = float(sigma.item()) + video_sigma_for_x0: float | torch.Tensor = sigma_val + if ctx.denoise_mask is not None: + video_sigma_for_x0 = sigma.to( + device=ctx.latents.device, dtype=torch.float32 + ) * ctx.denoise_mask.squeeze(-1) + + denoised_video = self._ltx2_velocity_to_x0( + ctx.latents, v_pos, video_sigma_for_x0 + ) + denoised_audio = self._ltx2_velocity_to_x0( + ctx.audio_latents, a_v_pos, sigma_val + ) + denoised_video_neg = None + denoised_audio_neg = None + denoised_video_perturbed = None + denoised_audio_perturbed = None + denoised_video_modality = None + denoised_audio_modality = None + + if ( + (stage1_guider_params is not None or batch.do_classifier_free_guidance) + and v_neg is not None + and a_v_neg is not None + ): + denoised_video_neg = self._ltx2_velocity_to_x0( + ctx.latents, v_neg, video_sigma_for_x0 + ) + denoised_audio_neg = self._ltx2_velocity_to_x0( + ctx.audio_latents, a_v_neg, sigma_val + ) + if stage1_guider_params is not None: + video_skip = self._ltx2_should_skip_step( + step.step_index, int(stage1_guider_params["video_skip_step"]) + ) + audio_skip = self._ltx2_should_skip_step( + step.step_index, int(stage1_guider_params["audio_skip_step"]) + ) + + need_perturbed = ( + float(stage1_guider_params["video_stg_scale"]) != 0.0 + or float(stage1_guider_params["audio_stg_scale"]) != 0.0 + ) + if need_perturbed: + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + v_ptb, a_v_ptb = step.current_model( + **build_model_kwargs( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + skip_video_self_attn_blocks=tuple( + stage1_guider_params["video_stg_blocks"] + ), + skip_audio_self_attn_blocks=tuple( + stage1_guider_params["audio_stg_blocks"] + ), + ) + ) + denoised_video_perturbed = self._ltx2_velocity_to_x0( + ctx.latents, v_ptb.float(), video_sigma_for_x0 + ) + denoised_audio_perturbed = self._ltx2_velocity_to_x0( + ctx.audio_latents, a_v_ptb.float(), sigma_val + ) + + need_modality = ( + float(stage1_guider_params["video_modality_scale"]) != 1.0 + or float(stage1_guider_params["audio_modality_scale"]) != 1.0 + ) + if need_modality: + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + v_mod, a_v_mod = step.current_model( + **build_model_kwargs( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + disable_a2v_cross_attn=True, + disable_v2a_cross_attn=True, + ) + ) + denoised_video_modality = self._ltx2_velocity_to_x0( + ctx.latents, v_mod.float(), video_sigma_for_x0 + ) + denoised_audio_modality = self._ltx2_velocity_to_x0( + ctx.audio_latents, a_v_mod.float(), sigma_val + ) + + if not video_skip: + denoised_video = self._ltx2_calculate_guided_x0( + cond=denoised_video, + uncond_text=( + denoised_video_neg + if denoised_video_neg is not None + else denoised_video + ), + uncond_perturbed=( + denoised_video_perturbed + if denoised_video_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_video_modality + if denoised_video_modality is not None + else 0.0 + ), + cfg_scale=float(stage1_guider_params["video_cfg_scale"]), + stg_scale=float(stage1_guider_params["video_stg_scale"]), + rescale_scale=float(stage1_guider_params["video_rescale_scale"]), + modality_scale=float(stage1_guider_params["video_modality_scale"]), + ) + if not audio_skip: + denoised_audio = self._ltx2_calculate_guided_x0( + cond=denoised_audio, + uncond_text=( + denoised_audio_neg + if denoised_audio_neg is not None + else denoised_audio + ), + uncond_perturbed=( + denoised_audio_perturbed + if denoised_audio_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_audio_modality + if denoised_audio_modality is not None + else 0.0 + ), + cfg_scale=float(stage1_guider_params["audio_cfg_scale"]), + stg_scale=float(stage1_guider_params["audio_stg_scale"]), + rescale_scale=float(stage1_guider_params["audio_rescale_scale"]), + modality_scale=float(stage1_guider_params["audio_modality_scale"]), + ) + elif ( + batch.do_classifier_free_guidance + and denoised_video_neg is not None + and denoised_audio_neg is not None + ): + denoised_video = denoised_video + (batch.guidance_scale - 1.0) * ( + denoised_video - denoised_video_neg + ) + denoised_audio = denoised_audio + (batch.guidance_scale - 1.0) * ( + denoised_audio - denoised_audio_neg + ) + + if ctx.denoise_mask is not None and ctx.clean_latent is not None: + denoised_video = ( + denoised_video * ctx.denoise_mask + + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask) + ) + + # 6. Convert x0 predictions back to velocity and update both latent streams. + if sigma_val == 0.0: + v_video = torch.zeros_like(denoised_video) + v_audio = torch.zeros_like(denoised_audio) + else: + v_video = ((ctx.latents.float() - denoised_video.float()) / sigma_val).to( + ctx.latents.dtype + ) + v_audio = ( + (ctx.audio_latents.float() - denoised_audio.float()) / sigma_val + ).to(ctx.audio_latents.dtype) + + ctx.latents = (ctx.latents.float() + v_video.float() * dt).to( + dtype=ctx.latents.dtype + ) + ctx.audio_latents = (ctx.audio_latents.float() + v_audio.float() * dt).to( + dtype=ctx.audio_latents.dtype + ) + ctx.latents = self.post_forward_for_ti2v_task( + batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z + ) + + def _record_trajectory( + self, + ctx: LTX2DenoisingContext, + step: DenoisingStepState, + batch: Req, + server_args: ServerArgs, + ) -> None: + """Record audio trajectory alongside the base video trajectory.""" + super()._record_trajectory(ctx, step, batch, server_args) + if batch.return_trajectory_latents and ctx.audio_latents is not None: + ctx.trajectory_audio_latents.append(ctx.audio_latents) + + def _finalize_denoising_loop( + self, ctx: LTX2DenoisingContext, batch: Req, server_args: ServerArgs + ) -> None: + """Expose audio latents before delegating to AV-aware postprocessing.""" + batch.audio_latents = ctx.audio_latents + self._post_denoising_loop( + batch=batch, + latents=ctx.latents, + trajectory_latents=ctx.trajectory_latents, + trajectory_timesteps=ctx.trajectory_timesteps, + trajectory_audio_latents=ctx.trajectory_audio_latents, + server_args=server_args, + is_warmup=ctx.is_warmup, + ) + + def _post_denoising_loop( + self, + batch: Req, + latents: torch.Tensor, + trajectory_latents: list, + trajectory_timesteps: list, + server_args: ServerArgs, + trajectory_audio_latents: list | None = None, + is_warmup: bool = False, + *args, + **kwargs, + ): + """Trim SP token padding before delegating to the base finalizer.""" + if trajectory_audio_latents: + batch.trajectory_audio_latents = torch.stack( + trajectory_audio_latents, dim=1 + ).cpu() + latents = self._truncate_sp_padded_token_latents(batch, latents) + super()._post_denoising_loop( + batch=batch, + latents=latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + server_args=server_args, + is_warmup=is_warmup, + ) + + def _get_prompt_embeds_validator(self, batch: Req): + """Allow either tensor or list prompt embeddings for LTX-2 prompts.""" + del batch + return lambda x: V.is_tensor(x) or V.list_not_empty(x) + + def _get_negative_prompt_embeds_validator(self, batch: Req): + """Allow either tensor or list negative prompt embeddings for LTX-2 CFG.""" + return ( + lambda x: (not batch.do_classifier_free_guidance) + or V.is_tensor(x) + or V.list_not_empty(x) + ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py new file mode 100644 index 000000000..7ac2d6cbe --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py @@ -0,0 +1 @@ +"""Model-specific helpers and stages for diffusion pipeline components.""" diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py new file mode 100644 index 000000000..b26d00fa6 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py @@ -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 diff --git a/scripts/ci/utils/diffusion/comparison_configs.json b/scripts/ci/utils/diffusion/comparison_configs.json index b1b766591..db51d53fa 100644 --- a/scripts/ci/utils/diffusion/comparison_configs.json +++ b/scripts/ci/utils/diffusion/comparison_configs.json @@ -135,6 +135,23 @@ } } }, + { + "id": "ltx2.3_twostage_t2v_2gpus", + "model": "Lightricks/LTX-2.3", + "task": "text-to-video", + "prompt": "A cat and a dog baking a cake together in a kitchen.", + "width": 768, + "height": 512, + "num_frames": 121, + "seed": 42, + "num_gpus": 2, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --pipeline-class-name LTX2TwoStagePipeline", + "extra_env": {} + } + } + }, { "id": "wan22_i2v_a14b_720p", "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",