[Model][Bugfix] Improve LTX audio parity and similarity guards (#6342)

Signed-off-by: mglyn <1203789601@qq.com>

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -11,6 +11,31 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
+    import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+
+    class FakeBaseVocoder(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.input_dtype = None
+
+        def forward(self, value):
+            self.input_dtype = value.dtype
+            return value
+
+    monkeypatch.setattr(
+        ltx_runtime.torch,
+        "autocast",
+        lambda **_kwargs: (_ for _ in ()).throw(AssertionError("base vocoder must not use autocast")),
+    )
+    vocoder = FakeBaseVocoder()
+
+    output = ltx_runtime._run_ltx_vocoder(vocoder, torch.ones(1, dtype=torch.bfloat16))
+
+    assert vocoder.input_dtype == torch.bfloat16
+    assert output.dtype == torch.bfloat16
+
+
 class TestLTXOutputRank:
     @pytest.mark.parametrize(
         ("distributed_vae_state", "expected_decode_calls"),

tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CUDA regression tests for LTX-2 BWE vocoder precision."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+class _BWEConvVocoder(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv1d(
+            1,
+            1,
+            kernel_size=1,
+            bias=False,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        self.bwe_generator = torch.nn.Identity()
+        self.input_dtype = None
+        self.conv_output_dtype = None
+
+    def forward(self, value):
+        self.input_dtype = value.dtype
+        output = self.conv(value)
+        self.conv_output_dtype = output.dtype
+        return output
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32():
+    from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder
+
+    vocoder = _BWEConvVocoder()
+    generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16)
+
+    output = _run_ltx_vocoder(vocoder, generated_mel)
+
+    assert next(vocoder.parameters()).dtype == torch.bfloat16
+    assert vocoder.input_dtype == torch.float32
+    assert vocoder.conv_output_dtype == torch.float32
+    assert output.dtype == torch.bfloat16

tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py

@@ -22,6 +22,7 @@ import shutil
 import subprocess
 import sys
 import time
+from dataclasses import asdict, dataclass
 from pathlib import Path
 
 import numpy as np
@@ -83,35 +84,27 @@ NEGATIVE_PROMPT = (
 # Release parity exercises the accelerated BF16 cuDNN path in both runtimes.
 # Degenerate one-token attention falls back to PyTorch SDPA MATH.
 ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16"
-VIDEO_SSIM_MEAN_THRESHOLD = 0.95
-VIDEO_SSIM_MIN_THRESHOLD = 0.90
-VIDEO_PSNR_MEAN_THRESHOLD = 30.0
-AUDIO_RELATIVE_L2_THRESHOLD = 0.2
-AUDIO_COSINE_THRESHOLD = 0.95
-
-# Decoded-output gates for the pinned official/Omni BF16 cuDNN comparison.
-# The metrics artifact preserves the exact observed values.
-LTX25_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
-LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
-LTX25_VIDEO_SSIM_MIN_THRESHOLD = 0.99
-LTX25_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
-LTX25_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
-LTX25_AUDIO_RELATIVE_L2_THRESHOLD = 0.32
-LTX25_AUDIO_COSINE_THRESHOLD = 0.95
-
-# Full/SFT uses independently converted dev weights, but the release contract
-# still requires decoded-video parity at the same 0.99 SSIM floor as the
-# distilled paths when both runtimes use BF16 cuDNN attention.
-LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
-LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD = 0.99
-LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
-LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD = 0.3
-LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
-LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD = 0.02
-LTX25_FULL_AUDIO_COSINE_THRESHOLD = AUDIO_COSINE_THRESHOLD
 LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0]
 
 
+@dataclass(frozen=True)
+class LTXAccuracyThresholds:
+    video_ssim_mean: float
+    video_ssim_min: float
+    video_psnr_mean_db: float
+    audio_relative_l2: float
+    audio_cosine_similarity: float
+
+
+STRICT_THRESHOLDS = LTXAccuracyThresholds(
+    video_ssim_mean=0.99,
+    video_ssim_min=0.99,
+    video_psnr_mean_db=40.0,
+    audio_relative_l2=0.10,
+    audio_cosine_similarity=0.99,
+)
+
+
 def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None:
     from vllm_omni.outputs import OmniRequestOutput
 
@@ -419,8 +412,6 @@ def test_ltx25_i2v_requests_pin_official_crf() -> None:
 
 
 def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
-    from benchmarks.diffusion.quantization_quality import compute_lpips_video
-
     assert reference.shape == prediction.shape
     assert reference.ndim == 4 and reference.shape[-1] == 3
     ssim_scores: list[float] = []
@@ -435,7 +426,6 @@ def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
         "ssim_mean": float(np.mean(ssim_scores)),
         "ssim_min": float(np.min(ssim_scores)),
         "psnr_mean_db": float(np.mean(psnr_scores)),
-        "lpips_alex_mean": compute_lpips_video(reference, prediction),
         "max_abs": float(difference.max()),
         "mean_abs": float(difference.mean()),
     }
@@ -463,6 +453,17 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
     }
 
 
