[Feat] Support online quantization (fp8/int8) for LTX-2 (#3700)

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>

benchmarks/diffusion/quantization_quality.py

@@ -34,6 +34,16 @@ Video example (text-to-video):
         --height 720 --width 1280 \
         --num-frames 81 --num-inference-steps 40 --seed 42
 
+LTX-2 example (text-to-video; audio output is generated but not scored by LPIPS):
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Lightricks/LTX-Video-2 \
+        --task t2v \
+        --quantization fp8 int8 \
+        --prompts \
+            "A serene lakeside sunrise with mist over the water" \
+        --height 704 --width 1216 \
+        --num-frames 121 --num-inference-steps 40 --seed 42
+
 Multiple quantization methods:
     python benchmarks/diffusion/quantization_quality.py \
         --model Tongyi-MAI/Z-Image-Turbo \
@@ -204,10 +214,19 @@ def _generate_video(omni, args, prompt, seed):
         else:
             frames = inner
     elif hasattr(first, "images") and first.images:
-        frames = first.images
+        frames = first.images[0]
     else:
         raise ValueError("Could not extract video frames from output.")
 
+    # LTX-2 (and similar audio+video models) may surface a dict or (video, audio) tuple
+    if isinstance(frames, dict):
+        frames = frames.get("video") or frames.get("frames")
+    elif isinstance(frames, tuple) and len(frames) == 2:
+        frames = frames[0]
+
+    if frames is None:
+        raise ValueError("Could not extract video frames from output.")
+
     if isinstance(frames, torch.Tensor):
         video = frames.detach().cpu()
         if video.dim() == 5:
@@ -225,9 +244,11 @@ def _generate_video(omni, args, prompt, seed):
     return frames_array, elapsed, peak_mem
 
 
-def _unload_omni(omni):
-    """Delete Omni instance and free GPU memory."""
-    del omni
+def _free_gpu_memory():
+    """Force GC and release cached GPU memory.
+
+    Must be called AFTER the caller has dropped (i.e., via `del`)
+    every reference to the Omni instance"""
     gc.collect()
     if torch.cuda.is_available():
         torch.accelerator.empty_cache()
@@ -267,7 +288,9 @@ def run_benchmark(args):
 
     bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
     bl_mem = baseline_outputs[prompts[0]][2]  # use first prompt's memory
-    _unload_omni(omni_bl)
+    omni_bl.shutdown()
+    del omni_bl
+    _free_gpu_memory()
 
     # Save baseline outputs
     bl_dir = output_dir / "baseline"
@@ -307,7 +330,9 @@ def run_benchmark(args):
 
         qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
         qt_mem = qt_outputs[prompts[0]][2]
-        _unload_omni(omni_qt)
+        omni_qt.shutdown()
+        del omni_qt
+        _free_gpu_memory()
 
         # Save quantized outputs
         qt_dir = output_dir / config_label.replace(" ", "_")

examples/offline_inference/text_to_video/text_to_video.py

@@ -237,6 +237,27 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
+def _extract_peak_memory_mb(result: Any) -> float:
+    """Pull worker-reported peak VRAM (MiB) generation result.
+
+    Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
+    """
+    if isinstance(result, list):
+        result = result[0] if result else None
+    if result is None:
+        return 0.0
+    val = getattr(result, "peak_memory_mb", 0.0)
+    if not val:
+        inner = getattr(result, "request_output", None)
+        if isinstance(inner, list):
+            inner = inner[0] if inner else None
+        val = getattr(inner, "peak_memory_mb", 0.0)
+    try:
+        return float(val or 0.0)
+    except (TypeError, ValueError):
+        return 0.0
+
+
 def main():
     args = parse_args()
     model_class_name = args.model_class_name
@@ -348,6 +369,10 @@ def main():
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
+    peak_mb = _extract_peak_memory_mb(frames)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
+
     audio = None
     if isinstance(frames, list):
         frames = frames[0] if frames else None

tests/diffusion/quantization/test_quantization_fp8.py

@@ -109,6 +109,85 @@ def _generate_single_stage_image(
         return images, peak_mem
 
 
+def _generate_single_stage_video(
+    model: str,
+    quantization: str | None = None,
+    height: int = 256,
+    width: int = 256,
+    num_frames: int = 25,
+    num_inference_steps: int = 8,
+    guidance_scale: float = 4.0,
+    seed: int = 42,
+    prompt: str = "A serene lakeside sunrise with mist over the water",
+    **extra_omni_kwargs: Any,
+) -> tuple[int, float]:
+    """Generate a t2v output with a single-stage diffusion model.
+
+    Returns (num_frames_produced, peak_memory_gib)
+    """
+    omni_kwargs: dict[str, Any] = dict(extra_omni_kwargs)
+    if quantization:
+        omni_kwargs["quantization"] = quantization
+
+    with OmniRunner(model, **omni_kwargs) as runner:
+        torch.accelerator.reset_peak_memory_stats()
+
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+        outputs = runner.omni.generate(
+            {"prompt": prompt, "negative_prompt": ""},
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+            ),
+        )
+
+        first = outputs[0]
+
+        # Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput.request_output).
+        frames: Any = None
+        if hasattr(first, "request_output") and isinstance(first.request_output, list):
+            inner = first.request_output[0]
+            if isinstance(inner, OmniRequestOutput) and inner.images:
+                frames = inner.images[0]
+        if frames is None and hasattr(first, "images") and first.images:
+            frames = first.images[0]
+        assert frames is not None, "No video frames returned from generate()"
+
+        # LTX-2 (audio+video) may surface (video, audio) tuples or {"video": ...} dicts.
+        if isinstance(frames, dict):
+            frames = frames.get("video") or frames.get("frames")
+        elif isinstance(frames, tuple) and len(frames) == 2:
+            frames = frames[0]
+        assert frames is not None, "Could not extract video frames from output"
+
+        if isinstance(frames, torch.Tensor):
+            video = frames.detach().cpu()
+            if video.dim() == 5:
+                video = video[0]
+            if video.dim() == 4 and video.shape[0] in (3, 4):
+                video = video.permute(1, 2, 3, 0)
+            num_frames_produced = int(video.shape[0])
+        else:
+            import numpy as np
+
+            arr = np.asarray(frames)
+            if arr.ndim == 5:
+                arr = arr[0]
+            num_frames_produced = int(arr.shape[0])
+
+        peak_mem_mb = getattr(first, "peak_memory_mb", None)
+        if peak_mem_mb:
+            peak_mem = float(peak_mem_mb) / 1024.0
+        else:
+            peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
+
+        return num_frames_produced, peak_mem
+
+
 def _generate_bagel_image(
     quantization_config: str | None = None,
     num_inference_steps: int = 15,
@@ -277,6 +356,25 @@ def test_single_stage_flux_fp8_uses_less_memory():
     assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
 
 
+@hardware_test(res={"cuda": "H100"})
+def test_single_stage_ltx2_fp8_uses_less_memory():
+    """FP8 should use less peak memory than BF16 for LTX-2."""
+    _, mem_bf16 = _generate_single_stage_video(
+        model="Lightricks/LTX-2",
+        quantization=None,
+    )
+    torch.accelerator.empty_cache()
+
+    _, mem_fp8 = _generate_single_stage_video(
+        model="Lightricks/LTX-2",
+        quantization="fp8",
+    )
+
+    print(f"LTX-2 BF16 peak memory: {mem_bf16:.2f} GiB")
+    print(f"LTX-2 FP8 peak memory:  {mem_fp8:.2f} GiB")
+    assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
+
+
 # ─── Multi-stage model tests (BAGEL) ─────────────────────────────────────────
 
 

tests/diffusion/quantization/test_quantization_quality.py

@@ -143,6 +143,18 @@ QUALITY_CONFIGS = [
         seed=142,
         num_inference_steps=20,
     ),
+    QualityTestConfig(
+        id="fp8_ltx2",
+        model="Lightricks/LTX-2",
+        quantization="fp8",
+        task="t2v",
+        prompt="A serene lakeside sunrise with mist over the water",
+        max_lpips=0.10,
+        height=256,
+        width=256,
+        num_frames=25,
+        num_inference_steps=8,
+    ),
 ]
 
 
@@ -261,10 +273,19 @@ def _generate_video(omni, config: QualityTestConfig):
         else:
             frames = inner
     elif hasattr(first, "images") and first.images:
-        frames = first.images
+        frames = first.images[0]
     else:
         raise ValueError("Could not extract video frames from output.")
 
+    # LTX-2 (audio+video) may surface (video, audio) tuples or {"video": ...} dicts
+    if isinstance(frames, dict):
+        frames = frames.get("video") or frames.get("frames")
+    elif isinstance(frames, tuple) and len(frames) == 2:
+        frames = frames[0]
+
+    if frames is None:
+        raise ValueError("Could not extract video frames from output.")
+
     if isinstance(frames, torch.Tensor):
         video = frames.detach().cpu()
         if video.dim() == 5:
@@ -275,7 +296,11 @@ def _generate_video(omni, config: QualityTestConfig):
             video = video.clamp(-1, 1) * 0.5 + 0.5
         return video.float().numpy(), peak_mem
 
-    return np.asarray(frames), peak_mem
+    frames_array = np.asarray(frames)
+    if frames_array.ndim == 5:
+        # strip the leading batch dim
+        frames_array = frames_array[0]
+    return frames_array, peak_mem
 
 
 def _compute_lpips(baseline, quantized, task: str) -> float:
@@ -318,8 +343,7 @@ def _compute_psnr_and_mae(baseline, quantized, task: str) -> tuple[float, float]
     return psnr, mae
 
 
-def _unload(omni):
-    del omni
+def _free_gpu_memory():
     gc.collect()
     if torch.cuda.is_available():
         torch.accelerator.empty_cache()
@@ -395,7 +419,9 @@ def test_quantization_quality(config: QualityTestConfig):
     # --- BF16 baseline ---
     omni_bl = Omni(model=config.baseline_ref())
     baseline_out, bl_mem = generate_fn(omni_bl, config)
-    _unload(omni_bl)
+    omni_bl.shutdown()
+    del omni_bl
+    _free_gpu_memory()
     _maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out)
 
     # --- Quantized ---
