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

[diffusion] model: support Ltx 2.3 two stage ti2v (#22667)

Browse Source
This commit is contained in:
Mick
2026-04-14 22:10:08 +08:00
committed by GitHub
parent 680bd4b429
commit c5e95080d2
15 changed files with 837 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py
+18
View File
@@ -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,
python/sglang/multimodal_gen/configs/sample/sampling_params.py
+2 -1
View File
@@ -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",
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+14
View File
@@ -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())
python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py
+2 -1
View File
@@ -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,
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+142 -12
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
+5
View File
@@ -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):
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py
+46 -11
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py
+456 -145
View File
@@ -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:
python/sglang/multimodal_gen/test/server/accuracy_config.py
-13
View File
@@ -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"
python/sglang/multimodal_gen/test/server/consistency_threshold.json
-12
View File
@@ -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,
python/sglang/multimodal_gen/test/server/perf_baselines.json
+57 -1
View File
@@ -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
}
}
}
python/sglang/multimodal_gen/test/server/testcase_configs.py
+10 -26
View File
@@ -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",
)
python/sglang/multimodal_gen/test/unit/test_sampling_params.py
+33
View File
@@ -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:
python/sglang/multimodal_gen/test/unit/test_vae_loader.py
+48
View File
@@ -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()