+def _assert_strict_similarity(
+    video_metrics: dict[str, float],
+    audio_metrics: dict[str, float | bool],
+) -> None:
+    assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
+    assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
+    assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
+    assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
+    assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity
+
+
 @pytest.mark.slow
 @pytest.mark.benchmark
 @pytest.mark.diffusion
@@ -569,6 +570,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
         "omni_model": omni_model_source,
         "omni_model_revision": omni_model_revision,
         "resolved_omni_model_path": str(omni_model),
+        "thresholds": asdict(STRICT_THRESHOLDS),
         "peak_memory_mb": {
             "official_allocated": official_metadata["peak_memory_allocated_mb"],
             "official_reserved": official_metadata["peak_memory_reserved_mb"],
@@ -582,13 +584,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
     (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
     print(json.dumps(result, indent=2))
 
-    ssim_mean_threshold = LTX25_VIDEO_SSIM_MEAN_THRESHOLD if task == "t2v" else LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD
-    assert video_metrics["ssim_mean"] >= ssim_mean_threshold
-    assert video_metrics["ssim_min"] >= LTX25_VIDEO_SSIM_MIN_THRESHOLD
-    assert video_metrics["psnr_mean_db"] >= LTX25_VIDEO_PSNR_MEAN_THRESHOLD
-    assert video_metrics["lpips_alex_mean"] <= LTX25_VIDEO_LPIPS_MEAN_THRESHOLD
-    assert audio_metrics["cosine_similarity"] >= LTX25_AUDIO_COSINE_THRESHOLD
-    assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
+    _assert_strict_similarity(video_metrics, audio_metrics)
 
 
 @pytest.mark.slow
@@ -729,6 +725,7 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
         "omni_model": omni_model_source,
         "omni_model_revision": omni_model_revision,
         "resolved_omni_model_path": str(omni_model),
+        "thresholds": asdict(STRICT_THRESHOLDS),
         "peak_memory_mb": {
             "official_allocated": official_metadata["peak_memory_allocated_mb"],
             "official_reserved": official_metadata["peak_memory_reserved_mb"],
@@ -742,12 +739,4 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
     (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
     print(json.dumps(result, indent=2))
 
-    assert video_metrics["ssim_mean"] >= LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD
-    assert video_metrics["ssim_min"] >= LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD
-    assert video_metrics["psnr_mean_db"] >= LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD
-    lpips_threshold = (
-        LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD if task == "t2v" else LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD
-    )
-    assert audio_metrics["relative_l2"] <= LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD
-    assert video_metrics["lpips_alex_mean"] <= lpips_threshold
-    assert audio_metrics["cosine_similarity"] >= LTX25_FULL_AUDIO_COSINE_THRESHOLD
+    _assert_strict_similarity(video_metrics, audio_metrics)

tests/e2e/accuracy/ltx/test_ltx_official_similarity.py

@@ -3,9 +3,8 @@
 
 """E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
 
-The original reduced one-stage guards remain unchanged. Four complementary
-default-shape cases cover both model versions, both two-stage weight families,
-and T2V/I2V without expanding every combination into a duplicate golden.
+Four complementary default-shape cases cover both model versions, both
+two-stage weight families, and T2V/I2V without duplicating one-stage coverage.
 """
 
 from __future__ import annotations
@@ -69,25 +68,23 @@ class LTXAccuracyThresholds:
     video_ssim_mean: float
     video_ssim_min: float
     video_psnr_mean_db: float
-    # Zero disables gating for waveform-sensitive audio that is not expected
-    # to align with the official output.
     audio_relative_l2: float
     audio_cosine_similarity: float
 
 
 STRICT_THRESHOLDS = LTXAccuracyThresholds(
-    video_ssim_mean=0.95,
-    video_ssim_min=0.90,
-    video_psnr_mean_db=30.0,
-    audio_relative_l2=0.20,
-    audio_cosine_similarity=0.95,
+    video_ssim_mean=0.99,
+    video_ssim_min=0.99,
+    video_psnr_mean_db=40.0,
+    audio_relative_l2=0.10,
+    audio_cosine_similarity=0.99,
 )
 
 
 @dataclass(frozen=True)
 class LTXAccuracyCase:
     name: str
-    pipeline_kind: Literal["one_stage", "distilled", "two_stage"]
+    pipeline_kind: Literal["distilled", "two_stage"]
     model_id: str
     model_revision: str
     model_env: str
@@ -99,8 +96,6 @@ class LTXAccuracyCase:
     num_inference_steps: int
     seed: int
     stg_block: int | None
-    thresholds: LTXAccuracyThresholds
-    prompt: str = PROMPT
     gemma_model_id: str | None = None
     gemma_model_revision: str | None = None
     gemma_model_env: str | None = None
@@ -168,59 +163,6 @@ I2V_IMAGE = LTXArtifact(
 )
 
 
-LEGACY_CASES = (
-    LTXAccuracyCase(
-        name="ltx2",
-        pipeline_kind="one_stage",
-        model_id="Lightricks/LTX-2",
-        model_revision=LTX2_REVISION,
-        model_env="VLLM_TEST_LTX2_MODEL",
-        model_class_name="LTX2Pipeline",
-        checkpoint=LTX2_CHECKPOINT,
-        width=512,
-        height=384,
-        num_frames=25,
-        num_inference_steps=20,
-        seed=42,
-        stg_block=29,
-        thresholds=STRICT_THRESHOLDS,
-    ),
-    LTXAccuracyCase(
-        name="ltx2_3",
-        pipeline_kind="one_stage",
-        model_id="diffusers/LTX-2.3-Diffusers",
-        model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
-        model_env="VLLM_TEST_LTX23_MODEL",
-        model_class_name="LTX2Pipeline",
-        checkpoint=LTX23_CHECKPOINT,
-        width=512,
-        height=384,
-        num_frames=25,
-        num_inference_steps=20,
-        seed=42,
-        stg_block=28,
-        thresholds=STRICT_THRESHOLDS,
-    ),
-    LTXAccuracyCase(
-        name="ltx2_3_i2v",
-        pipeline_kind="one_stage",
-        model_id="diffusers/LTX-2.3-Diffusers",
-        model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
-        model_env="VLLM_TEST_LTX23_MODEL",
-        model_class_name="LTX2Pipeline",
-        checkpoint=LTX23_CHECKPOINT,
-        width=512,
-        height=384,
-        num_frames=25,
-        num_inference_steps=20,
-        seed=42,
-        stg_block=28,
-        thresholds=STRICT_THRESHOLDS,
-        image=I2V_IMAGE,
-    ),
-)
-
-
 TWO_STAGE_CASES = (
     LTXAccuracyCase(
         name="ltx2_distilled_t2v",
@@ -240,7 +182,6 @@ TWO_STAGE_CASES = (
         gemma_model_id="Lightricks/LTX-2",
         gemma_model_revision=LTX2_REVISION,
         gemma_model_env="VLLM_TEST_LTX2_MODEL",
-        thresholds=STRICT_THRESHOLDS,
     ),
     LTXAccuracyCase(
         name="ltx23_distilled_i2v",
@@ -260,7 +201,6 @@ TWO_STAGE_CASES = (
         gemma_model_id="diffusers/LTX-2.3-Diffusers",
         gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
         gemma_model_env="VLLM_TEST_LTX23_MODEL",
-        thresholds=STRICT_THRESHOLDS,
         image=I2V_IMAGE,
     ),
     LTXAccuracyCase(
@@ -280,7 +220,6 @@ TWO_STAGE_CASES = (
         seed=10,
         stg_block=29,
         enable_layerwise_offload=True,
-        thresholds=STRICT_THRESHOLDS,
     ),
     LTXAccuracyCase(
         name="ltx23_two_stage_layer_fused_i2v",
@@ -299,12 +238,11 @@ TWO_STAGE_CASES = (
         seed=10,
         stg_block=28,
         enable_layerwise_offload=True,
-        thresholds=STRICT_THRESHOLDS,
         image=I2V_IMAGE,
     ),
 )
 
-CASES = (*LEGACY_CASES, *TWO_STAGE_CASES)
+WEEKLY_CASES = TWO_STAGE_CASES
 
 
 def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
@@ -460,10 +398,49 @@ def _resolve_image(case: LTXAccuracyCase) -> Path | None:
     return _resolve_artifact(case.image)
 
 
+def _omni_model_with_pinned_artifacts(
+    model: Path,
+    output_root: Path,
+    artifacts: tuple[Path | None, ...],
+) -> Path:
+    """Expose the already-resolved sidecars to Omni without copying the model."""
+    resolved_artifacts = tuple(artifact for artifact in artifacts if artifact is not None)
+    if not resolved_artifacts:
+        return model
+
+    overlay = output_root / "omni-model"
+    overlay.mkdir()
+    for source in model.iterdir():
+        (overlay / source.name).symlink_to(source.resolve(), target_is_directory=source.is_dir())
+    for artifact in resolved_artifacts:
+        target = overlay / artifact.name
+        if target.exists():
+            assert target.samefile(artifact), f"Model sidecar does not match pinned artifact: {target} != {artifact}"
+        else:
+            target.symlink_to(artifact.resolve())
+    return overlay
+
+
+def test_omni_model_overlay_pins_artifacts(tmp_path: Path) -> None:
+    model = tmp_path / "model"
+    model.mkdir()
+    component = model / "model_index.json"
+    component.write_text("{}\n")
+    artifact = tmp_path / "sidecar.safetensors"
+    artifact.write_bytes(b"pinned")
+    output_root = tmp_path / "output"
+    output_root.mkdir()
+
+    overlay = _omni_model_with_pinned_artifacts(model, output_root, (artifact,))
+
+    assert (overlay / component.name).samefile(component)
+    assert (overlay / artifact.name).samefile(artifact)
+
+
 def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
     request: dict[str, object] = {
         "pipeline_kind": case.pipeline_kind,
-        "prompt": case.prompt,
+        "prompt": PROMPT,
         "negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
         "width": case.width,
         "height": case.height,
@@ -473,14 +450,17 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
         "seed": case.seed,
     }
     if case.pipeline_kind != "distilled":
-        assert case.stg_block is not None
         request.update(
             video_cfg_scale=3.0,
             audio_cfg_scale=7.0,
-            video_stg_scale=1.0,
-            audio_stg_scale=1.0,
             video_modality_scale=3.0,
             audio_modality_scale=3.0,
+        )
+    if case.pipeline_kind == "two_stage":
+        assert case.stg_block is not None
+        request.update(
+            video_stg_scale=1.0,
+            audio_stg_scale=1.0,
             video_rescale_scale=0.7,
             audio_rescale_scale=0.7,
             video_stg_blocks=[case.stg_block],
@@ -543,7 +523,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
     }
 
 
-@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
+@pytest.mark.parametrize("case", WEEKLY_CASES, ids=lambda case: case.name)
 @pytest.mark.slow
 @pytest.mark.benchmark
 @pytest.mark.diffusion
@@ -561,9 +541,11 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
     checkpoint = _resolve_artifact(case.checkpoint, model)
     spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
     distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
+    omni_model = _omni_model_with_pinned_artifacts(model, output_root, (spatial_upsampler, distilled_lora))
     image = _resolve_image(case)
+    request = _request(case, image)
     request_path = output_root / "request.json"
-    request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
+    request_path.write_text(json.dumps(request, indent=2) + "\n")
 
     runner = Path(__file__).with_name("run_ltx_reference.py")
     runner_args = [
@@ -619,7 +601,7 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
             "--output-dir",
             str(omni_output),
             "--model",
-            str(model),
+            str(omni_model),
             "--model-class-name",
             case.model_class_name,
         ]
@@ -651,18 +633,16 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
         "spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
         "distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
         "enable_layerwise_offload": enable_layerwise_offload,
-        "thresholds": asdict(case.thresholds),
-        "request": _request(case, image),
+        "thresholds": asdict(STRICT_THRESHOLDS),
+        "request": request,
         "video": video_metrics,
         "audio": audio_metrics,
     }
     (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
     print(json.dumps(result, indent=2))
 
-    assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean
-    assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
-    assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
-    if case.thresholds.audio_relative_l2 > 0:
-        assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
-    if case.thresholds.audio_cosine_similarity > 0:
-        assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity
+    assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
+    assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
+    assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
+    assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
+    assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -637,14 +637,18 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         upsampler_config = os.path.join(model, "latent_upsampler", "config.json")
         if os.path.isfile(upsampler_config) or not local_files_only:
             try:
-                pipeline.latent_upsampler = _load_component(
-                    LTX2LatentUpsamplerModel,
-                    model,
-                    "latent_upsampler",
-                    local_files_only=local_files_only,
-                    dtype=dtype,
-                    revision=revision,
-                )
+                # Diffusers builds the rational-resampler kernel with Long
+                # tensors; constructing it under a CUDA default-device context
+                # attempts an unsupported Long addmm before weights are loaded.
+                with torch.device("cpu"):
+                    pipeline.latent_upsampler = _load_component(
+                        LTX2LatentUpsamplerModel,
+                        model,
+                        "latent_upsampler",
+                        local_files_only=local_files_only,
+                        dtype=dtype,
+                        revision=revision,
+                    )
             except (OSError, ValueError):
                 if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None:
                     raise

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -60,6 +60,29 @@ from .ltx2_request import (
 )
 
 
+def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.Tensor:
+    """Run the BWE vocoder in FP32, matching the official LTX pipeline."""
+    if not hasattr(vocoder, "bwe_generator"):
+        return vocoder(generated_mel)
+
+    input_dtype = generated_mel.dtype
+    device_type = generated_mel.device.type
+    module_dtype = next(vocoder.parameters()).dtype
+    if device_type == "mps":
+        if module_dtype != torch.float32:
+            vocoder.float()
+        try:
+            return vocoder(generated_mel.float()).to(input_dtype)
+        finally:
+            if module_dtype != torch.float32:
+                vocoder.to(module_dtype)
+
+    # BF16 errors compound through the BWE model's long convolution stack.
+    # FP32 autocast upcasts each op without materializing a second FP32 model.
+    with torch.autocast(device_type=device_type, dtype=torch.float32):
+        return vocoder(generated_mel.float()).to(input_dtype)
+
+
 def _expand_per_prompt_decode_value(
     value: float | list[float],
     *,
@@ -558,7 +581,7 @@ class LTXRuntime(
         if video.numel() > 0:
             video = self.video_processor.postprocess_video(video, output_type=output_type)
         generated_mel = self.audio_vae.decode(audio_latents.to(self.audio_vae.dtype), return_dict=False)[0]
-        audio = self.vocoder(generated_mel)
+        audio = _run_ltx_vocoder(self.vocoder, generated_mel)
         return self._make_output((video, audio))
 
     def _prepare_video_latents_stage(