Motivation

Support LTX-2 Video & Audio Joint model.

This PR mainly fix bugs in Sglang LTX-2 pipelines

How to use

generate

• T2V

sglang generate \
  --model-path LTX-2_Model \
  --prompt "An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot. [Audio Description] Subtle ambient space atmosphere with deep bass hum, gentle cracking and crunching sounds as the eggshell breaks apart, soft whooshing of lunar dust particles floating in low gravity, quiet breathing sounds from inside the spacesuit, a faint heartbeat rhythm building tension, ethereal cinematic synth pads creating a sense of wonder and discovery, all blended with vast empty silence of space." \
  --negative-prompt "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \
  --width 768 --height 512 --num-frames 121 --fps 24 \
  --num-inference-steps 40 --guidance-scale 4.0 \
  --seed 1024 \
  --output-path ./output/ --output-file-name ltx2_sample.mp4

• TI2V

sglang generate \
  --model-path LTX-2_Model \
  --prompt "An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot. [Audio Description] Subtle ambient space atmosphere with deep bass hum, gentle cracking and crunching sounds as the eggshell breaks apart, soft whooshing of lunar dust particles floating in low gravity, quiet breathing sounds from inside the spacesuit, a faint heartbeat rhythm building tension, ethereal cinematic synth pads creating a sense of wonder and discovery, all blended with vast empty silence of space." \
  --image-path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" \
  --negative-prompt "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \
  --width 768 --height 512 --num-frames 121 --fps 24 \
  --num-inference-steps 40 --guidance-scale 4.0 \
  --seed 1024 \
  --output-path ./output/ --output-file-name ltx2_sample.mp4

• SP, TP & Cache DiT (8 GPUs)

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
SGLANG_CACHE_DIT_ENABLED=true \
SGLANG_CACHE_DIT_SCM_PRESET=medium \
sglang generate \
  --model-path LTX-2_Model \
  --num-gpus 8 \
  --tp-size 2 \
  --sp-degree 4 \
  --ulysses-degree 2 \
  --ring-degree 2 \
  --prompt "An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot. [Audio Description] Subtle ambient space atmosphere with deep bass hum, gentle cracking and crunching sounds as the eggshell breaks apart, soft whooshing of lunar dust particles floating in low gravity, quiet breathing sounds from inside the spacesuit, a faint heartbeat rhythm building tension, ethereal cinematic synth pads creating a sense of wonder and discovery, all blended with vast empty silence of space." \
  --image-path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" \
  --negative-prompt "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \
  --width 768 --height 512 --num-frames 121 --fps 24 \
  --num-inference-steps 40 --guidance-scale 4.0 \
  --seed 1024 \
  --output-path ./output/ --output-file-name ltx2_sample.mp4

Bugs fix

configs/models/dits/ltx_2.py

• Changed audio_positional_embedding_max_pos default from [2048] to [20] to match Official LTX-2 audio positional embedding configuration

configs/pipeline_configs/ltx_2.py

• Added prepare_latent_shape() method to return packed latent shape [B, seq, C] directly, matching Official's latent preparation
• Changed prepare_audio_latent_shape() return shape from (B, C, L, M) to (B, L, C*M) to match Official's packed audio latent format
• Changed prepare_sigmas() from np.linspace(1.0, 1.0/steps, steps) to torch.linspace(1.0, 0.0, steps+1)[:-1] to match Official's sigma schedule
• Added early return in maybe_pack_latents() to skip packing if latents are already 3D [B, seq, C]
• Added early return in maybe_pack_audio_latents() to skip packing if audio latents are already 3D [B, T, C*F]

runtime/models/adapter/ltx_2_connector.py

• Removed .float() upcast in apply_split_rotary_emb to keep computation in native dtype (bfloat16) instead of upcasting to float32
• Added dtype parameter to LTX2RotaryPosEmbed1d.forward() to allow explicit dtype control for cos/sin freqs
• Added dtype casting for cos_freqs and sin_freqs when dtype is specified
• Changed self.rope() call to pass dtype=hidden_states.dtype for consistent precision

runtime/models/dits/ltx_2.py

• Removed .float() upcast in apply_split_rotary_emb to keep rotary embedding computation in native dtype
• Added self.coords_dtype attribute: bfloat16 for video, float32 for audio, to control coordinate precision per modality
• Added explicit coords = coords.to(self.coords_dtype) before coordinate averaging
• Changed RoPE output to explicitly cast to torch.bfloat16 instead of returning in computed dtype
• Changed rope_double_precision to read from HuggingFace config first (hf_config.get(...)) with arch config as fallback
• Changed adaln_single calls to pass hidden_dtype=hidden_states.dtype for explicit dtype control in AdaLN timestamp embeddings
• Changed timestep * ts_ca_mult to timestep * self.av_ca_timestep_scale_multiplier to avoid intermediate float64 computation
• Wrapped self.norm_out() and self.audio_norm_out() in torch.autocast(enabled=False) to disable autocast during final layer norm, ensuring deterministic precision

