[diffusion] chore: align LTX-2 with official (#24313)

2026-05-07 08:46:28 +08:00
parent 7fa48916f5
commit 2e642ea187
11 changed files with 369 additions and 208 deletions
@@ -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

@@ -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
@@ -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:
@@ -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)
@@ -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)

@@ -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,
            )
@@ -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

@@ -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,
                )
@@ -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",
@@ -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,
@@ -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
        )