tzhouamandClaude Fable 5.1
580a201631
vllm_omni/edge/decoder_export.py exports a stateless windowed decoder (context 72 frames + chunk, one program per ramp chunk size) via torch.export (.pt2, verified here on CPU for chunks 2 and 25 with the real 0.6B decoder: exact match vs eager, 460 MB fp32 per program) or ExecuTorch XNNPACK (.pte; written against the executorch 1.x API, not executed: package not installed, install is permission gated; QNN partitioner export-only/unverified). The WP5 client gains --decoder-backend executorch --pte-manifest for phone/NPU targets. benchmarks/tts/edge_results/summarize.py renders all committed JSON as tables. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
63 lines
2.2 KiB
Python
63 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
|