diff --git a/.github/workflows/diffusion-ci-gt-gen.yml b/.github/workflows/diffusion-ci-gt-gen.yml index 5bbc8aa7b..909cfff54 100644 --- a/.github/workflows/diffusion-ci-gt-gen.yml +++ b/.github/workflows/diffusion-ci-gt-gen.yml @@ -39,10 +39,15 @@ on: default: '' type: string official_source_group: - description: 'Official GT source group filter: all, diffusers, wan21, or ltx23. Used only when run_official_cases is true.' + description: 'Official GT source group filter: all, diffusers, wan21, ltx, or ltx23. Used only when run_official_cases is true.' required: false default: '' type: string + ci_data_ref: + description: 'ci-data ref to use for repro scripts when running official GT cases.' + required: false + default: 'main' + type: string concurrency: group: diffusion-ci-gt-gen-${{ github.ref }}-${{ inputs.output_name || inputs.case_ids || inputs.official_case_ids || inputs.official_source_group || inputs.run_official_cases || 'default' }} @@ -90,6 +95,13 @@ jobs: "zimage_image_t2i", ], "wan21": ["wan2_1_t2v_1.3b"], + "ltx": [ + "ltx_2_two_stage_t2v", + "ltx_2.3_two_stage_t2v_2gpus", + "ltx_2_3_two_stage_ti2v_2gpus", + "ltx_2.3_one_stage_ti2v", + "ltx_2_3_hq_pipeline", + ], "ltx23": [ "ltx_2.3_two_stage_t2v_2gpus", "ltx_2.3_one_stage_ti2v", @@ -97,7 +109,7 @@ jobs: } source_group = os.environ["OFFICIAL_SOURCE_GROUP"].strip() or "all" if source_group == "all": - selected_groups = list(groups) + selected_groups = ["diffusers", "wan21", "ltx"] elif source_group in groups: selected_groups = [source_group] else: @@ -109,7 +121,6 @@ jobs: "flux_2_image_t2i", "flux_2_klein_image_t2i", "flux_2_ti2i", - "ltx_2.3_two_stage_t2v_2gpus", } include = [] for group in selected_groups: @@ -154,9 +165,13 @@ jobs: uses: actions/checkout@v4 with: repository: sgl-project/ci-data - ref: main + ref: ${{ inputs.ci_data_ref || 'main' }} path: ci-data token: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} + sparse-checkout: | + diffusion-ci/repro_scripts + diffusion-ci/consistency_gt/official_generated/case_map.json + sparse-checkout-cone-mode: false - name: Prepare sgl-kernel/dist for prebuilt wheel if: inputs.kernel_artifact_run_id != '' @@ -179,8 +194,8 @@ jobs: CUSTOM_BUILD_SGL_KERNEL="${{ inputs.kernel_artifact_run_id != '' && 'true' || 'false' }}" \ bash scripts/ci/cuda/ci_install_dependency.sh diffusion - - name: Install official LTX-2.3 repro dependencies - if: matrix.source_group == 'ltx23' + - name: Install official LTX repro dependencies + if: matrix.source_group == 'ltx' || matrix.source_group == 'ltx23' run: | UV_SYSTEM_PYTHON=1 uv pip install \ "transformers==4.52.4" \ @@ -201,7 +216,7 @@ jobs: git -C /tmp/mmgen-official-code/Wan2.1 rev-parse HEAD - name: Checkout official LTX-2 repo - if: matrix.source_group == 'ltx23' + if: matrix.source_group == 'ltx' || matrix.source_group == 'ltx23' run: | mkdir -p /tmp/mmgen-official-code if [ ! -d /tmp/mmgen-official-code/LTX-2/.git ]; then @@ -211,65 +226,6 @@ jobs: git -C /tmp/mmgen-official-code/LTX-2 checkout 41d924371612b692c0fd1e4d9d94c3dfb3c02cb3 git -C /tmp/mmgen-official-code/LTX-2 rev-parse HEAD - - name: Prepare official LTX-2.3 assets - if: matrix.source_group == 'ltx23' - env: - HF_TOKEN: ${{ secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }} - HUGGING_FACE_HUB_TOKEN: ${{ secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }} - run: | - python - <<'PY' - import os - from pathlib import Path - - from huggingface_hub import hf_hub_download, snapshot_download - - root = Path.home() / ".cache/sgl_diffusion/materialized_models/Lightricks__LTX-2.3-official-assets" - root.mkdir(parents=True, exist_ok=True) - text_encoder = root / "text_encoder" - tokenizer = root / "tokenizer" - text_encoder.mkdir(exist_ok=True) - tokenizer.mkdir(exist_ok=True) - - gemma_root = Path( - snapshot_download( - repo_id="google/gemma-3-12b-it-qat-q4_0-unquantized", - ignore_patterns=["*.onnx", "*.msgpack"], - max_workers=8, - ) - ) - tokenizer_names = { - "tokenizer.json", - "tokenizer.model", - "tokenizer_config.json", - "special_tokens_map.json", - } - - def link(src: Path, dst: Path) -> None: - if dst.exists() or dst.is_symlink(): - return - try: - dst.symlink_to(src) - except OSError: - os.link(src, dst) - - for src in gemma_root.iterdir(): - if not src.is_file(): - continue - link(src, text_encoder / src.name) - if src.name in tokenizer_names: - link(src, tokenizer / src.name) - - for filename in ( - "ltx-2.3-22b-distilled-lora-384.safetensors", - "ltx-2.3-spatial-upscaler-x2-1.1.safetensors", - ): - src = Path(hf_hub_download("Lightricks/LTX-2.3", filename)) - link(src, root / filename) - - print(root) - print(sorted(p.name for p in root.iterdir())) - PY - - name: Generate official output env: HF_TOKEN: ${{ secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }} @@ -301,17 +257,22 @@ jobs: --dtype bf16 \ --device-map none \ --generator-device cuda - elif [ "$SOURCE_GROUP" = "ltx23" ]; then + elif [ "$SOURCE_GROUP" = "ltx" ] || [ "$SOURCE_GROUP" = "ltx23" ]; then cd python + extra_ltx_args=() + if [ "$CASE_ID" = "ltx_2_3_hq_pipeline" ]; then + extra_ltx_args+=(--num-frames 24) + fi set +e PYTHONPATH=/tmp/mmgen-official-code/LTX-2/packages/ltx-core/src:/tmp/mmgen-official-code/LTX-2/packages/ltx-pipelines/src:$PWD:$PYTHONPATH \ python ../ci-data/diffusion-ci/repro_scripts/gen_official_ltx23.py \ --out-dir ./${{ env.OUTPUT_NAME }} \ - --case-ids "$CASE_ID" + --case-ids "$CASE_ID" \ + "${extra_ltx_args[@]}" status=$? set -e if [ "$status" -ne 0 ]; then - find ./${{ env.OUTPUT_NAME }} -name 'official_ltx23_manifest.json' -print -exec cat {} \; + find ./${{ env.OUTPUT_NAME }} \( -name 'official_ltx_manifest.json' -o -name 'official_ltx23_manifest.json' \) -print -exec cat {} \; exit "$status" fi else @@ -327,6 +288,7 @@ jobs: python/${{ env.OUTPUT_NAME }}/*.jpg python/${{ env.OUTPUT_NAME }}/*.png python/${{ env.OUTPUT_NAME }}/official_gt_manifest_*.json + python/${{ env.OUTPUT_NAME }}/official_ltx_manifest.json python/${{ env.OUTPUT_NAME }}/official_ltx23_manifest.json retention-days: 7 diff --git a/python/sglang/multimodal_gen/configs/sample/ltx_2.py b/python/sglang/multimodal_gen/configs/sample/ltx_2.py index e074ace6e..81233cb27 100644 --- a/python/sglang/multimodal_gen/configs/sample/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/sample/ltx_2.py @@ -47,6 +47,7 @@ class LTX2SamplingParams(SamplingParams): class LTX23SamplingParams(LTX2SamplingParams): """Sampling parameters matching official LTX-2.3 one-stage defaults.""" + seed: int = 42 generator_device: str = "cuda" guidance_scale: float = 3.0 num_inference_steps: int = 30 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 cde2357c4..f7013fb3a 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 @@ -1,6 +1,8 @@ +import functools import math from typing import Optional, Tuple, Union +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F @@ -19,8 +21,7 @@ def apply_interleaved_rotary_emb( cos, sin = freqs x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) # [B, S, C // 2] x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) - out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) - return out + return x * cos + x_rotated * sin def apply_split_rotary_emb( @@ -47,7 +48,7 @@ def apply_split_rotary_emb( r = last // 2 # (..., 2, r) - split_x = x.reshape(*x.shape[:-1], 2, r).float() + split_x = x.reshape(*x.shape[:-1], 2, r) first_x = split_x[..., :1, :] # (..., 1, r) second_x = split_x[..., 1:, :] # (..., 1, r) @@ -70,6 +71,19 @@ def apply_split_rotary_emb( return out +@functools.lru_cache(maxsize=5) +def _ltx2_connector_rope_freq_grid_np( + theta: float, num_pos_dims: int, dim: int +) -> torch.Tensor: + # Official LTX uses NumPy float64 for double-precision RoPE frequencies. + n_elem = 2 * num_pos_dims + pow_indices = np.power( + theta, + np.linspace(0.0, 1.0, dim // n_elem, dtype=np.float64), + ) + return torch.tensor(pow_indices * math.pi / 2.0, dtype=torch.float32) + + class LTX2Attention(torch.nn.Module): r""" Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key @@ -261,18 +275,22 @@ class LTX2RotaryPosEmbed1d(nn.Module): # 2. Calculate 1D RoPE frequencies num_rope_elems = 2 # 1 (because 1D) * 2 (for cos, sin) = 2 - freqs_dtype = torch.float64 if self.double_precision else torch.float32 - pow_indices = torch.pow( - self.theta, - torch.linspace( - start=0.0, - end=1.0, - steps=self.dim // num_rope_elems, - dtype=freqs_dtype, - device=device, - ), - ) - freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) + if self.double_precision: + freqs = _ltx2_connector_rope_freq_grid_np(self.theta, 1, self.dim).to( + device=device + ) + else: + pow_indices = torch.pow( + self.theta, + torch.linspace( + start=0.0, + end=1.0, + steps=self.dim // num_rope_elems, + dtype=torch.float32, + device=device, + ), + ) + freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) # 3. Matrix-vector outer product between pos ids of shape (batch_size, seq_len) and freqs vector of shape # (self.dim // 2,). @@ -487,7 +505,12 @@ class LTX2ConnectorTransformer1d(nn.Module): attention_mask = torch.zeros_like(attention_mask) # 2. Calculate 1D RoPE positional embeddings - rotary_emb = self.rope(batch_size, seq_len, device=hidden_states.device) + rotary_emb = self.rope( + batch_size, + seq_len, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) # 3. Run 1D transformer blocks for block in self.transformer_blocks: 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 7a4c649ae..80c955611 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -4,8 +4,11 @@ from __future__ import annotations +import functools +import math from typing import Any, Optional, Tuple, Union +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F @@ -89,6 +92,17 @@ def _ltx2_batched_perturbation_mask( return mask.view(mask.numel(), *([1] * (values.ndim - 1))), False +@functools.lru_cache(maxsize=5) +def _ltx2_rope_freq_grid_np(theta: float, num_pos_dims: int, dim: int) -> torch.Tensor: + # Official LTX uses NumPy float64 for double-precision RoPE frequencies. + n_elem = 2 * num_pos_dims + pow_indices = np.power( + theta, + np.linspace(0.0, 1.0, dim // n_elem, dtype=np.float64), + ) + return torch.tensor(pow_indices * math.pi / 2.0, dtype=torch.float32) + + def apply_interleaved_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: @@ -331,18 +345,22 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): ).to(device) num_rope_elems = num_pos_dims * 2 - freqs_dtype = torch.float64 if self.double_precision else torch.float32 - pow_indices = torch.pow( - self.theta, - torch.linspace( - start=0.0, - end=1.0, - steps=self.dim // num_rope_elems, - dtype=freqs_dtype, - device=device, - ), - ) - freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) + if self.double_precision: + freqs = _ltx2_rope_freq_grid_np(self.theta, num_pos_dims, self.dim).to( + device=device + ) + else: + pow_indices = torch.pow( + self.theta, + torch.linspace( + start=0.0, + end=1.0, + steps=self.dim // num_rope_elems, + dtype=torch.float32, + device=device, + ), + ) + freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs freqs = freqs.transpose(-1, -2).flatten(2) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py index 7047af233..c5e6ae354 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py @@ -213,18 +213,11 @@ class Gemma3Attention(nn.Module): is_neox_style=True, ) - # NOTE(gmixiaojin): The shared RotaryEmbedding above computes inv_freq on - # GPU and uses the x1*cos - x2*sin formula, which causes slight - # numerical differences vs HuggingFace (see the NOTE in - # rotary_embedding.py:_compute_inv_freq). For HF-exact alignment we - # precompute inv_freq on CPU and use rotate_half in self.rotary_emb(). - freq_indices = ( - torch.arange(0, self.head_dim, 2, dtype=torch.int64).float() / self.head_dim + self.rope_scaling_factor = ( + float(rope_scaling["factor"]) + if rope_scaling and rope_scaling.get("factor") + else None ) - inv_freq = 1.0 / (self.rope_theta**freq_indices) - if rope_scaling and rope_scaling.get("factor"): - inv_freq = inv_freq / float(rope_scaling["factor"]) - self.register_buffer("_hf_inv_freq", inv_freq, persistent=False) # Local Attention not support attention mask, we use global attention instead. # self.attn = LocalAttention( @@ -246,12 +239,21 @@ class Gemma3Attention(nn.Module): ) def rotary_emb(self, positions, q, k): - """Apply RoPE using HF-exact formula with precomputed inv_freq.""" + """Apply RoPE using the same device-side inv_freq materialization as LTX.""" positions_flat = positions.flatten().float() num_tokens = positions_flat.shape[0] with torch.autocast(device_type=q.device.type, enabled=False): - freqs = torch.outer(positions_flat, self._hf_inv_freq.float()) + freq_indices = ( + torch.arange( + 0, self.head_dim, 2, dtype=torch.int64, device=q.device + ).float() + / self.head_dim + ) + inv_freq = 1.0 / (self.rope_theta**freq_indices) + if self.rope_scaling_factor is not None: + inv_freq = inv_freq / self.rope_scaling_factor + freqs = torch.outer(positions_flat, inv_freq) emb = freqs.repeat(1, 2) cos = emb.cos().to(q.dtype).unsqueeze(1) sin = emb.sin().to(q.dtype).unsqueeze(1) @@ -290,11 +292,10 @@ class Gemma3Attention(nn.Module): key = k.transpose(1, 2) value = v.transpose(1, 2) - min_val = torch.finfo(query.dtype).min - attn_mask = torch.zeros( + attn_mask = torch.ones( (seq_len, seq_len), device=hidden_states.device, - dtype=query.dtype, + dtype=torch.bool, ) causal = torch.triu( torch.ones( @@ -302,30 +303,39 @@ class Gemma3Attention(nn.Module): ), diagonal=1, ) - attn_mask = attn_mask.masked_fill(causal, min_val) + attn_mask = attn_mask.masked_fill(causal, False) if self.is_sliding and self.sliding_window is not None: idx = torch.arange(seq_len, device=hidden_states.device) dist = idx[None, :] - idx[:, None] too_far = dist > self.sliding_window - attn_mask = attn_mask.masked_fill(too_far, min_val) + attn_mask = attn_mask.masked_fill(too_far, False) - key_pad = ~attention_mask.to(torch.bool) attn_mask = attn_mask[None, None, :, :].expand(batch_size, 1, seq_len, seq_len) - attn_mask = attn_mask.masked_fill( - key_pad[:, None, None, :].expand(batch_size, 1, seq_len, seq_len), - min_val, - ) + attn_mask = attn_mask & attention_mask.to(torch.bool)[:, None, None, :] - attn_kwargs = { - "attn_mask": attn_mask, - "dropout_p": 0.0, - "is_causal": False, - "scale": self.scaling, - } if query.shape[1] != key.shape[1]: - attn_kwargs["enable_gqa"] = True + num_key_value_groups = query.shape[1] // key.shape[1] + key = key[:, :, None, :, :].expand( + batch_size, key.shape[1], num_key_value_groups, seq_len, self.head_dim + ) + value = value[:, :, None, :, :].expand( + batch_size, + value.shape[1], + num_key_value_groups, + seq_len, + self.head_dim, + ) + key = key.reshape(batch_size, query.shape[1], seq_len, self.head_dim) + value = value.reshape(batch_size, query.shape[1], seq_len, self.head_dim) + attn_output = torch.nn.functional.scaled_dot_product_attention( - query, key, value, **attn_kwargs + query, + key, + value, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=False, + scale=self.scaling, ) attn_output = attn_output.transpose(1, 2) 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 8e57ab99c..3259ab363 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -193,6 +193,7 @@ class LTX2SigmaPreparationStage(PipelineStage): int(batch.num_inference_steps), number_of_tokens=latent_num_frames * latent_height * latent_width, ) + batch.sigmas.append(0.0011) else: batch.sigmas = build_official_ltx2_sigmas( int(batch.num_inference_steps) @@ -631,13 +632,14 @@ class LTX2SnapshotResidencyStrategy(LTX2TwoStageResidencyStrategy): def _pin_stage1_transformer_if_beneficial(self) -> None: """Optionally pin stage-1 DiT on GPU to remove first-stage cold H2D stall. - We only do this on high-VRAM CUDA machines with CPU offload enabled and - without FSDP inference. It trades extra steady-state VRAM for lower - request latency before the first denoise step. + We only do this outside low-VRAM mode on high-VRAM CUDA machines with + CPU offload enabled and without FSDP inference. It trades extra + steady-state VRAM for lower request latency before the first denoise step. """ if ( not self.server_args.dit_cpu_offload or self.server_args.use_fsdp_inference + or self._snapshot_low_vram_mode or not current_platform.is_cuda() or current_platform.get_device_total_memory() / BYTES_PER_GB < 70 ): @@ -934,6 +936,11 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): set_lora_kwargs["merge_weights"] = ( self._should_merge_stage2_distilled_lora(self.server_args) ) + elif phase == "stage1" and self.pipeline_name == "LTX2TwoStageHQPipeline": + # Official HQ also builds stage 1 with distilled LoRA fused. + set_lora_kwargs["merge_weights"] = ( + self._should_merge_stage2_distilled_lora(self.server_args) + ) self.set_lora( **set_lora_kwargs, ) 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 1f22533da..df26a38b5 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 @@ -361,11 +361,9 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): scheduler_sigmas = self.distilled_sigmas scheduler.sigmas = scheduler_sigmas - num_steps = len(self.distilled_sigmas) - 1 + num_steps = len(scheduler_sigmas) - 1 scheduler.num_inference_steps = num_steps - scheduler.timesteps = (self.distilled_sigmas[:num_steps] * 1000).to( - distilled_device - ) + scheduler.timesteps = (scheduler_sigmas[:num_steps] * 1000).to(distilled_device) scheduler._step_index = None scheduler._begin_index = None 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 54dba7f21..c6a0fb492 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 @@ -1,3 +1,4 @@ +import math from contextlib import contextmanager from dataclasses import dataclass, field @@ -281,9 +282,7 @@ class LTX2DenoisingStage(DenoisingStage): ) if generator is None: raise ValueError("LTX-2 res2s noise generator was not initialized.") - return cls._ltx2_res2s_new_noise(reference_tensor, generator).to( - dtype=reference_tensor.dtype - ) + return cls._ltx2_res2s_new_noise(reference_tensor, generator) @staticmethod def _ltx2_apply_clean_latent_mask( @@ -319,6 +318,34 @@ class LTX2DenoisingStage(DenoisingStage): b1 = cls._ltx2_phi_1(-h) - b2 return a21, b1, b2 + @staticmethod + def _ltx2_phi_scalar(j: int, neg_h: float) -> float: + if abs(neg_h) < 1e-10: + return 1.0 / math.factorial(j) + remainder = sum(neg_h**k / math.factorial(k) for k in range(j)) + return (math.exp(neg_h) - remainder) / (neg_h**j) + + @classmethod + def _ltx2_get_res2s_coefficients_scalar( + cls, h: float, c2: float = 0.5 + ) -> tuple[float, float, float]: + a21 = c2 * cls._ltx2_phi_scalar(1, -h * c2) + b2 = cls._ltx2_phi_scalar(2, -h) / c2 + b1 = cls._ltx2_phi_scalar(1, -h) - b2 + return a21, b1, b2 + + @staticmethod + def _ltx2_res2s_step_size_scalar( + sigma: torch.Tensor, sigma_next: torch.Tensor + ) -> float: + return float( + ( + -torch.log( + sigma_next.detach().double().cpu() / sigma.detach().double().cpu() + ) + ).item() + ) + @staticmethod def _ltx2_get_sde_coeff( sigma_next: torch.Tensor, @@ -328,31 +355,29 @@ class LTX2DenoisingStage(DenoisingStage): sigma_max: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if sigma_down is not None: - alpha_ratio = (1.0 - sigma_next) / (1.0 - sigma_down) - sigma_up = torch.sqrt( - torch.clamp( - sigma_next.square() - sigma_down.square() * alpha_ratio.square(), - min=0.0, - ) - ) + alpha_ratio = (1 - sigma_next) / (1 - sigma_down) + sigma_up = (sigma_next**2 - sigma_down**2 * alpha_ratio**2).clamp( + min=0 + ) ** 0.5 elif sigma_up is not None: - sigma_up = torch.minimum(sigma_up, sigma_next * 0.9999) + sigma_up.clamp_(max=sigma_next * 0.9999) sigmax = sigma_max if sigma_max is not None else torch.ones_like(sigma_next) sigma_signal = sigmax - sigma_next - sigma_residual = torch.sqrt( - torch.clamp(sigma_next.square() - sigma_up.square(), min=0.0) - ) + sigma_residual = (sigma_next**2 - sigma_up**2).clamp(min=0) ** 0.5 alpha_ratio = sigma_signal + sigma_residual sigma_down = sigma_residual / alpha_ratio else: alpha_ratio = torch.ones_like(sigma_next) sigma_down = sigma_next sigma_up = torch.zeros_like(sigma_next) - return ( - torch.nan_to_num(alpha_ratio), - torch.nan_to_num(sigma_down), - torch.nan_to_num(sigma_up), + + sigma_up = torch.nan_to_num( + sigma_up if sigma_up is not None else torch.zeros_like(sigma_next), 0.0 ) + nan_mask = torch.isnan(sigma_down) + sigma_down[nan_mask] = sigma_next[nan_mask].to(sigma_down.dtype) + alpha_ratio = torch.nan_to_num(alpha_ratio, 1.0) + return alpha_ratio, sigma_down, sigma_up @classmethod def _ltx2_res2s_sde_step( @@ -391,6 +416,8 @@ class LTX2DenoisingStage(DenoisingStage): sigma_next: torch.Tensor, model_video_velocity: torch.Tensor, model_audio_velocity: torch.Tensor, + model_video_timestep: torch.Tensor | None, + model_audio_timestep: torch.Tensor | None, midpoint_model_call, ) -> tuple[torch.Tensor, torch.Tensor]: """res2s RK2 step for unguided stage-2 refinement (HQ pipeline). @@ -407,10 +434,24 @@ class LTX2DenoisingStage(DenoisingStage): denoised_video = ctx.latents.float() denoised_audio = ctx.audio_latents.float() else: - denoised_video = ctx.latents.float() - sigma * model_video_velocity.float() - denoised_audio = ( - ctx.audio_latents.float() - sigma * model_audio_velocity.float() + video_sigma_for_x0 = ( + model_video_timestep + if ctx.use_ltx23_hq_timestep_semantics + and model_video_timestep is not None + else sigma ) + audio_sigma_for_x0 = ( + model_audio_timestep + if ctx.use_ltx23_hq_timestep_semantics + and model_audio_timestep is not None + else sigma + ) + denoised_video = self._ltx2_velocity_to_x0( + ctx.latents, model_video_velocity, video_sigma_for_x0 + ).float() + denoised_audio = self._ltx2_velocity_to_x0( + ctx.audio_latents, model_audio_velocity, audio_sigma_for_x0 + ).float() if sigma_val == 0.0 or sigma_next_val == 0.0: next_video = denoised_video.to(dtype=ctx.latents.dtype) @@ -420,8 +461,14 @@ class LTX2DenoisingStage(DenoisingStage): sigma_d = sigma.double() sigma_next_d = sigma_next.double() - h = -torch.log(torch.clamp(sigma_next_d / sigma_d, min=1e-12)) - a21, b1, b2 = self._ltx2_get_res2s_coefficients(h) + if ctx.use_ltx23_hq_timestep_semantics: + h = self._ltx2_res2s_step_size_scalar(sigma_d, sigma_next_d) + a21, b1, b2 = self._ltx2_get_res2s_coefficients_scalar(h) + h_value = h + else: + h = -torch.log(torch.clamp(sigma_next_d / sigma_d, min=1e-12)) + a21, b1, b2 = self._ltx2_get_res2s_coefficients(h) + h_value = float(h.item()) sub_sigma = torch.sqrt(torch.clamp(sigma_d * sigma_next_d, min=0.0)) anchor_video = ctx.latents.double() @@ -433,12 +480,12 @@ class LTX2DenoisingStage(DenoisingStage): midpoint_audio_det = anchor_audio + h * a21 * eps1_audio sub_noise_video = ( - self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True).float() + self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True) if ctx.use_native_hq_res2s_sde_noise else self._randn_like_with_batch_generators(ctx.latents, batch).float() ) sub_noise_audio = ( - self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True).float() + self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True) if ctx.use_native_hq_res2s_sde_noise else self._randn_like_with_batch_generators( ctx.audio_latents, batch @@ -460,15 +507,15 @@ class LTX2DenoisingStage(DenoisingStage): noise=sub_noise_audio, terminal=False, ) - midpoint_video_latents = self._ltx2_apply_clean_latent_mask( + midpoint_video_model_latents = self._ltx2_apply_clean_latent_mask( midpoint_video_latents.to(dtype=ctx.latents.dtype), ctx ) - midpoint_audio_latents = midpoint_audio_latents.to( + midpoint_audio_model_latents = midpoint_audio_latents.to( dtype=ctx.audio_latents.dtype ) # Bongmath anchor refinement for the first stage-2 step. - if float(h.item()) < 0.5 and sigma_val > 0.03: + if h_value < 0.5 and sigma_val > 0.03: x_mid_v = midpoint_video_latents.double() x_mid_a = midpoint_audio_latents.double() for _ in range(100): @@ -477,12 +524,26 @@ class LTX2DenoisingStage(DenoisingStage): anchor_audio = x_mid_a - h * a21 * eps1_audio eps1_audio = denoised_audio.double() - anchor_audio - mid_v, mid_a = midpoint_model_call( - midpoint_video_latents, midpoint_audio_latents, sub_sigma + mid_v, mid_a, mid_video_timestep, mid_audio_timestep = midpoint_model_call( + midpoint_video_model_latents, midpoint_audio_model_latents, sub_sigma ) - midpoint_denoised_video = midpoint_video_latents.float() - sub_sigma * mid_v - midpoint_denoised_audio = midpoint_audio_latents.float() - sub_sigma * mid_a + mid_video_sigma_for_x0 = ( + mid_video_timestep + if ctx.use_ltx23_hq_timestep_semantics and mid_video_timestep is not None + else sub_sigma + ) + mid_audio_sigma_for_x0 = ( + mid_audio_timestep + if ctx.use_ltx23_hq_timestep_semantics and mid_audio_timestep is not None + else sub_sigma + ) + midpoint_denoised_video = self._ltx2_velocity_to_x0( + midpoint_video_latents, mid_v, mid_video_sigma_for_x0 + ).float() + midpoint_denoised_audio = self._ltx2_velocity_to_x0( + midpoint_audio_latents, mid_a, mid_audio_sigma_for_x0 + ).float() eps2_video = midpoint_denoised_video.double() - anchor_video eps2_audio = midpoint_denoised_audio.double() - anchor_audio @@ -491,30 +552,34 @@ class LTX2DenoisingStage(DenoisingStage): next_audio_det = anchor_audio + h * (b1 * eps1_audio + b2 * eps2_audio) step_noise_video = ( - self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False).float() + self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False) if ctx.use_native_hq_res2s_sde_noise else self._randn_like_with_batch_generators(ctx.latents, batch).float() ) step_noise_audio = ( - self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False).float() + self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False) if ctx.use_native_hq_res2s_sde_noise else self._randn_like_with_batch_generators( ctx.audio_latents, batch ).float() ) + sde_sigma = sigma if ctx.use_ltx23_hq_timestep_semantics else sigma_d + sde_sigma_next = ( + sigma_next if ctx.use_ltx23_hq_timestep_semantics else sigma_next_d + ) next_video = self._ltx2_res2s_sde_step( sample=anchor_video, denoised_sample=next_video_det, - sigma=sigma_d, - sigma_next=sigma_next_d, + sigma=sde_sigma, + sigma_next=sde_sigma_next, noise=step_noise_video, terminal=False, ) next_audio = self._ltx2_res2s_sde_step( sample=anchor_audio, denoised_sample=next_audio_det, - sigma=sigma_d, - sigma_next=sigma_next_d, + sigma=sde_sigma, + sigma_next=sde_sigma_next, noise=step_noise_audio, terminal=False, ) @@ -1345,7 +1410,12 @@ class LTX2DenoisingStage(DenoisingStage): video_latents: torch.Tensor, audio_latents: torch.Tensor, sigma_value: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + ]: original_video_latents = ctx.latents original_audio_latents = ctx.audio_latents ctx.latents = video_latents @@ -1419,7 +1489,12 @@ class LTX2DenoisingStage(DenoisingStage): mid_a_u, mid_a_t = mid_a.chunk(2) mid_v = mid_v_u + batch.guidance_scale * (mid_v_t - mid_v_u) mid_a = mid_a_u + batch.guidance_scale * (mid_a_t - mid_a_u) - return mid_v, mid_a + return ( + mid_v, + mid_a, + model_inputs_local.timestep_video, + model_inputs_local.timestep_audio, + ) finally: ctx.latents = original_video_latents ctx.audio_latents = original_audio_latents @@ -1431,6 +1506,8 @@ class LTX2DenoisingStage(DenoisingStage): sigma_next=sigma_next, model_video_velocity=model_video, model_audio_velocity=model_audio, + model_video_timestep=model_inputs.timestep_video, + model_audio_timestep=model_inputs.timestep_audio, midpoint_model_call=_stage2_midpoint_model_call, ) else: @@ -1748,23 +1825,28 @@ class LTX2DenoisingStage(DenoisingStage): v_mod, a_v_mod = pass_outputs.get("modality", (None, None)) sigma_value_float = float(sigma_value.item()) - video_sigma_for_x0: float | torch.Tensor = sigma_value_float - if ctx.denoise_mask is not None: + audio_sigma_for_x0: float | torch.Tensor = sigma_value_float + if ctx.use_ltx23_hq_timestep_semantics: + video_sigma_for_x0 = model_inputs_local.timestep_video + audio_sigma_for_x0 = model_inputs_local.timestep_audio + elif ctx.denoise_mask is not None: video_sigma_for_x0 = sigma_value.to( device=video_latents.device, dtype=torch.float32 ) * ctx.denoise_mask.squeeze(-1) + else: + video_sigma_for_x0 = sigma_value_float denoised_video_local = self._ltx2_velocity_to_x0( video_latents, v_pos, video_sigma_for_x0 ) denoised_audio_local = self._ltx2_velocity_to_x0( - audio_latents, a_v_pos, sigma_value_float + audio_latents, a_v_pos, audio_sigma_for_x0 ) denoised_video_neg = self._ltx2_velocity_to_x0( video_latents, v_neg, video_sigma_for_x0 ) denoised_audio_neg = self._ltx2_velocity_to_x0( - audio_latents, a_v_neg, sigma_value_float + audio_latents, a_v_neg, audio_sigma_for_x0 ) denoised_video_perturbed = ( None @@ -1777,7 +1859,7 @@ class LTX2DenoisingStage(DenoisingStage): None if a_v_ptb is None else self._ltx2_velocity_to_x0( - audio_latents, a_v_ptb, sigma_value_float + audio_latents, a_v_ptb, audio_sigma_for_x0 ) ) denoised_video_modality = ( @@ -1791,7 +1873,7 @@ class LTX2DenoisingStage(DenoisingStage): None if a_v_mod is None else self._ltx2_velocity_to_x0( - audio_latents, a_v_mod, sigma_value_float + audio_latents, a_v_mod, audio_sigma_for_x0 ) ) @@ -1867,8 +1949,14 @@ class LTX2DenoisingStage(DenoisingStage): else: sigma_d = sigma.double() sigma_next_d = sigma_next.double() - h = -torch.log(torch.clamp(sigma_next_d / sigma_d, min=1e-12)) - a21, b1, b2 = self._ltx2_get_res2s_coefficients(h) + if ctx.use_ltx23_hq_timestep_semantics: + h = self._ltx2_res2s_step_size_scalar(sigma_d, sigma_next_d) + a21, b1, b2 = self._ltx2_get_res2s_coefficients_scalar(h) + h_value = h + else: + h = -torch.log(torch.clamp(sigma_next_d / sigma_d, min=1e-12)) + a21, b1, b2 = self._ltx2_get_res2s_coefficients(h) + h_value = float(h.item()) sub_sigma = torch.sqrt(torch.clamp(sigma_d * sigma_next_d, min=0.0)) anchor_video = ctx.latents.double() @@ -1880,16 +1968,14 @@ class LTX2DenoisingStage(DenoisingStage): midpoint_audio_deterministic = anchor_audio + h * a21 * eps1_audio substep_video_noise = ( - self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True).float() + self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True) if ctx.use_native_hq_res2s_sde_noise else self._randn_like_with_batch_generators( ctx.latents, batch ).float() ) substep_audio_noise = ( - self._ltx2_res2s_noise_like( - ctx.audio_latents, ctx, substep=True - ).float() + self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True) if ctx.use_native_hq_res2s_sde_noise else self._randn_like_with_batch_generators( ctx.audio_latents, batch @@ -1913,15 +1999,15 @@ class LTX2DenoisingStage(DenoisingStage): terminal=False, ) - midpoint_video_latents = self._ltx2_apply_clean_latent_mask( + midpoint_video_model_latents = self._ltx2_apply_clean_latent_mask( midpoint_video_latents.to(dtype=ctx.latents.dtype), ctx, ) - midpoint_audio_latents = midpoint_audio_latents.to( + midpoint_audio_model_latents = midpoint_audio_latents.to( dtype=ctx.audio_latents.dtype ) - if float(h.item()) < 0.5 and sigma_val > 0.03: + if h_value < 0.5 and sigma_val > 0.03: x_mid_v = midpoint_video_latents.double() x_mid_a = midpoint_audio_latents.double() for _ in range(100): @@ -1932,8 +2018,8 @@ class LTX2DenoisingStage(DenoisingStage): midpoint_denoised_video, midpoint_denoised_audio = ( evaluate_stage1_guided_x0( - video_latents=midpoint_video_latents, - audio_latents=midpoint_audio_latents, + video_latents=midpoint_video_model_latents, + audio_latents=midpoint_audio_model_latents, sigma_value=sub_sigma, update_skip_cache=False, ) @@ -1949,34 +2035,36 @@ class LTX2DenoisingStage(DenoisingStage): ) step_video_noise = ( - self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False).float() + self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False) if ctx.use_native_hq_res2s_sde_noise else self._randn_like_with_batch_generators( ctx.latents, batch ).float() ) step_audio_noise = ( - self._ltx2_res2s_noise_like( - ctx.audio_latents, ctx, substep=False - ).float() + self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False) if ctx.use_native_hq_res2s_sde_noise else self._randn_like_with_batch_generators( ctx.audio_latents, batch ).float() ) + sde_sigma = sigma if ctx.use_ltx23_hq_timestep_semantics else sigma_d + sde_sigma_next = ( + sigma_next if ctx.use_ltx23_hq_timestep_semantics else sigma_next_d + ) next_video_latents = self._ltx2_res2s_sde_step( sample=anchor_video, denoised_sample=next_video_deterministic, - sigma=sigma_d, - sigma_next=sigma_next_d, + sigma=sde_sigma, + sigma_next=sde_sigma_next, noise=step_video_noise, terminal=False, ) next_audio_latents = self._ltx2_res2s_sde_step( sample=anchor_audio, denoised_sample=next_audio_deterministic, - sigma=sigma_d, - sigma_next=sigma_next_d, + sigma=sde_sigma, + sigma_next=sde_sigma_next, noise=step_audio_noise, terminal=False, ) diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index ac3aab1d4..6e31b35b9 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -23,6 +23,7 @@ from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.models.encoders import T5Config from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + LTX2PipelineConfig, is_ltx23_native_variant, ) from sglang.multimodal_gen.configs.quantization.nunchaku import NunchakuSVDQuantArgs @@ -487,6 +488,22 @@ class ServerArgs(DisaggArgsMixin): self._parse_attention_backend_config(self.attention_backend_config) ) + if self.backend != Backend.DIFFUSERS and isinstance( + self.pipeline_config, LTX2PipelineConfig + ): + text_backend = self.component_attention_backends.get("text_encoder") + if text_backend != "torch_sdpa": + if text_backend is None: + logger.info( + "Automatically set torch_sdpa backend for component text_encoder to preserve LTX2 official attention semantics" + ) + else: + logger.warning( + "Overriding %s backend with torch_sdpa for component text_encoder to preserve LTX2 official attention semantics", + text_backend, + ) + self.component_attention_backends["text_encoder"] = "torch_sdpa" + if self.ring_degree > 1: if self.attention_backend is not None and self.attention_backend not in ( "fa", diff --git a/python/sglang/multimodal_gen/test/server/consistency_threshold.json b/python/sglang/multimodal_gen/test/server/consistency_threshold.json index 33a93f94d..f3e85286f 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_threshold.json +++ b/python/sglang/multimodal_gen/test/server/consistency_threshold.json @@ -103,6 +103,12 @@ "psnr_threshold": 25.0, "mean_abs_diff_threshold": 8.0 }, + "ltx_2_two_stage_t2v": { + "clip_threshold": 0.90, + "ssim_threshold": 0.89, + "psnr_threshold": 24.0, + "mean_abs_diff_threshold": 10.0 + }, "ltx_2.3_one_stage_ti2v": { "clip_threshold": 0.64, "ssim_threshold": 0.42, @@ -244,8 +250,8 @@ "ltx_2_3_hq_pipeline": { "clip_threshold": 0.78, "ssim_threshold": 0.48, - "psnr_threshold": 12.0, - "mean_abs_diff_threshold": 48.0 + "psnr_threshold": 13.0, + "mean_abs_diff_threshold": 45.0 }, "ltx_2_3_two_stage_ti2v_2gpus": { "clip_threshold": 0.92, diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 4a708e80e..5e96bbcb0 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -33,7 +33,7 @@ if TYPE_CHECKING: logger = init_logger(__name__) -SGL_TEST_FILES_CI_DATA_REVISION = "4d9eff3b05b0ffe1d3529e8bb148b63af11a4b92" +SGL_TEST_FILES_CI_DATA_REVISION = "6e7b99e16b857c98285277fe3b4ffef30559bde9" SGL_TEST_FILES_CONSISTENCY_GT_ROOT = ( "https://raw.githubusercontent.com/" f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/" @@ -50,6 +50,14 @@ SGL_TEST_FILES_CONSISTENCY_GT_BASES = ( SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_BASE, SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE, ) +# Keep non-comparable LTX CI scenarios on sglang_generated rather than hiding +# remaining semantic gaps behind very loose official thresholds. +SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES = frozenset( + { + "ltx_2.3_one_stage_ti2v", + "ltx_2.3_two_stage_t2v_2gpus", + } +) CONSISTENCY_THRESHOLD_JSON_PATH = ( Path(__file__).resolve().parent / "server" / "consistency_threshold.json" ) @@ -955,10 +963,28 @@ def _remote_consistency_gt_candidates( def _remote_file_exists(url: str) -> bool: - try: - return requests.head(url, timeout=10, allow_redirects=True).status_code == 200 - except requests.RequestException: - return False + for method in ("head", "get"): + try: + if method == "head": + resp = requests.head(url, timeout=10, allow_redirects=True) + else: + resp = requests.get( + url, + timeout=10, + allow_redirects=True, + headers={"Range": "bytes=0-0"}, + stream=True, + ) + try: + if resp.status_code in (200, 206): + return True + if resp.status_code not in (403, 405, 429) and resp.status_code < 500: + return False + finally: + resp.close() + except requests.RequestException: + pass + return False def _find_remote_consistency_gt_files( @@ -967,7 +993,12 @@ def _find_remote_consistency_gt_files( is_video: bool, output_format: str | None = None, ) -> list[tuple[str, str]]: - for base_url in SGL_TEST_FILES_CONSISTENCY_GT_BASES: + if case_id in SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES: + bases = SGL_TEST_FILES_CONSISTENCY_GT_BASES + else: + # Avoid accidentally comparing non-comparable CI cases against official GT. + bases = (SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE,) + for base_url in bases: candidates = _remote_consistency_gt_candidates( base_url, case_id, num_gpus, is_video, output_format )