52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""GPU 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__()
|
|
# Match the production LTX vocoder's conv_pre geometry. The previous
|
|
# single-channel 1x1 probe intermittently had no MIOpen algorithm.
|
|
self.conv = torch.nn.Conv1d(
|
|
128,
|
|
1024,
|
|
kernel_size=7,
|
|
padding=3,
|
|
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", "rocm": "MI325"}, 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, 128, 64), 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
|