62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
|
"""WP4 decoder export: windowed wrapper (CPU, fake decoder) and ExecuTorch lowering (skipped without executorch)."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import importlib.util
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
|
||
|
|
from vllm_omni.edge import decoder_export as dx
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||
|
|
|
||
|
|
HOP = 4
|
||
|
|
|
||
|
|
|
||
|
|
class _FakeConfig:
|
||
|
|
num_quantizers = 3
|
||
|
|
codebook_size = 8
|
||
|
|
output_sample_rate = 24000
|
||
|
|
|
||
|
|
|
||
|
|
class _FakeDecoder(nn.Module):
|
||
|
|
"""Stands in for Qwen3TTSTokenizerV2Decoder: ``_forward_exact`` maps each frame to HOP samples."""
|
||
|
|
|
||
|
|
total_upsample = HOP
|
||
|
|
config = _FakeConfig()
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.scale = nn.Parameter(torch.tensor(0.5))
|
||
|
|
|
||
|
|
def _forward_exact(self, codes: torch.Tensor) -> torch.Tensor:
|
||
|
|
# [1, Q, T] -> [1, 1, T*HOP]; value depends on the frame's codes so slicing is testable
|
||
|
|
per_frame = codes.to(torch.float32).sum(dim=1, keepdim=True) * self.scale # [1, 1, T]
|
||
|
|
return per_frame.repeat_interleave(HOP, dim=-1).clamp(min=-1000, max=1000)
|
||
|
|
|
||
|
|
|
||
|
|
def test_windowed_wrapper_returns_only_the_new_chunk():
|
||
|
|
module = dx.WindowedCode2Wav(_FakeDecoder(), chunk_frames=2, context_frames=3)
|
||
|
|
assert module.window_frames == 5 and module.hop == HOP and module.num_quantizers == 3
|
||
|
|
codes = torch.arange(15, dtype=torch.long).reshape(1, 3, 5)
|
||
|
|
out = module(codes)
|
||
|
|
assert out.shape == (1, 2 * HOP)
|
||
|
|
full = _FakeDecoder()._forward_exact(codes).reshape(1, -1)
|
||
|
|
assert torch.equal(out, full[:, -2 * HOP :])
|
||
|
|
|
||
|
|
|
||
|
|
def test_torch_export_of_windowed_wrapper_matches_eager():
|
||
|
|
module = dx.WindowedCode2Wav(_FakeDecoder(), chunk_frames=2, context_frames=3)
|
||
|
|
ep = dx.export_program(module)
|
||
|
|
codes = torch.randint(0, 8, module.example_input().shape, dtype=torch.long)
|
||
|
|
with torch.inference_mode():
|
||
|
|
assert torch.allclose(ep.module()(codes), module(codes))
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.skipif(importlib.util.find_spec("executorch") is None, reason="executorch not installed")
|
||
|
|
def test_executorch_lowering_of_windowed_wrapper():
|
||
|
|
module = dx.WindowedCode2Wav(_FakeDecoder(), chunk_frames=2, context_frames=3)
|
||
|
|
buf = dx.lower_executorch(dx.export_program(module), partitioner="xnnpack")
|
||
|
|
assert isinstance(buf, (bytes, bytearray)) and len(buf) > 0
|