From 2bc9cb7f13e2db487c35f05a5973db620043ec9d Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Thu, 21 May 2026 19:52:21 +0800 Subject: [PATCH] [Feat] Support online quantization (fp8/int8) for LTX-2 (#3700) Signed-off-by: Yuanheng Zhao Signed-off-by: yuanheng --- benchmarks/diffusion/quantization_quality.py | 37 +++++-- .../text_to_video/text_to_video.py | 25 +++++ .../quantization/test_quantization_fp8.py | 98 +++++++++++++++++++ .../quantization/test_quantization_quality.py | 40 ++++++-- .../diffusion/models/ltx2/ltx2_transformer.py | 87 ++++++++++++++-- .../diffusion/models/ltx2/pipeline_ltx2.py | 18 +++- .../diffusion/models/ltx2/pipeline_ltx2_3.py | 3 +- 7 files changed, 285 insertions(+), 23 deletions(-) diff --git a/benchmarks/diffusion/quantization_quality.py b/benchmarks/diffusion/quantization_quality.py index 09a9fa74e..b47fae41f 100644 --- a/benchmarks/diffusion/quantization_quality.py +++ b/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(" ", "_") diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 8bbc871ea..23b9afe82 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/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 diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 85466f68e..58f91742d 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/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) ───────────────────────────────────────── diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 02ef0400f..692d83c12 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/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 --- diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index cd8eee9e0..2a65ad6b0 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/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) ] ) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index 3c2322940..c92fd687f 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/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 diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index 84ef29305..ec3b495f5 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/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(