@@ -405,7 +431,9 @@ def test_quantization_quality(config: QualityTestConfig):
     else:
         omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization)
     quant_out, qt_mem = generate_fn(omni_qt, config)
-    _unload(omni_qt)
+    omni_qt.shutdown()
+    del omni_qt
+    _free_gpu_memory()
     _maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out)
 
     # --- Similarity metrics ---

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -17,7 +17,7 @@ import inspect
 from collections.abc import Iterable
 from dataclasses import dataclass
 from types import SimpleNamespace
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.distributed
@@ -39,6 +39,9 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
@@ -173,7 +176,16 @@ class LTX2AdaLayerNormSingle(nn.Module):
 
 
 class ColumnParallelApproxGELU(nn.Module):
-    def __init__(self, dim_in: int, dim_out: int, *, approximate: str, bias: bool = True):
+    def __init__(
+        self,
+        dim_in: int,
+        dim_out: int,
+        *,
+        approximate: str,
+        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ):
         super().__init__()
         self.proj = ColumnParallelLinear(
             dim_in,
@@ -181,6 +193,8 @@ class ColumnParallelApproxGELU(nn.Module):
             bias=bias,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj" if prefix else "proj",
         )
         self.approximate = approximate
 
@@ -200,6 +214,8 @@ class LTX2FeedForward(nn.Module):
         bias: bool = True,
         dropout: float = 0.0,
         final_dropout: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
 
