[Edge][WP4] Windowed Code2Wav decoder export (torch.export / ExecuTorch) + client backend; results summarizer
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>
This commit is contained in:
tzhouamandClaude Fable 5.1 committed
1 parent
37c0ea9a6b
commit
580a201631
5 files changed
+549
-1
No files matched your search
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"model": "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
|
||||
"backend": "torch_export",
|
||||
"partitioner": null,
|
||||
"context_frames": 72,
|
||||
"num_quantizers": 16,
|
||||
"hop": 1920,
|
||||
"sample_rate": 24000,
|
||||
"programs": {
|
||||
"2": {
|
||||
"chunk_frames": 2,
|
||||
"window_frames": 74,
|
||||
"max_abs_diff_vs_eager": 0.0,
|
||||
"out_samples": 3840,
|
||||
"path": "/tmp/claude-2005/-data-zhoutaichang-embedding-infer/e695f57b-0163-4287-b1be-1d0ec9330768/scratchpad/code2wav_export/code2wav_c2_ctx72.pt2",
|
||||
"export_s": 8.396002520807087,
|
||||
"bytes": 459787784
|
||||
},
|
||||
"25": {
|
||||
"chunk_frames": 25,
|
||||
"window_frames": 97,
|
||||
"max_abs_diff_vs_eager": 0.0,
|
||||
"out_samples": 48000,
|
||||
"path": "/tmp/claude-2005/-data-zhoutaichang-embedding-infer/e695f57b-0163-4287-b1be-1d0ec9330768/scratchpad/code2wav_export/code2wav_c25_ctx72.pt2",
|
||||
"export_s": 6.752497379668057,
|
||||
"bytes": 459791131
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Render the committed edge results (stream_latency / init_profile / hw_emulation / cpu_sweep /
|
||||
codec_stream) as markdown tables.
|
||||
|
||||
python benchmarks/tts/edge_results/summarize.py [--root benchmarks/tts/edge_results] [--out summary.md]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
GiB = 1 << 30
|
||||
|
||||
|
||||
def _load(p: Path) -> dict[str, Any] | None:
|
||||
try:
|
||||
return json.loads(p.read_text())
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _f(v: Any, nd: int = 1) -> str:
|
||||
if v is None:
|
||||
return ""
|
||||
if isinstance(v, float):
|
||||
return f"{v:.{nd}f}"
|
||||
return str(v)
|
||||
|
||||
|
||||
def stream_latency_rows(root: Path) -> list[dict[str, Any]]:
|
||||
rows = []
|
||||
for f in sorted((root / "stream_latency").glob("*/*/*.json")):
|
||||
d = _load(f)
|
||||
if not d or "summary" not in d:
|
||||
continue
|
||||
s = d["summary"]
|
||||
meta = d.get("meta", {})
|
||||
res = d.get("resources", {}) or {}
|
||||
tag = f.stem.split("_", 1)[1] if "_" in f.stem else ""
|
||||
st = d.get("step_stats") or {}
|
||||
rows.append(
|
||||
{
|
||||
"model": f.parents[1].name.replace("Qwen3-TTS-12Hz-", "").replace("-CustomVoice", ""),
|
||||
"deploy": f.parent.name,
|
||||
"tag": tag,
|
||||
"run": f.stem.split("_", 1)[0],
|
||||
"n": s.get("n_requests"),
|
||||
"init_s": d.get("init", {}).get("init_s"),
|
||||
"ttfa_p50": s["ttfa_ms"]["p50"],
|
||||
"ttfa_p90": s["ttfa_ms"]["p90"],
|
||||
"playback_p50": s["playback_start_ms"]["p50"],
|
||||
"stall_p50": s["stall_at_ttfa_ms"]["p50"],
|
||||
"rtf_stream_p50": s["rtf_streamed"]["p50"],
|
||||
"rtf_total_p50": s["rtf_total"]["p50"],
|
||||
"rss_gib": (res.get("peak_rss_bytes") or 0) / GiB,
|
||||
"gpu_mib": res.get("peak_gpu_mem_mib"),
|
||||
"sha": (meta.get("git_sha") or "")[:8],
|
||||
"step_stats": st.get("attribution", [])[:6] if st else None,
|
||||
"path": str(f.relative_to(root)),
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def table(rows: list[dict[str, Any]], cols: list[str], nd: int = 1) -> str:
|
||||
out = ["| " + " | ".join(cols) + " |", "|" + "---|" * len(cols)]
|
||||
for r in rows:
|
||||
out.append("| " + " | ".join(_f(r.get(c), nd) for c in cols) + " |")
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def render(root: Path) -> str:
|
||||
parts = ["# Edge results summary\n"]
|
||||
rows = stream_latency_rows(root)
|
||||
parts.append("## Streaming latency (stream_latency_bench.py)\n")
|
||||
parts.append(
|
||||
table(
|
||||
rows,
|
||||
[
|
||||
"model",
|
||||
"deploy",
|
||||
"tag",
|
||||
"run",
|
||||
"n",
|
||||
"init_s",
|
||||
"ttfa_p50",
|
||||
"ttfa_p90",
|
||||
"playback_p50",
|
||||
"stall_p50",
|
||||
"rtf_stream_p50",
|
||||
"rtf_total_p50",
|
||||
"rss_gib",
|
||||
"gpu_mib",
|
||||
"sha",
|
||||
],
|
||||
)
|
||||
)
|
||||
for r in rows:
|
||||
if r.get("step_stats"):
|
||||
parts.append(f"\n### Step attribution: {r['model']} {r['deploy']} {r['tag']} ({r['run']})\n")
|
||||
parts.append(
|
||||
table(r["step_stats"], ["role", "counter", "count", "mean_ms", "p95_ms", "share_of_frame"], nd=3)
|
||||
)
|
||||
for name, cols in (
|
||||
("init_profile", ["config", "init_s", "ttfa_ms", "gpu_mib"]),
|
||||
(
|
||||
"hw_emulation",
|
||||
[
|
||||
"case",
|
||||
"hardware_class",
|
||||
"returncode",
|
||||
"init_s",
|
||||
"ttfa_p50",
|
||||
"playback_start_p50",
|
||||
"stall_p50",
|
||||
"rtf_total_p50",
|
||||
"peak_rss_gib",
|
||||
"peak_gpu_mib",
|
||||
],
|
||||
),
|
||||
(
|
||||
"cpu_sweep",
|
||||
[
|
||||
"case",
|
||||
"threads",
|
||||
"inductor",
|
||||
"returncode",
|
||||
"init_s",
|
||||
"ttfa_p50",
|
||||
"playback_start_p50",
|
||||
"rtf_total_p50",
|
||||
"peak_rss_gib",
|
||||
],
|
||||
),
|
||||
):
|
||||
for f in sorted((root / name).glob("*/summary.json")):
|
||||
d = _load(f)
|
||||
if not d:
|
||||
continue
|
||||
parts.append(f"\n## {name} {f.parent.name}\n")
|
||||
rows_n = d.get("rows") or d.get("configs") or []
|
||||
if name == "init_profile" and isinstance(rows_n, dict):
|
||||
rows_n = [{"config": k, **v} for k, v in rows_n.items()]
|
||||
parts.append(table(rows_n, cols))
|
||||
if name == "cpu_sweep":
|
||||
for r in rows_n:
|
||||
if r.get("step_stats"):
|
||||
parts.append(f"\n### Step attribution: {r['case']}\n")
|
||||
parts.append(
|
||||
table(
|
||||
r["step_stats"],
|
||||
["role", "counter", "count", "mean_ms", "p95_ms", "share_of_frame"],
|
||||
nd=3,
|
||||
)
|
||||
)
|
||||
for f in sorted((root / "codec_stream").glob("**/*.json")):
|
||||
d = _load(f)
|
||||
if d and ("snr_db" in d or "parity" in d):
|
||||
parts.append(f"\n## codec_stream {f.relative_to(root)}\n")
|
||||
parts.append("```json\n" + json.dumps(d, indent=1)[:3000] + "\n```")
|
||||
return "\n".join(parts) + "\n"
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--root", default=str(HERE))
|
||||
p.add_argument("--out", default=None)
|
||||
args = p.parse_args(argv)
|
||||
text = render(Path(args.root))
|
||||
if args.out:
|
||||
Path(args.out).write_text(text)
|
||||
else:
|
||||
print(text)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -109,6 +109,63 @@ class LocalDecoder:
|
||||
return wav.float().cpu().numpy().reshape(-1)
|
||||
|
||||
|
||||
class ExecutorchWindowedDecoder:
|
||||
"""Decoder backend for phone/NPU targets: exported windowed programs (see vllm_omni.edge.decoder_export).
|
||||
|
||||
Keeps the last ``context_frames`` decoded frames as explicit state and runs
|
||||
the ``.pte`` program matching the chunk size (falls back to the closest
|
||||
larger program with zero-padded context when a size is missing). Untested
|
||||
in this checkout (executorch not installed); mirrors ``LocalDecoder``'s API.
|
||||
"""
|
||||
|
||||
def __init__(self, manifest_path: str):
|
||||
import json as _json
|
||||
|
||||
from executorch.runtime import Runtime # type: ignore[import-not-found]
|
||||
|
||||
self.manifest = _json.loads(Path(manifest_path).read_text())
|
||||
self.num_quantizers = int(self.manifest["num_quantizers"])
|
||||
self.sample_rate = int(self.manifest["sample_rate"])
|
||||
self.hop = int(self.manifest["hop"])
|
||||
self.context = int(self.manifest["context_frames"])
|
||||
rt = Runtime.get()
|
||||
self.methods = {}
|
||||
for chunk, entry in self.manifest["programs"].items():
|
||||
program = rt.load_program(entry["path"])
|
||||
self.methods[int(chunk)] = program.load_method("forward")
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self.history = np.zeros((0, self.num_quantizers), dtype=np.int64)
|
||||
self.frames_decoded = 0
|
||||
|
||||
def decode(self, codes: np.ndarray) -> np.ndarray:
|
||||
if codes.shape[0] == 0:
|
||||
return np.zeros(0, dtype=np.float32)
|
||||
n = int(codes.shape[0])
|
||||
sizes = sorted(self.methods)
|
||||
size = next((s for s in sizes if s >= n), sizes[-1])
|
||||
window = np.concatenate([self.history, codes], axis=0)[-(self.context + size) :]
|
||||
pad = self.context + size - window.shape[0]
|
||||
if pad > 0:
|
||||
window = np.concatenate([np.zeros((pad, self.num_quantizers), dtype=np.int64), window], axis=0)
|
||||
inp = torch.as_tensor(window.T[None], dtype=torch.long) # [1, Q, ctx+size]
|
||||
out = self.methods[size].execute([inp])[0]
|
||||
wav = out.reshape(-1).float().numpy()[-size * self.hop :]
|
||||
self.history = np.concatenate([self.history, codes], axis=0)[-self.context :]
|
||||
self.frames_decoded += n
|
||||
return wav[-n * self.hop :] if n < size else wav
|
||||
|
||||
|
||||
def build_decoder(args: argparse.Namespace, steady: int):
|
||||
backend = getattr(args, "decoder_backend", "torch")
|
||||
if backend == "executorch":
|
||||
if not getattr(args, "pte_manifest", None):
|
||||
raise SystemExit("--decoder-backend executorch needs --pte-manifest <export dir>/manifest.json")
|
||||
return ExecutorchWindowedDecoder(args.pte_manifest)
|
||||
return LocalDecoder(args.model, device=args.device, left_context=args.left_context, chunk_frames=steady)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- client
|
||||
async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
import websockets
|
||||
@@ -116,7 +173,7 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
ramp = args.ramp or DEFAULT_RAMP
|
||||
schedule = parse_chunk_ramp({"codec_chunk_ramp": ramp}, steady=ramp[-1]) or ramp
|
||||
steady = schedule[-1]
|
||||
decoder = LocalDecoder(args.model, device=args.device, left_context=args.left_context, chunk_frames=steady)
|
||||
decoder = build_decoder(args, steady)
|
||||
url = f"ws://{args.host}:{args.port}/v1/audio/speech/stream"
|
||||
config: dict[str, Any] = {
|
||||
"type": "session.config",
|
||||
@@ -235,6 +292,8 @@ def main(argv: list[str] | None = None) -> int:
|
||||
p.add_argument("--max-new-tokens", dest="max_new_tokens", type=int, default=None)
|
||||
p.add_argument("--seed", type=int, default=None)
|
||||
p.add_argument("--device", default="cpu")
|
||||
p.add_argument("--decoder-backend", dest="decoder_backend", choices=["torch", "executorch"], default="torch")
|
||||
p.add_argument("--pte-manifest", dest="pte_manifest", default=None, help="manifest.json from decoder_export")
|
||||
p.add_argument("--left-context", type=int, default=72)
|
||||
p.add_argument("--ramp", type=int, nargs="*", default=None, help="chunk schedule in frames, default 2 4 8 16 25")
|
||||
p.add_argument("--out", default=None)
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
"""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
|
||||
@@ -0,0 +1,216 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Export the Qwen3-TTS Code2Wav decoder for on-device (NPU / phone) runtimes (WP4).
|
||||
|
||||
vLLM-Omni cannot run on phone NPUs; the ``npu_phone`` hardware class is served
|
||||
by the codec-token stream (WP5) plus an exported decoder artifact. This module
|
||||
exports a *windowed, stateless* variant of ``Qwen3TTSTokenizerV2Decoder``:
|
||||
|
||||
* input ``codes``: ``int64[1, num_quantizers, context_frames + chunk_frames]``
|
||||
(the last ``context_frames`` frames already decoded plus the new chunk);
|
||||
* output ``wav``: ``float32[1, chunk_frames * hop]`` (only the new chunk's audio,
|
||||
``hop = total_upsample`` samples per frame, 1920 at 24 kHz / 12.5 Hz).
|
||||
|
||||
Fixed shapes (one program per chunk size of the ramp schedule) keep the export
|
||||
static, which is what ExecuTorch / QNN / CoreML partitioners need; the
|
||||
72-frame window is explicit input state instead of the decoder's Python
|
||||
caches, so the artifact carries no mutable state. Backends:
|
||||
|
||||
* ``torch_export`` - ``torch.export`` + ``.pt2`` (always available);
|
||||
* ``executorch`` - XNNPACK-lowered ``.pte`` (needs ``pip install executorch``;
|
||||
the QNN partitioner is export-only and listed as unverified).
|
||||
|
||||
Status: the ExecuTorch path is written against the executorch 1.x API but was
|
||||
not executed in this checkout (package not installed; install is permission
|
||||
gated). ``tests/edge/test_decoder_export.py`` skips the lowering without it and
|
||||
still checks the windowed wrapper against the eager decoder.
|
||||
|
||||
CLI:
|
||||
python -m vllm_omni.edge.decoder_export --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
|
||||
--ramp 2 4 8 16 25 --context 72 --backend torch_export --out-dir ./code2wav_export
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
DEFAULT_RAMP = (2, 4, 8, 16, 25)
|
||||
DEFAULT_CONTEXT = 72
|
||||
|
||||
|
||||
class WindowedCode2Wav(nn.Module):
|
||||
"""Stateless windowed decode: full window in, trailing ``chunk_frames`` of audio out."""
|
||||
|
||||
def __init__(self, decoder: nn.Module, chunk_frames: int, context_frames: int = DEFAULT_CONTEXT):
|
||||
super().__init__()
|
||||
self.decoder = decoder
|
||||
self.chunk_frames = int(chunk_frames)
|
||||
self.context_frames = int(context_frames)
|
||||
self.hop = int(getattr(decoder, "total_upsample", 1920))
|
||||
self.num_quantizers = int(getattr(getattr(decoder, "config", None), "num_quantizers", 16))
|
||||
|
||||
@property
|
||||
def window_frames(self) -> int:
|
||||
return self.context_frames + self.chunk_frames
|
||||
|
||||
def example_input(self) -> torch.Tensor:
|
||||
return torch.zeros(1, self.num_quantizers, self.window_frames, dtype=torch.long)
|
||||
|
||||
def forward(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
wav = self.decoder._forward_exact(codes) # [1, 1, samples] or [1, samples]
|
||||
wav = wav.reshape(wav.shape[0], -1)
|
||||
return wav[:, -self.chunk_frames * self.hop :]
|
||||
|
||||
|
||||
def resolve_model_dir(model: str) -> Path:
|
||||
p = Path(model)
|
||||
if p.is_dir():
|
||||
return p
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
return Path(snapshot_download(model))
|
||||
|
||||
|
||||
def load_code2wav_decoder(model: str, dtype: torch.dtype = torch.float32) -> nn.Module:
|
||||
"""Load ``Qwen3TTSTokenizerV2Decoder`` from ``<model>/speech_tokenizer`` (eager, CPU)."""
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
|
||||
Qwen3TTSTokenizerV2Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
|
||||
Qwen3TTSTokenizerV2Decoder,
|
||||
)
|
||||
|
||||
model_dir = resolve_model_dir(model)
|
||||
cfg = Qwen3TTSTokenizerV2Config.from_pretrained(str(model_dir), subfolder="speech_tokenizer")
|
||||
decoder = Qwen3TTSTokenizerV2Decoder._from_config(cfg.decoder_config)
|
||||
state = load_file(str(model_dir / "speech_tokenizer" / "model.safetensors"))
|
||||
dec_state = {k[len("decoder.") :]: v for k, v in state.items() if k.startswith("decoder.")}
|
||||
_, unexpected = decoder.load_state_dict(dec_state, strict=False)
|
||||
if unexpected:
|
||||
raise RuntimeError(f"unexpected decoder keys: {unexpected[:5]}")
|
||||
decoder.eval().to(dtype=dtype)
|
||||
if hasattr(decoder, "precompute_snake_caches"):
|
||||
decoder.precompute_snake_caches()
|
||||
return decoder
|
||||
|
||||
|
||||
def export_program(module: WindowedCode2Wav) -> Any:
|
||||
"""``torch.export`` the windowed module with its fixed example input."""
|
||||
example = module.example_input()
|
||||
with torch.inference_mode():
|
||||
return torch.export.export(module.eval(), (example,))
|
||||
|
||||
|
||||
def lower_executorch(exported_program: Any, partitioner: str = "xnnpack") -> bytes:
|
||||
"""Lower an exported program to an ExecuTorch ``.pte`` buffer (requires executorch)."""
|
||||
from executorch.exir import to_edge_transform_and_lower # type: ignore[import-not-found]
|
||||
|
||||
partitioners = []
|
||||
if partitioner == "xnnpack":
|
||||
from executorch.backends.xnnpack.partition.xnnpack_partitioner import ( # type: ignore[import-not-found]
|
||||
XnnpackPartitioner,
|
||||
)
|
||||
|
||||
partitioners = [XnnpackPartitioner()]
|
||||
elif partitioner == "qnn": # unverified: needs the Qualcomm SDK build of executorch
|
||||
from executorch.backends.qualcomm.partition.qnn_partitioner import ( # type: ignore[import-not-found]
|
||||
QnnPartitioner,
|
||||
)
|
||||
|
||||
partitioners = [QnnPartitioner()]
|
||||
elif partitioner not in ("none", ""):
|
||||
raise ValueError(f"unknown partitioner {partitioner!r}")
|
||||
edge = to_edge_transform_and_lower(exported_program, partitioner=partitioners or None)
|
||||
return edge.to_executorch().buffer
|
||||
|
||||
|
||||
def export_windowed_decoder(
|
||||
model: str,
|
||||
out_dir: str | Path,
|
||||
*,
|
||||
ramp: tuple[int, ...] = DEFAULT_RAMP,
|
||||
context_frames: int = DEFAULT_CONTEXT,
|
||||
backend: str = "torch_export",
|
||||
partitioner: str = "xnnpack",
|
||||
check: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
"""Export one program per chunk size; returns a manifest (also written as ``manifest.json``)."""
|
||||
out = Path(out_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
decoder = load_code2wav_decoder(model)
|
||||
manifest: dict[str, Any] = {
|
||||
"model": model,
|
||||
"backend": backend,
|
||||
"partitioner": partitioner if backend == "executorch" else None,
|
||||
"context_frames": context_frames,
|
||||
"num_quantizers": int(decoder.config.num_quantizers),
|
||||
"hop": int(decoder.total_upsample),
|
||||
"sample_rate": int(getattr(decoder.config, "output_sample_rate", 24000) or 24000),
|
||||
"programs": {},
|
||||
}
|
||||
for chunk in sorted(set(int(c) for c in ramp)):
|
||||
module = WindowedCode2Wav(decoder, chunk, context_frames)
|
||||
t0 = time.perf_counter()
|
||||
ep = export_program(module)
|
||||
entry: dict[str, Any] = {"chunk_frames": chunk, "window_frames": module.window_frames}
|
||||
if check:
|
||||
codes = torch.randint(0, int(decoder.config.codebook_size), module.example_input().shape, dtype=torch.long)
|
||||
with torch.inference_mode():
|
||||
ref = module(codes)
|
||||
got = ep.module()(codes)
|
||||
entry["max_abs_diff_vs_eager"] = float((ref - got).abs().max())
|
||||
entry["out_samples"] = int(got.shape[-1])
|
||||
if backend == "torch_export":
|
||||
path = out / f"code2wav_c{chunk}_ctx{context_frames}.pt2"
|
||||
torch.export.save(ep, str(path))
|
||||
elif backend == "executorch":
|
||||
path = out / f"code2wav_c{chunk}_ctx{context_frames}.pte"
|
||||
path.write_bytes(lower_executorch(ep, partitioner))
|
||||
else:
|
||||
raise ValueError(f"unknown backend {backend!r}")
|
||||
entry["path"] = str(path)
|
||||
entry["export_s"] = time.perf_counter() - t0
|
||||
entry["bytes"] = path.stat().st_size
|
||||
manifest["programs"][str(chunk)] = entry
|
||||
print(
|
||||
f"[decoder_export] chunk={chunk} -> {path} ({entry['bytes'] / 1e6:.1f} MB, {entry['export_s']:.1f} s)",
|
||||
flush=True,
|
||||
)
|
||||
(out / "manifest.json").write_text(json.dumps(manifest, indent=1))
|
||||
return manifest
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--model", required=True)
|
||||
p.add_argument("--out-dir", required=True)
|
||||
p.add_argument("--ramp", type=int, nargs="*", default=list(DEFAULT_RAMP))
|
||||
p.add_argument("--context", type=int, default=DEFAULT_CONTEXT)
|
||||
p.add_argument("--backend", choices=["torch_export", "executorch"], default="torch_export")
|
||||
p.add_argument("--partitioner", choices=["xnnpack", "qnn", "none"], default="xnnpack")
|
||||
p.add_argument("--no-check", action="store_true")
|
||||
args = p.parse_args(argv)
|
||||
export_windowed_decoder(
|
||||
args.model,
|
||||
args.out_dir,
|
||||
ramp=tuple(args.ramp),
|
||||
context_frames=args.context,
|
||||
backend=args.backend,
|
||||
partitioner=args.partitioner,
|
||||
check=not args.no_check,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in new issue
Block a user