runtime/models/encoders/gemma_3.py

• Added _rotate_half() helper function implementing HuggingFace's rotate_half formula
• Added CPU-precomputed _hf_inv_freq buffer in Gemma3Attention.__init__() to match HuggingFace's exact inverse frequency computation
• Added custom rotary_emb() method that uses _hf_inv_freq with rotate_half formula, replacing SGLang's shared RotaryEmbedding which uses a slightly different numerical formula
• Changed attention mask dtype from torch.float32 to query.dtype for consistent precision
• Changed float("-inf") mask fill to torch.finfo(query.dtype).min to avoid overflow in bfloat16
• Changed embed_scale multiplication to use explicit torch.tensor(self.embed_scale, device, dtype) to avoid precision mismatch
• Commented out position_ids = position_ids + 1 to match Official LTX-2's position indexing

runtime/pipelines/ltx_2_pipeline.py

• Added LTX2FlowMatchScheduler class that overrides _time_shift_exponential to use torch float32 arithmetic instead of numpy float64, matching Official LTX-2's scheduler precision
• Added initialize_pipeline() method to replace the default FlowMatchEulerDiscreteScheduler with LTX2FlowMatchScheduler
• Changed prepare_mu() to use fixed token count 4096 (matching Official's MAX_SHIFT_ANCHOR default) instead of computing from actual latent dimensions

runtime/pipelines_core/stages/denoising_av.py

• Changed image conditioning encode path to use encode_dtype = batch.latents.dtype (bfloat16) for both input tensor and VAE weights, instead of hardcoded torch.float32
• Added VAE dtype save/restore: saves original_dtype from vae_precision, converts VAE to encode_dtype before encoding, restores original_dtype after per-channel normalization
• Autocast remains disabled during encode (since vae_precision stays "fp32", vae_autocast_enabled is False), matching Official's pure bf16 encode without autocast

runtime/pipelines_core/stages/decoding_av.py

• Added self.vae.to(original_dtype) after decode to restore VAE dtype (previously VAE stayed in bfloat16 after the forced .to(torch.bfloat16) conversion)

Precision Alignment [verified on a single GPU]

To reproduce pixel-perfect alignment with Official LTX-2 (ltx-pipelines / ltx-core), the following deterministic torch settings must be applied to both pipelines before any model execution:

import torch
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(False)
torch.backends.cuda.enable_cudnn_sdp(False)

These settings disable FlashAttention and Memory-Efficient SDPA backends, forcing PyTorch to use the deterministic Math SDPA backend. Without them, non-deterministic attention kernels produce different results across runs even with identical inputs and seeds.

How to apply

Both Official and SGLang pipelines must load these settings before any torch operation. This is done via a sitecustomize.py + torch_settings.py pair placed on PYTHONPATH:

mkdir -p /path/to/_site

cat > /path/to/_site/torch_settings.py << 'EOF'
import torch
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(False)
torch.backends.cuda.enable_cudnn_sdp(False)
EOF

cat > /path/to/_site/sitecustomize.py << 'EOF'
import torch_settings
EOF

Then prepend this directory to PYTHONPATH when launching either pipeline:

PYTHONPATH="/path/to/_site:$PYTHONPATH" python -m ltx_pipelines.ti2vid_one_stage ...
PYTHONPATH="/path/to/_site:$PYTHONPATH" sglang generate ...

Python automatically imports sitecustomize.py at startup, which triggers torch_settings.py and applies the deterministic settings before any model code runs.

Verified alignment

With these settings applied and identical generation parameters (prompt, seed, resolution, num_frames, guidance_scale, num_inference_steps, image conditioning), the two pipelines produce bit-for-bit identical float pixel outputs (max_diff = 0.0).

Checklist

LTX-2

• Accuracy alignment achieved
• Support TI2V inference
• Implement Gemma 3 support using SGLang (not Transformers)
• Gemma 3: support TP (tensor parallelism)
• Support RoPE / SP
• Wire up / validate Cache-DiT
• Support TP for the primary model
• Verify correctness of generate and the server
• Ensure structural and code-style consistency across modules
• Polish code style; remove redundant comments
• Clean up debug code
• Fix lint and update to latest main
• Follow up on PR reviews

LTX-2 Stack PR

• Support one-stage generation
• Support distilled generation
• Support IC LoRA
• Support keyframe_interpolation
• Support two-stage generation

Review Process

1. Ping Merge Oncalls to start the PR flow. See [PR Merge Process].
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • \tag-run-ci-label, \rerun-failed-ci, \tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.