@@ -211,13 +227,22 @@ class LTX2FeedForward(nn.Module):
         dropout_layer: nn.Module = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
 
         layers: list[nn.Module] = [
-            ColumnParallelApproxGELU(dim, inner_dim, approximate="tanh", bias=bias),
+            ColumnParallelApproxGELU(
+                dim,
+                inner_dim,
+                approximate="tanh",
+                bias=bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.net.0" if prefix else "net.0",
+            ),
             dropout_layer,
             RowParallelLinear(
                 inner_dim,
                 dim_out,
                 input_is_parallel=True,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.net.2" if prefix else "net.2",
             ),
         ]
         if final_dropout:
@@ -515,6 +540,9 @@ class LTX2Attention(torch.nn.Module):
         rope_type: str = "interleaved",
         apply_gated_attention: bool = False,
         processor=None,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+        disable_kv_quant: bool = False,
     ):
         super().__init__()
         # LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both
@@ -548,6 +576,8 @@ class LTX2Attention(torch.nn.Module):
                 head_size=self.head_dim,
                 total_num_heads=heads,
                 bias=bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
             )
             self.query_num_heads = self.to_qkv.num_heads
             self.kv_num_heads = self.to_qkv.num_kv_heads
@@ -562,6 +592,8 @@ class LTX2Attention(torch.nn.Module):
                 bias=bias,
                 gather_output=False,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_q" if prefix else "to_q",
             )
             self.to_k = ColumnParallelLinear(
                 self.cross_attention_dim,
@@ -569,6 +601,8 @@ class LTX2Attention(torch.nn.Module):
                 bias=bias,
                 gather_output=False,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_k" if prefix else "to_k",
             )
             self.to_v = ColumnParallelLinear(
                 self.cross_attention_dim,
@@ -576,6 +610,8 @@ class LTX2Attention(torch.nn.Module):
                 bias=bias,
                 gather_output=False,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_v" if prefix else "to_v",
             )
 
         self.heads = self.query_num_heads
@@ -617,6 +653,8 @@ class LTX2Attention(torch.nn.Module):
                     bias=out_bias,
                     input_is_parallel=True,
                     return_bias=False,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0" if prefix else "to_out.0",
                 ),
                 torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(),
             ]
@@ -627,9 +665,12 @@ class LTX2Attention(torch.nn.Module):
             num_kv_heads=self.kv_num_heads,
             softmax_scale=1.0 / (dim_head**0.5),
             causal=False,
