From d3c990dca8c3787ffdfe1d1424cfda9584d6bff9 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Wed, 26 Aug 2026 20:59:10 +0800 Subject: [PATCH] [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 | 25 +++ .../models/ltx2/test_ltx2_vocoder_cuda.py | 48 ++++++ .../ltx/test_ltx25_official_similarity.py | 79 ++++----- .../ltx/test_ltx_official_similarity.py | 154 ++++++++---------- .../diffusion/models/ltx2/ltx2_components.py | 20 ++- .../diffusion/models/ltx2/ltx2_runtime.py | 25 ++- 6 files changed, 210 insertions(+), 141 deletions(-) create mode 100644 tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 108bb960d..2e0647ccf 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/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"), diff --git a/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py new file mode 100644 index 000000000..d3d885d57 --- /dev/null +++ b/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 diff --git a/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py index d72d14022..d7c82ef3a 100644 --- a/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py +++ b/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) diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py index 59a8540aa..63ce73923 100644 --- a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py +++ b/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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 37c3f05dc..d4d98e7c4 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index f560f678b..d3addd432 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/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(