From c5e95080d21bd2619348f2662ad86cda99663fb0 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 14 Apr 2026 22:10:08 +0800 Subject: [PATCH] [diffusion] model: support Ltx 2.3 two stage ti2v (#22667) --- .../configs/pipeline_configs/ltx_2.py | 18 + .../configs/sample/sampling_params.py | 3 +- .../loader/component_loaders/vae_loader.py | 14 + .../runtime/models/adapter/ltx_2_connector.py | 3 +- .../runtime/models/dits/ltx_2.py | 154 ++++- .../runtime/pipelines/ltx_2_pipeline.py | 5 + .../pipelines_core/stages/denoising_av.py | 57 +- .../pipelines_core/stages/ltx_2_denoising.py | 601 +++++++++++++----- .../test/server/accuracy_config.py | 13 - .../test/server/consistency_threshold.json | 12 - .../test/server/perf_baselines.json | 58 +- .../test/server/testcase_configs.py | 36 +- .../test/unit/test_sampling_params.py | 33 + .../test/unit/test_vae_loader.py | 48 ++ .../utils/diffusion/comparison_configs.json | 10 +- 15 files changed, 837 insertions(+), 228 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_vae_loader.py 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 00b5dc80c..03478f9ae 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -118,6 +118,24 @@ def is_ltx23_native_variant(arch_config: object) -> bool: return str(getattr(arch_config, "ltx_variant", "ltx_2")) == "ltx_2_3" +def sync_ltx23_runtime_vae_markers( + arch_config: object, + loaded_vae_config: object | None, +) -> None: + if loaded_vae_config is None: + return + source = getattr(loaded_vae_config, "arch_config", loaded_vae_config) + for key in ( + "ltx_variant", + "condition_encoder_subdir", + "video_decoder_variant", + "video_decoder_config", + ): + value = getattr(source, key, None) + if value is not None: + setattr(arch_config, key, value) + + def _gemma_postprocess_func( outputs: BaseEncoderOutput, text_inputs: dict, diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 33ffc8387..3fae7d089 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -609,6 +609,7 @@ class SamplingParams: user_kwargs = dict(kwargs) user_kwargs.pop("diffusers_kwargs", None) + user_sampling_params = SamplingParams(*args, **user_kwargs) # TODO: refactor sampling_params._merge_with_user_params( @@ -793,7 +794,7 @@ class SamplingParams: "--cfg-normalization", type=float, dest="cfg_normalization", - help=("CFG renormalization factor (for Z-Image). "), + help="CFG renormalization factor (for Z-Image). ", ) add_argument( "--boundary-ratio", diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index c49aa4d67..2ff38095e 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -27,6 +27,19 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) +def _backfill_ltx2_audio_vae_latent_stats( + loaded: dict[str, torch.Tensor], component_name: str +) -> None: + if component_name != "audio_vae": + return + mean_key = "per_channel_statistics.mean-of-means" + std_key = "per_channel_statistics.std-of-means" + if "latents_mean" not in loaded and mean_key in loaded: + loaded["latents_mean"] = loaded[mean_key] + if "latents_std" not in loaded and std_key in loaded: + loaded["latents_std"] = loaded[std_key] + + def _convert_conv3d_weights_to_channels_last_3d(module: nn.Module) -> int: """ Convert Conv3d weights to channels_last_3d (NDHWC) memory format. @@ -142,6 +155,7 @@ class VAELoader(ComponentLoader): loaded = {} for sf_path in safetensors_list: loaded.update(safetensors_load_file(sf_path)) + _backfill_ltx2_audio_vae_latent_stats(loaded, component_name) vae.load_state_dict(loaded, strict=False) state_keys = set(vae.state_dict().keys()) diff --git a/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py b/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py index e629e2aad..cde2357c4 100644 --- a/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py +++ b/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py @@ -517,6 +517,7 @@ class LTX2TextConnectors(nn.Module): ): super().__init__() caption_channels = config.caption_channels + self.caption_channels = caption_channels text_proj_in_factor = config.text_proj_in_factor video_connector_num_attention_heads = config.video_connector_num_attention_heads video_connector_attention_head_dim = config.video_connector_attention_head_dim @@ -641,7 +642,7 @@ class LTX2TextConnectors(nn.Module): audio_hidden_states = audio_hidden_states.to( self.audio_aggregate_embed.weight.dtype ) - source_dim = self.video_aggregate_embed.out_features + source_dim = self.caption_channels video_hidden_states = self._rescale_v2_features( video_hidden_states, self.video_aggregate_embed.out_features, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 18dee2ddf..84bd3cd81 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -48,13 +48,54 @@ def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int: ) +def _ltx2_is_perturbed( + perturbation_config: dict[str, object], + key: str, + block_idx: int, +) -> bool: + value = perturbation_config.get(key) + if value is None: + return False + if key.endswith("_blocks"): + return block_idx in value + return bool(value) + + +def _ltx2_batched_perturbation_mask( + perturbation_configs: tuple[dict[str, object], ...] | None, + key: str, + block_idx: int, + values: torch.Tensor, +) -> tuple[torch.Tensor | None, bool]: + if not perturbation_configs: + return None, False + + mask = torch.ones( + (len(perturbation_configs),), device=values.device, dtype=values.dtype + ) + any_perturbed = False + all_perturbed = True + for batch_idx, config in enumerate(perturbation_configs): + perturbed = _ltx2_is_perturbed(config, key, block_idx) + any_perturbed = any_perturbed or perturbed + all_perturbed = all_perturbed and perturbed + if perturbed: + mask[batch_idx] = 0 + + if not any_perturbed: + return None, False + if all_perturbed: + return None, True + return mask.view(mask.numel(), *([1] * (values.ndim - 1))), False + + def apply_interleaved_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: cos, sin = freqs x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) - return (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + return x * cos + x_rotated * sin def apply_split_rotary_emb( @@ -76,7 +117,7 @@ def apply_split_rotary_emb( ) r = last // 2 - split_x = x.reshape(*x.shape[:-1], 2, r).float() + split_x = x.reshape(*x.shape[:-1], 2, r) first_x = split_x[..., :1, :] second_x = split_x[..., 1:, :] @@ -248,9 +289,13 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): return self.prepare_audio_coords(*args, **kwargs) def forward( - self, coords: torch.Tensor, device: Optional[Union[str, torch.device]] = None + self, + coords: torch.Tensor, + device: Optional[Union[str, torch.device]] = None, + out_dtype: Optional[torch.dtype] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: device = device or coords.device + out_dtype = out_dtype or coords.dtype num_pos_dims = coords.shape[1] if coords.ndim == 4: @@ -317,7 +362,7 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): cos_freqs = torch.swapaxes(cos_freq, 1, 2) sin_freqs = torch.swapaxes(sin_freq, 1, 2) - return cos_freqs, sin_freqs + return cos_freqs.to(dtype=out_dtype), sin_freqs.to(dtype=out_dtype) def rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor: @@ -653,6 +698,8 @@ class LTX2Attention(nn.Module): ) if perturbation_mask is not None: + if perturbation_mask.ndim == out.ndim - 1: + perturbation_mask = perturbation_mask.unsqueeze(-1) out = out * perturbation_mask + v * (1 - perturbation_mask) if not use_attention: @@ -918,6 +965,10 @@ class LTX2TransformerBlock(nn.Module): skip_audio_self_attn: bool = False, skip_a2v_cross_attn: bool = False, skip_v2a_cross_attn: bool = False, + video_self_attn_perturbation_mask: Optional[torch.Tensor] = None, + audio_self_attn_perturbation_mask: Optional[torch.Tensor] = None, + a2v_cross_attn_perturbation_mask: Optional[torch.Tensor] = None, + v2a_cross_attn_perturbation_mask: Optional[torch.Tensor] = None, audio_replicated_for_sp: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -934,6 +985,7 @@ class LTX2TransformerBlock(nn.Module): norm_hidden_states, mask=video_self_attention_mask, pe=video_rotary_emb, + perturbation_mask=video_self_attn_perturbation_mask, all_perturbed=skip_video_self_attn, gather_context_kv_for_sp=audio_replicated_for_sp, ) @@ -949,6 +1001,7 @@ class LTX2TransformerBlock(nn.Module): norm_audio_hidden_states, mask=audio_self_attention_mask, pe=audio_rotary_emb, + perturbation_mask=audio_self_attn_perturbation_mask, all_perturbed=skip_audio_self_attn, skip_sequence_parallel_override=audio_replicated_for_sp, ) @@ -1097,6 +1150,10 @@ class LTX2TransformerBlock(nn.Module): mask=a2v_cross_attention_mask, skip_sequence_parallel_override=audio_replicated_for_sp, ) + if a2v_cross_attn_perturbation_mask is not None: + a2v_attn_hidden_states = ( + a2v_attn_hidden_states * a2v_cross_attn_perturbation_mask + ) hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states # V2A @@ -1116,6 +1173,10 @@ class LTX2TransformerBlock(nn.Module): mask=v2a_cross_attention_mask, gather_context_kv_for_sp=audio_replicated_for_sp, ) + if v2a_cross_attn_perturbation_mask is not None: + v2a_attn_hidden_states = ( + v2a_attn_hidden_states * v2a_cross_attn_perturbation_mask + ) audio_hidden_states = ( audio_hidden_states + v2a_gate * v2a_attn_hidden_states ) @@ -1532,6 +1593,12 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): raise ValueError( "audio_num_frames must be provided for RoPE coordinate generation." ) + perturbation_configs = kwargs.get("perturbation_configs") + if perturbation_configs is not None and len(perturbation_configs) != batch_size: + raise ValueError( + "perturbation_configs length must match batch size, got " + f"{len(perturbation_configs)=} {batch_size=}." + ) if video_coords is None: # Wan-style SP-RoPE: when SP is enabled, each rank runs on its local @@ -1566,15 +1633,25 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): video_coords, hidden_states.device, hidden_states.dtype ) audio_coords = audio_coords.to(device=audio_hidden_states.device) - video_rotary_emb = self.rope(video_coords, device=hidden_states.device) + video_rotary_emb = self.rope( + video_coords, + device=hidden_states.device, + out_dtype=hidden_states.dtype, + ) audio_rotary_emb = self.audio_rope( - audio_coords, device=audio_hidden_states.device + audio_coords, + device=audio_hidden_states.device, + out_dtype=audio_hidden_states.dtype, ) ca_video_rotary_emb = self.cross_attn_rope( - video_coords[:, 0:1, :], device=hidden_states.device + video_coords[:, 0:1, :], + device=hidden_states.device, + out_dtype=hidden_states.dtype, ) ca_audio_rotary_emb = self.cross_attn_audio_rope( - audio_coords[:, 0:1, :], device=audio_hidden_states.device + audio_coords[:, 0:1, :], + device=audio_hidden_states.device, + out_dtype=audio_hidden_states.dtype, ) # 2. Patchify input projections @@ -1673,6 +1750,55 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ()) skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ()) for block in self.transformer_blocks: + video_self_attn_perturbation_mask = None + audio_self_attn_perturbation_mask = None + a2v_cross_attn_perturbation_mask = None + v2a_cross_attn_perturbation_mask = None + skip_video_self_attn = block.idx in skip_video_self_attn_blocks + skip_audio_self_attn = block.idx in skip_audio_self_attn_blocks + skip_a2v_cross_attn = disable_a2v_cross_attn + skip_v2a_cross_attn = disable_v2a_cross_attn + if perturbation_configs is not None: + if not skip_video_self_attn: + ( + video_self_attn_perturbation_mask, + skip_video_self_attn, + ) = _ltx2_batched_perturbation_mask( + perturbation_configs, + "skip_video_self_attn_blocks", + block.idx, + hidden_states, + ) + if not skip_audio_self_attn: + ( + audio_self_attn_perturbation_mask, + skip_audio_self_attn, + ) = _ltx2_batched_perturbation_mask( + perturbation_configs, + "skip_audio_self_attn_blocks", + block.idx, + audio_hidden_states, + ) + if not skip_a2v_cross_attn: + ( + a2v_cross_attn_perturbation_mask, + skip_a2v_cross_attn, + ) = _ltx2_batched_perturbation_mask( + perturbation_configs, + "skip_a2v_cross_attn", + block.idx, + hidden_states, + ) + if not skip_v2a_cross_attn: + ( + v2a_cross_attn_perturbation_mask, + skip_v2a_cross_attn, + ) = _ltx2_batched_perturbation_mask( + perturbation_configs, + "skip_v2a_cross_attn", + block.idx, + audio_hidden_states, + ) hidden_states, audio_hidden_states = block( hidden_states, audio_hidden_states, @@ -1699,10 +1825,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): audio_self_attention_mask=audio_self_attention_mask, a2v_cross_attention_mask=a2v_cross_attention_mask, v2a_cross_attention_mask=v2a_cross_attention_mask, - skip_video_self_attn=block.idx in skip_video_self_attn_blocks, - skip_audio_self_attn=block.idx in skip_audio_self_attn_blocks, - skip_a2v_cross_attn=disable_a2v_cross_attn, - skip_v2a_cross_attn=disable_v2a_cross_attn, + skip_video_self_attn=skip_video_self_attn, + skip_audio_self_attn=skip_audio_self_attn, + skip_a2v_cross_attn=skip_a2v_cross_attn, + skip_v2a_cross_attn=skip_v2a_cross_attn, + video_self_attn_perturbation_mask=video_self_attn_perturbation_mask, + audio_self_attn_perturbation_mask=audio_self_attn_perturbation_mask, + a2v_cross_attn_perturbation_mask=a2v_cross_attn_perturbation_mask, + v2a_cross_attn_perturbation_mask=v2a_cross_attn_perturbation_mask, audio_replicated_for_sp=audio_replicated_for_sp, ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index 406a9e9c1..2c7763207 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -7,6 +7,7 @@ from diffusers import FlowMatchEulerDiscreteScheduler from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( is_ltx23_native_variant, + sync_ltx23_runtime_vae_markers, ) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PipelineComponentLoader, @@ -249,6 +250,10 @@ class _BaseLTX2Pipeline(LoRAPipeline): def initialize_pipeline(self, server_args: ServerArgs): orig = self.get_module("scheduler") self.modules["scheduler"] = LTX2FlowMatchScheduler.from_config(orig.config) + sync_ltx23_runtime_vae_markers( + server_args.pipeline_config.vae_config.arch_config, + getattr(self.get_module("vae"), "config", None), + ) class LTX2Pipeline(_BaseLTX2Pipeline): 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 e735960bd..3fe53cd58 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 @@ -3,9 +3,7 @@ import copy import torch from diffusers.utils.torch_utils import randn_tensor -from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( - is_ltx23_native_variant, -) +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import is_ltx23_native_variant from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( LTX2DenoisingStage, @@ -156,23 +154,59 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): 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" + original_clean_latent_background = getattr( + batch, "ltx2_ti2v_clean_latent_background", None + ) + is_native_ti2v = ( + is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config) + and batch.image_path is not None + and isinstance(batch.latents, torch.Tensor) + ) + if is_native_ti2v: + # Official two-stage TI2V keeps the upsampled stage-2 latent as the + # clean background and only overwrites the conditioned frame tokens. + batch.ltx2_ti2v_clean_latent_background = batch.latents.detach().clone() + else: + batch.ltx2_ti2v_clean_latent_background = None if self._should_reset_stage2_generators(server_args): self._reset_stage2_generators(batch) - noise_scale = self.distilled_sigmas[0].to(batch.latents.device) - video_noise = self._randn_like_with_batch_generators(batch.latents, batch) - batch.latents = video_noise * noise_scale + batch.latents * (1 - noise_scale) + noise_scale = float(self.distilled_sigmas[0].item()) + if is_native_ti2v: + prepared_latents, denoise_mask, _ = self._prepare_ltx2_ti2v_clean_state( + latents=batch.latents, + image_latent=batch.image_latent, + num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)), + zero_clean_latent=True, + clean_latent_background=batch.ltx2_ti2v_clean_latent_background, + ) + video_noise = self._randn_like_with_batch_generators( + prepared_latents, batch + ) + scaled_mask = ( + denoise_mask.to(device=prepared_latents.device, dtype=torch.float32) + * noise_scale + ) + batch.latents = ( + video_noise * scaled_mask + prepared_latents * (1 - scaled_mask) + ).to(prepared_latents.dtype) + else: + video_noise = self._randn_like_with_batch_generators(batch.latents, batch) + batch.latents = ( + video_noise * noise_scale + batch.latents * (1 - noise_scale) + ).to(batch.latents.dtype) if isinstance(batch.audio_latents, torch.Tensor): audio_noise = self._randn_like_with_batch_generators( batch.audio_latents, batch ) - audio_noise_scale = noise_scale.to( - batch.audio_latents.device, batch.audio_latents.dtype + audio_scaled_mask = ( + torch.ones_like(batch.audio_latents[..., :1], dtype=torch.float32) + * noise_scale ) batch.audio_latents = ( - audio_noise * audio_noise_scale - + batch.audio_latents * (1 - audio_noise_scale) - ) + audio_noise * audio_scaled_mask + + batch.audio_latents * (1 - audio_scaled_mask) + ).to(batch.audio_latents.dtype) if not is_ltx23_native_variant( server_args.pipeline_config.vae_config.arch_config ): @@ -214,5 +248,6 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): batch.timesteps = original_batch_timesteps batch.num_inference_steps = original_batch_num_inference_steps batch.do_classifier_free_guidance = original_do_cfg + batch.ltx2_ti2v_clean_latent_background = original_clean_latent_background return batch 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 index a644fc693..2c0540115 100644 --- 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 @@ -59,6 +59,8 @@ class LTX2DenoisingContext(DenoisingContext): latent_width: int = 0 denoise_mask: torch.Tensor | None = None clean_latent: torch.Tensor | None = None + last_denoised_video: torch.Tensor | None = None + last_denoised_audio: torch.Tensor | None = None trajectory_audio_latents: list[torch.Tensor] = field(default_factory=list) @@ -167,6 +169,7 @@ class LTX2DenoisingStage(DenoisingStage): image_latent: torch.Tensor, num_img_tokens: int, zero_clean_latent: bool, + clean_latent_background: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: latents = latents.clone() conditioned = image_latent[:, :num_img_tokens, :].to( @@ -179,7 +182,13 @@ class LTX2DenoisingStage(DenoisingStage): dtype=torch.float32, ) denoise_mask[:, :num_img_tokens, :] = 0.0 - if zero_clean_latent: + if clean_latent_background is not None: + clean_latent = ( + clean_latent_background.detach() + .clone() + .to(device=latents.device, dtype=latents.dtype) + ) + elif zero_clean_latent: clean_latent = torch.zeros_like(latents) else: clean_latent = latents.detach().clone() @@ -642,6 +651,14 @@ class LTX2DenoisingStage(DenoisingStage): 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].") + clean_latent_background = getattr( + batch, "ltx2_ti2v_clean_latent_background", None + ) + if not ( + isinstance(clean_latent_background, torch.Tensor) + and clean_latent_background.shape == ctx.latents.shape + ): + clean_latent_background = None # 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( @@ -649,6 +666,7 @@ class LTX2DenoisingStage(DenoisingStage): image_latent=batch.image_latent, num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)), zero_clean_latent=ctx.is_ltx23_variant, + clean_latent_background=clean_latent_background, ) ) return ctx @@ -948,6 +966,11 @@ class LTX2DenoisingStage(DenoisingStage): ctx.audio_latents = ctx.audio_scheduler.step( model_audio, step.t_device, ctx.audio_latents, return_dict=False )[0] + if ctx.denoise_mask is not None and ctx.clean_latent is not None: + ctx.latents = ( + ctx.latents.float() * ctx.denoise_mask + + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask) + ).to(dtype=ctx.latents.dtype) ctx.latents = self.post_forward_for_ti2v_task( batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z ) @@ -956,88 +979,57 @@ class LTX2DenoisingStage(DenoisingStage): 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, - ) - ) + negative_encoder_hidden_states = batch.negative_prompt_embeds[0] + negative_audio_encoder_hidden_states = batch.negative_audio_prompt_embeds[0] + negative_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, + ) - if stage1_guider_params is not None or batch.do_classifier_free_guidance: - v_neg, a_v_neg = step.current_model( + 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 + ) + need_modality = ( + float(stage1_guider_params["video_modality_scale"]) != 1.0 + or float(stage1_guider_params["audio_modality_scale"]) != 1.0 + ) + + if ctx.use_ltx23_legacy_one_stage: + 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=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, - ), + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ) + ) + v_neg, a_v_neg = step.current_model( + **build_model_kwargs( + encoder_hidden_states=negative_encoder_hidden_states, + audio_encoder_hidden_states=negative_audio_encoder_hidden_states, + encoder_attention_mask=negative_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_pos = v_pos.float() + a_v_pos = a_v_pos.float() 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 - ) + v_ptb = None + a_v_ptb = None if need_perturbed: with set_forward_context( current_timestep=step.step_index, attn_metadata=step.attn_metadata @@ -1055,17 +1047,11 @@ class LTX2DenoisingStage(DenoisingStage): ), ) ) - 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 - ) + v_ptb = v_ptb.float() + a_v_ptb = a_v_ptb.float() - need_modality = ( - float(stage1_guider_params["video_modality_scale"]) != 1.0 - or float(stage1_guider_params["audio_modality_scale"]) != 1.0 - ) + v_mod = None + a_v_mod = None if need_modality: with set_forward_context( current_timestep=step.step_index, attn_metadata=step.attn_metadata @@ -1079,76 +1065,401 @@ class LTX2DenoisingStage(DenoisingStage): disable_v2a_cross_attn=True, ) ) - denoised_video_modality = self._ltx2_velocity_to_x0( - ctx.latents, v_mod.float(), video_sigma_for_x0 + v_mod = v_mod.float() + a_v_mod = a_v_mod.float() + else: + # NOTE: this flag must be identical across all SP ranks so that + # every rank executes the same number of model-forward calls (each + # of which contains NCCL collectives). + # _should_apply_ltx2_ti2v() is SP-rank-dependent (only the rank owning the first latent + # frame returns True), so we must NOT use it here. + # Instead we check the rank-invariant attribute that is always set on every + # rank when the request is a TI2V request. + use_split_two_stage_ti2v_guider = ( + server_args.pipeline_class_name == "LTX2TwoStagePipeline" + and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 + ) + + def cat_or_none(items: list[torch.Tensor | None]) -> torch.Tensor | None: + if items[0] is None: + return None + return torch.cat(items, dim=0) + + pass_specs: list[ + tuple[ + str, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + dict[str, object], + ] + ] = [ + ( + "cond", + encoder_hidden_states, + audio_encoder_hidden_states, + encoder_attention_mask, + { + "skip_video_self_attn_blocks": (), + "skip_audio_self_attn_blocks": (), + "skip_a2v_cross_attn": False, + "skip_v2a_cross_attn": False, + }, + ), + ( + "neg", + negative_encoder_hidden_states, + negative_audio_encoder_hidden_states, + negative_encoder_attention_mask, + { + "skip_video_self_attn_blocks": (), + "skip_audio_self_attn_blocks": (), + "skip_a2v_cross_attn": False, + "skip_v2a_cross_attn": False, + }, + ), + ] + if need_perturbed: + pass_specs.append( + ( + "perturbed", + encoder_hidden_states, + audio_encoder_hidden_states, + 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"] + ), + "skip_a2v_cross_attn": False, + "skip_v2a_cross_attn": False, + }, + ) ) - denoised_audio_modality = self._ltx2_velocity_to_x0( - ctx.audio_latents, a_v_mod.float(), sigma_val + if need_modality: + pass_specs.append( + ( + "modality", + encoder_hidden_states, + audio_encoder_hidden_states, + encoder_attention_mask, + { + "skip_video_self_attn_blocks": (), + "skip_audio_self_attn_blocks": (), + "skip_a2v_cross_attn": True, + "skip_v2a_cross_attn": True, + }, + ) ) - 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 + num_passes = len(pass_specs) + expanded_batch_size = batch_size * num_passes + perturbation_configs = tuple( + perturbation_config + for _, _, _, _, perturbation_config in pass_specs + for _ in range(batch_size) ) - denoised_audio = denoised_audio + (batch.guidance_scale - 1.0) * ( - denoised_audio - denoised_audio_neg + batched_hidden_states = self._repeat_batch_dim( + latent_model_input, expanded_batch_size ) + batched_audio_hidden_states = self._repeat_batch_dim( + audio_latent_model_input, expanded_batch_size + ) + batched_encoder_hidden_states = torch.cat( + [item[1] for item in pass_specs], dim=0 + ) + batched_audio_encoder_hidden_states = torch.cat( + [item[2] for item in pass_specs], dim=0 + ) + batched_timestep_video = self._repeat_batch_dim( + timestep_video, expanded_batch_size + ) + batched_timestep_audio = self._repeat_batch_dim( + timestep_audio, expanded_batch_size + ) + batched_prompt_timestep_video = ( + None + if prompt_timestep_video is None + else self._repeat_batch_dim(prompt_timestep_video, expanded_batch_size) + ) + batched_prompt_timestep_audio = ( + None + if prompt_timestep_audio is None + else self._repeat_batch_dim(prompt_timestep_audio, expanded_batch_size) + ) + batched_encoder_attention_mask = cat_or_none( + [item[3] for item in pass_specs] + ) + batched_audio_encoder_attention_mask = cat_or_none( + [item[3] for item in pass_specs] + ) + batched_video_coords = ( + None + if video_coords is None + else self._repeat_batch_dim(video_coords, expanded_batch_size) + ) + batched_audio_coords = ( + None + if audio_coords is None + else self._repeat_batch_dim(audio_coords, expanded_batch_size) + ) + batched_video_self_attention_mask = ( + None + if video_self_attention_mask is None + else self._repeat_batch_dim( + video_self_attention_mask, expanded_batch_size + ) + ) + batched_audio_self_attention_mask = ( + None + if audio_self_attention_mask is None + else self._repeat_batch_dim( + audio_self_attention_mask, expanded_batch_size + ) + ) + batched_a2v_cross_attention_mask = ( + None + if a2v_cross_attention_mask is None + else self._repeat_batch_dim( + a2v_cross_attention_mask, expanded_batch_size + ) + ) + batched_v2a_cross_attention_mask = ( + None + if v2a_cross_attention_mask is None + else self._repeat_batch_dim( + v2a_cross_attention_mask, expanded_batch_size + ) + ) + if use_split_two_stage_ti2v_guider: + split_sizes = [1] * expanded_batch_size + + def split_or_none( + tensor: torch.Tensor | None, + ) -> list[torch.Tensor | None]: + if tensor is None: + return [None] * len(split_sizes) + return list(tensor.split(split_sizes, dim=0)) + + batched_video_chunks = [] + batched_audio_chunks = [] + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + for ( + hidden_states_chunk, + audio_hidden_states_chunk, + encoder_hidden_states_chunk, + audio_encoder_hidden_states_chunk, + timestep_video_chunk, + timestep_audio_chunk, + prompt_timestep_video_chunk, + prompt_timestep_audio_chunk, + encoder_attention_mask_chunk, + audio_encoder_attention_mask_chunk, + video_coords_chunk, + audio_coords_chunk, + video_self_attention_mask_chunk, + audio_self_attention_mask_chunk, + a2v_cross_attention_mask_chunk, + v2a_cross_attention_mask_chunk, + perturbation_config_chunk, + ) in zip( + batched_hidden_states.split(split_sizes, dim=0), + batched_audio_hidden_states.split(split_sizes, dim=0), + batched_encoder_hidden_states.split(split_sizes, dim=0), + batched_audio_encoder_hidden_states.split(split_sizes, dim=0), + batched_timestep_video.split(split_sizes, dim=0), + batched_timestep_audio.split(split_sizes, dim=0), + split_or_none(batched_prompt_timestep_video), + split_or_none(batched_prompt_timestep_audio), + split_or_none(batched_encoder_attention_mask), + split_or_none(batched_audio_encoder_attention_mask), + split_or_none(batched_video_coords), + split_or_none(batched_audio_coords), + split_or_none(batched_video_self_attention_mask), + split_or_none(batched_audio_self_attention_mask), + split_or_none(batched_a2v_cross_attention_mask), + split_or_none(batched_v2a_cross_attention_mask), + ((cfg,) for cfg in perturbation_configs), + strict=True, + ): + video_chunk, audio_chunk = step.current_model( + hidden_states=hidden_states_chunk, + audio_hidden_states=audio_hidden_states_chunk, + encoder_hidden_states=encoder_hidden_states_chunk, + audio_encoder_hidden_states=audio_encoder_hidden_states_chunk, + timestep=timestep_video_chunk, + audio_timestep=timestep_audio_chunk, + prompt_timestep=prompt_timestep_video_chunk, + audio_prompt_timestep=prompt_timestep_audio_chunk, + encoder_attention_mask=encoder_attention_mask_chunk, + audio_encoder_attention_mask=audio_encoder_attention_mask_chunk, + 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_chunk, + audio_coords=audio_coords_chunk, + video_self_attention_mask=video_self_attention_mask_chunk, + audio_self_attention_mask=audio_self_attention_mask_chunk, + a2v_cross_attention_mask=a2v_cross_attention_mask_chunk, + v2a_cross_attention_mask=v2a_cross_attention_mask_chunk, + audio_replicated_for_sp=ctx.replicate_audio_for_sp, + perturbation_configs=perturbation_config_chunk, + return_latents=False, + return_dict=False, + ) + batched_video_chunks.append(video_chunk) + batched_audio_chunks.append(audio_chunk) + + batched_video = torch.cat(batched_video_chunks, dim=0) + batched_audio = torch.cat(batched_audio_chunks, dim=0) + else: + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + batched_video, batched_audio = step.current_model( + hidden_states=batched_hidden_states, + audio_hidden_states=batched_audio_hidden_states, + encoder_hidden_states=batched_encoder_hidden_states, + audio_encoder_hidden_states=batched_audio_encoder_hidden_states, + timestep=batched_timestep_video, + audio_timestep=batched_timestep_audio, + prompt_timestep=batched_prompt_timestep_video, + audio_prompt_timestep=batched_prompt_timestep_audio, + encoder_attention_mask=batched_encoder_attention_mask, + audio_encoder_attention_mask=batched_audio_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=batched_video_coords, + audio_coords=batched_audio_coords, + video_self_attention_mask=batched_video_self_attention_mask, + audio_self_attention_mask=batched_audio_self_attention_mask, + a2v_cross_attention_mask=batched_a2v_cross_attention_mask, + v2a_cross_attention_mask=batched_v2a_cross_attention_mask, + audio_replicated_for_sp=ctx.replicate_audio_for_sp, + perturbation_configs=perturbation_configs, + return_latents=False, + return_dict=False, + ) + + batched_video = batched_video.float() + batched_audio = batched_audio.float() + pass_outputs = { + pass_name: ( + video_chunk, + audio_chunk, + ) + for (pass_name, _, _, _, _), video_chunk, audio_chunk in zip( + pass_specs, + batched_video.chunk(num_passes, dim=0), + batched_audio.chunk(num_passes, dim=0), + strict=True, + ) + } + v_pos, a_v_pos = pass_outputs["cond"] + v_neg, a_v_neg = pass_outputs["neg"] + v_ptb, a_v_ptb = pass_outputs.get("perturbed", (None, None)) + v_mod, a_v_mod = pass_outputs.get("modality", (None, None)) + + 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 = 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 + ) + denoised_video_perturbed = ( + None + if v_ptb is None + else self._ltx2_velocity_to_x0(ctx.latents, v_ptb, video_sigma_for_x0) + ) + denoised_audio_perturbed = ( + None + if a_v_ptb is None + else self._ltx2_velocity_to_x0(ctx.audio_latents, a_v_ptb, sigma_val) + ) + denoised_video_modality = ( + None + if v_mod is None + else self._ltx2_velocity_to_x0(ctx.latents, v_mod, video_sigma_for_x0) + ) + denoised_audio_modality = ( + None + if a_v_mod is None + else self._ltx2_velocity_to_x0(ctx.audio_latents, a_v_mod, sigma_val) + ) + + if not video_skip: + denoised_video = self._ltx2_calculate_guided_x0( + cond=denoised_video, + uncond_text=denoised_video_neg, + 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"]), + ) + ctx.last_denoised_video = denoised_video + elif ctx.last_denoised_video is not None: + denoised_video = ctx.last_denoised_video + + if not audio_skip: + denoised_audio = self._ltx2_calculate_guided_x0( + cond=denoised_audio, + uncond_text=denoised_audio_neg, + 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"]), + ) + ctx.last_denoised_audio = denoised_audio + elif ctx.last_denoised_audio is not None: + denoised_audio = ctx.last_denoised_audio 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) - ) + ).to(denoised_video.dtype) # 6. Convert x0 predictions back to velocity and update both latent streams. if sigma_val == 0.0: diff --git a/python/sglang/multimodal_gen/test/server/accuracy_config.py b/python/sglang/multimodal_gen/test/server/accuracy_config.py index 1f052264f..469f58644 100644 --- a/python/sglang/multimodal_gen/test/server/accuracy_config.py +++ b/python/sglang/multimodal_gen/test/server/accuracy_config.py @@ -68,11 +68,6 @@ SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = { "HF AutoencoderDC checkpoint leaves required to_qkv_multiscale weights missing, so VAE transfer would compare against partially initialized reference weights" ) }, - "mova_360p_1gpu": { - ComponentType.TRANSFORMER: ComponentSkip( - "HF reference transformer cannot be materialized from the video_dit repo layout" - ) - }, "qwen_image_t2i_cache_dit_enabled": { ComponentType.VAE: ComponentSkip( "Representative VAE accuracy is already covered by qwen_image_t2i for the same source component and topology" @@ -324,14 +319,6 @@ SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = { "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer" ), }, - "mova_360p_ring2_uly1": { - ComponentType.TRANSFORMER: ComponentSkip( - "HF reference transformer cannot be materialized from the MOVA video_dit repo layout" - ), - ComponentType.TEXT_ENCODER: ComponentSkip( - "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer" - ), - }, "flux_image_t2i_2_gpus": { ComponentType.TEXT_ENCODER: ComponentSkip( "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.47) after 100% matched weight transfer" diff --git a/python/sglang/multimodal_gen/test/server/consistency_threshold.json b/python/sglang/multimodal_gen/test/server/consistency_threshold.json index 45e04b413..bdba43a43 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_threshold.json +++ b/python/sglang/multimodal_gen/test/server/consistency_threshold.json @@ -79,12 +79,6 @@ "psnr_threshold": 28.0, "mean_abs_diff_threshold": 8.0 }, - "mova_360p_1gpu": { - "clip_threshold": 0.90, - "ssim_threshold": 0.87, - "psnr_threshold": 24.0, - "mean_abs_diff_threshold": 10.0 - }, "wan2_1_t2v_1_3b_lora_1gpu": { "clip_threshold": 0.88, "ssim_threshold": 0.75, @@ -223,12 +217,6 @@ "psnr_threshold": 22.2, "mean_abs_diff_threshold": 10.0 }, - "mova_360p_ring2_uly1": { - "clip_threshold": 0.90, - "ssim_threshold": 0.91, - "psnr_threshold": 24.0, - "mean_abs_diff_threshold": 10.0 - }, "wan2_1_i2v_14b_480P_2gpu": { "clip_threshold": 0.76, "ssim_threshold": 0.51, diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 04c93c430..3011ddbd7 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -2616,7 +2616,63 @@ }, "expected_e2e_ms": 34384.39, "expected_avg_denoise_ms": 782.76, - "expected_median_denoise_ms": 806.95 + "expected_median_denoise_ms": 806.95, + "estimated_full_test_time_s": 216.7 + }, + "ltx_2_3_two_stage_ti2v_2gpus": { + "stages_ms": { + "InputValidationStage": 3.61, + "TextEncodingStage": 1775.95, + "LTX2TextConnectorStage": 27.44, + "LTX2HalveResolutionStage": 0.06, + "LTX2LoRASwitchStage": 107.23, + "LTX2SigmaPreparationStage": 0.41, + "TimestepPreparationStage": 24.7, + "LTX2AVLatentPreparationStage": 0.2, + "LTX2AVDenoisingStage": 26851.54, + "LTX2UpsampleStage": 408.81, + "LTX2RefinementStage": 1561.91, + "LTX2AVDecodingStage": 862.29, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 812.4, + "1": 817.22, + "2": 1159.79, + "3": 852.65, + "4": 857.69, + "5": 864.02, + "6": 868.72, + "7": 859.11, + "8": 836.39, + "9": 841.12, + "10": 1093.34, + "11": 902.56, + "12": 830.56, + "13": 818.54, + "14": 837.61, + "15": 837.03, + "16": 853.87, + "17": 830.49, + "18": 833.02, + "19": 1179.55, + "20": 846.55, + "21": 856.85, + "22": 887.05, + "23": 836.35, + "24": 843.8, + "25": 836.38, + "26": 821.39, + "27": 826.6, + "28": 1165.21, + "29": 833.04, + "30": 540.96, + "31": 334.98, + "32": 312.19 + }, + "expected_e2e_ms": 40552.46, + "expected_avg_denoise_ms": 840.21, + "expected_median_denoise_ms": 837.61 } } } diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index cc7bf4da6..0ca40edca 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -549,17 +549,6 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [ extras={"enable_upscaling": True, "upscaling_scale": 4}, ), ), - DiffusionTestCase( - "mova_360p_1gpu", - DiffusionServerArgs( - model_path=DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST, - modality="video", - num_gpus=1, - dit_layerwise_offload=True, - ), - TI2V_sampling_params, - run_perf_check=False, - ), ] HUNYUAN3D_SHAPE_sampling_params = DiffusionSamplingParams( @@ -917,19 +906,6 @@ TWO_GPU_CASES_A = [ TI2V_sampling_params, run_perf_check=False, ), - DiffusionTestCase( - "mova_360p_ring2_uly1", - DiffusionServerArgs( - model_path=DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST, - modality="video", - num_gpus=2, - ring_degree=2, - ulysses_degree=1, - dit_layerwise_offload=True, - ), - TI2V_sampling_params, - run_perf_check=False, - ), DiffusionTestCase( "ltx_2_two_stage_t2v", DiffusionServerArgs( @@ -940,6 +916,16 @@ TWO_GPU_CASES_A = [ ), T2V_sampling_params, ), + DiffusionTestCase( + "ltx_2_3_two_stage_ti2v_2gpus", + DiffusionServerArgs( + model_path="Lightricks/LTX-2.3", + modality="video", + num_gpus=2, + extras=["--pipeline-class-name LTX2TwoStagePipeline"], + ), + TI2V_sampling_params, + ), ] TWO_GPU_CASES_B = [ @@ -1098,7 +1084,6 @@ ACCURACY_ONE_GPU_CASES_A_IDS = ( "qwen_image_edit_2511_ti2i", "qwen_image_layered_i2i", "flux_2_image_t2i_upscaling_4x", - "mova_360p_1gpu", ) ACCURACY_ONE_GPU_CASES_B_IDS = ( @@ -1130,7 +1115,6 @@ ACCURACY_TWO_GPU_CASES_A_IDS = ( "fsdp-inference", "mova_360p_tp2", "mova_360p_ring1_uly2", - "mova_360p_ring2_uly1", "ltx_2_two_stage_t2v", ) diff --git a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py b/python/sglang/multimodal_gen/test/unit/test_sampling_params.py index 8d53a53cb..4ccc7fcb8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py +++ b/python/sglang/multimodal_gen/test/unit/test_sampling_params.py @@ -1,8 +1,14 @@ import argparse import math import unittest +from types import SimpleNamespace from unittest.mock import MagicMock, patch +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + LTX2PipelineConfig, + is_ltx23_native_variant, + sync_ltx23_runtime_vae_markers, +) from sglang.multimodal_gen.configs.sample.diffusers_generic import ( DiffusersGenericSamplingParams, ) @@ -141,6 +147,33 @@ class TestSamplingParamsSubclass(unittest.TestCase): expected, ) + def test_ltx23_runtime_vae_markers_sync_variant_and_decoder_metadata(self): + arch_config = LTX2PipelineConfig().vae_config.arch_config + + self.assertFalse(is_ltx23_native_variant(arch_config)) + self.assertEqual(arch_config.video_decoder_variant, "ltx_2") + self.assertEqual(arch_config.condition_encoder_subdir, "") + + sync_ltx23_runtime_vae_markers( + arch_config, + SimpleNamespace( + arch_config=SimpleNamespace( + ltx_variant="ltx_2_3", + condition_encoder_subdir="ltx23_image_encoder", + video_decoder_variant="ltx_2_3", + video_decoder_config={"_class_name": "AutoencoderKLLTX2Video"}, + ) + ), + ) + + self.assertTrue(is_ltx23_native_variant(arch_config)) + self.assertEqual(arch_config.condition_encoder_subdir, "ltx23_image_encoder") + self.assertEqual(arch_config.video_decoder_variant, "ltx_2_3") + self.assertEqual( + arch_config.video_decoder_config, + {"_class_name": "AutoencoderKLLTX2Video"}, + ) + class TestSamplingParamsCliArgs(unittest.TestCase): def _parse_cli_kwargs(self, argv: list[str]) -> dict: diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py new file mode 100644 index 000000000..fd9b52f56 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py @@ -0,0 +1,48 @@ +import unittest + +import torch + +from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import ( + _backfill_ltx2_audio_vae_latent_stats, +) + + +class TestVAELoader(unittest.TestCase): + def test_backfill_ltx2_audio_vae_latent_stats_maps_official_keys(self): + loaded = { + "per_channel_statistics.mean-of-means": torch.tensor([1.0, 2.0]), + "per_channel_statistics.std-of-means": torch.tensor([3.0, 4.0]), + } + + _backfill_ltx2_audio_vae_latent_stats(loaded, "audio_vae") + + self.assertTrue(torch.equal(loaded["latents_mean"], torch.tensor([1.0, 2.0]))) + self.assertTrue(torch.equal(loaded["latents_std"], torch.tensor([3.0, 4.0]))) + + def test_backfill_ltx2_audio_vae_latent_stats_does_not_override_existing(self): + loaded = { + "per_channel_statistics.mean-of-means": torch.tensor([1.0, 2.0]), + "per_channel_statistics.std-of-means": torch.tensor([3.0, 4.0]), + "latents_mean": torch.tensor([5.0, 6.0]), + "latents_std": torch.tensor([7.0, 8.0]), + } + + _backfill_ltx2_audio_vae_latent_stats(loaded, "audio_vae") + + self.assertTrue(torch.equal(loaded["latents_mean"], torch.tensor([5.0, 6.0]))) + self.assertTrue(torch.equal(loaded["latents_std"], torch.tensor([7.0, 8.0]))) + + def test_backfill_ltx2_audio_vae_latent_stats_skips_non_audio_vae(self): + loaded = { + "per_channel_statistics.mean-of-means": torch.tensor([1.0]), + "per_channel_statistics.std-of-means": torch.tensor([2.0]), + } + + _backfill_ltx2_audio_vae_latent_stats(loaded, "vae") + + self.assertNotIn("latents_mean", loaded) + self.assertNotIn("latents_std", loaded) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/ci/utils/diffusion/comparison_configs.json b/scripts/ci/utils/diffusion/comparison_configs.json index db51d53fa..d9a8938ef 100644 --- a/scripts/ci/utils/diffusion/comparison_configs.json +++ b/scripts/ci/utils/diffusion/comparison_configs.json @@ -136,18 +136,16 @@ } }, { - "id": "ltx2.3_twostage_t2v_2gpus", + "id": "ltx2.3_twostage_ti2v_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, + "task": "image-to-video", + "prompt": "The cat starts walking slowly towards the camera.", "num_frames": 121, "seed": 42, "num_gpus": 2, "frameworks": { "sglang": { - "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --pipeline-class-name LTX2TwoStagePipeline", + "serve_args": "--enable-torch-compile --warmup --pipeline-class-name LTX2TwoStagePipeline", "extra_env": {} } }