+            prefix=prefix,
+            disable_kv_quant=disable_kv_quant,
         )
 
         # LTX-2.3: per-head gated attention
+        # leave unquantized for this linear
         if apply_gated_attention:
             self.to_gate_logits = nn.Linear(query_dim, heads, bias=True)
         else:
@@ -748,6 +789,8 @@ class LTX2VideoTransformerBlock(nn.Module):
         elementwise_affine: bool = False,
         rope_type: str = "interleaved",
         perturbed_attn: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.video_cross_attn_adaln = video_cross_attn_adaln
@@ -767,6 +810,8 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn1" if prefix else "attn1",
         )
 
         self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -781,6 +826,8 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.audio_attn1" if prefix else "audio_attn1",
         )
 
         # 2. Prompt Cross-Attention
@@ -796,6 +843,9 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn2" if prefix else "attn2",
+            disable_kv_quant=True,
         )
 
         self.audio_norm2 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -810,9 +860,14 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.audio_attn2" if prefix else "audio_attn2",
+            disable_kv_quant=True,
         )
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
+        # K/V here come from the other modality's latents (long sequences),
+        # so KV-cache quant remains enabled.
         self.audio_to_video_norm = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine)
         self.audio_to_video_attn = LTX2Attention(
             query_dim=dim,
@@ -825,6 +880,8 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.audio_to_video_attn" if prefix else "audio_to_video_attn",
         )
 
         self.video_to_audio_norm = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -839,14 +896,26 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.video_to_audio_attn" if prefix else "video_to_audio_attn",
         )
 
         # 4. Feedforward layers
         self.norm3 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine)
-        self.ff = LTX2FeedForward(dim, activation_fn=activation_fn)
+        self.ff = LTX2FeedForward(
+            dim,
+            activation_fn=activation_fn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ff" if prefix else "ff",
+        )
 
         self.audio_norm3 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
-        self.audio_ff = LTX2FeedForward(audio_dim, activation_fn=activation_fn)
+        self.audio_ff = LTX2FeedForward(
+            audio_dim,
+            activation_fn=activation_fn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.audio_ff" if prefix else "audio_ff",
+        )
 
         # 5. Per-Layer Modulation Parameters
         # LTX-2.3 with cross_attn_adaln uses 9 params (extra 3 for cross-attn modulation);
@@ -1365,6 +1434,9 @@ class LTX2VideoTransformer3DModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     _hsdp_shard_conditions = [is_transformer_block_module]
     _sp_plan: dict[str, Any] | None = None
+    packed_modules_mapping = {
+        "to_qkv": ["to_q", "to_k", "to_v"],
+    }
 
     @staticmethod
     def _build_sp_plan(rope_type: str) -> dict[str, Any]:
@@ -1453,6 +1525,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
         cross_attn_mod: bool = False,
         audio_gated_attn: bool = False,
         audio_cross_attn_mod: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
         self.perturbed_attn = perturbed_attn
@@ -1647,8 +1720,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
                     elementwise_affine=norm_elementwise_affine,
                     rope_type=rope_type,
                     perturbed_attn=perturbed_attn,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{layer_idx}",
                 )
-                for _ in range(num_layers)
+                for layer_idx in range(num_layers)
             ]
         )
 

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -9,10 +9,13 @@ import json
 import os
 from collections.abc import Iterable
 from contextlib import nullcontext
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 import torch
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
 from diffusers.pipelines.ltx2 import LTX2TextConnectors
 from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
@@ -67,14 +70,20 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
     return {}
 
 
-def create_transformer_from_config(config: dict) -> LTX2VideoTransformer3DModel:
+def create_transformer_from_config(
+    config: dict,
+    quant_config: QuantizationConfig | None = None,
+) -> LTX2VideoTransformer3DModel:
     """Create LTX2VideoTransformer3DModel from config dict."""
-    if not config:
+    if not config and quant_config is None:
         return LTX2VideoTransformer3DModel()
 
     signature = inspect.signature(LTX2VideoTransformer3DModel.__init__)
     allowed_keys = set(signature.parameters.keys())
     kwargs = {k: v for k, v in config.items() if k in allowed_keys}
+    if quant_config is not None:
+        kwargs["quant_config"] = quant_config
+
     return LTX2VideoTransformer3DModel(**kwargs)
 
 
@@ -214,7 +223,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         ).to(self.device)
 
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(transformer_config)
+        quant_config = getattr(self.od_config, "quantization_config", None)
+        self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py

@@ -174,7 +174,8 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
 
         # --- Transformer: created empty, weights loaded via AutoWeightsLoader ---
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(transformer_config)
+        quant_config = getattr(self.od_config, "quantization_config", None)
+        self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
 
         # --- Scheduler ---
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(