fb103ac5ae
Every registered string now lives in one table (codecs/registry.rs) with its kind, status, media class, registering SPEC section, edition and registration date. `wai registry export` writes registry/capabilities.json, its JSON Schema, the generated table in REGISTRY.md and vocab v2; --check fails on drift. vocab v1 stays stable. A candidate capability (a registered string whose payload format no revision has pinned) must not be emitted. The C packers, the browser writers and `wai wrap` refuse one, and `codecs::is_dispatchable` is the single dispatch and selection predicate (SPEC §4): candidates, companion-only capabilities, records and derivations are not dispatchable, and an unregistered string is decoded only where the sink advertises it. The codec glossary is linked to the registry through a reviewed table and held to its schema by a library test; `org` becomes `body`, from a closed allowlist. The SPEC gains §5 sections on kinds and candidates, a §8 edition rule, corrected citations and rewritten appendices. The quantum-ops vectors are regenerated with neutral device ids. CI declares the feature-gated test targets with required-features, checks the registry export, builds wai-web for wasm32, and runs two public-text checks (comparative prose, and private name patterns) with a self-test. The specification, extensions, research notes, READMEs, glossary and source comments are rewritten to name no company or product and to make no comparisons. Co-Authored-By: Claude Opus 5.5 <[email protected]>
113 lines
4.3 KiB
Python
113 lines
4.3 KiB
Python
"""WAI source-side neural encoder for `wai.neural.wavtokenizer`.
|
||
|
||
WavTokenizer-large-unify-40token (Jiang et al., ICLR 2025) — single
|
||
quantizer, 40 tokens per second of 24 kHz audio, single 4096-entry
|
||
codebook. UTMOS at 0.9 kbps reportedly exceeds DAC at 9 kbps.
|
||
|
||
WavTokenizer is research code, not a transformers integration. We add
|
||
the cloned repo to sys.path and use its `decoder.pretrained.WavTokenizer`
|
||
class directly. Set --wavtokenizer-src to the cloned repo path (default
|
||
/tmp/wavtokenizer from `git clone`).
|
||
|
||
Wire format (uniform with all wai.neural.<audio> capabilities):
|
||
<III> q t n_samples (little-endian u32 × 3)
|
||
q*t × u16 codes (row-major, LE)
|
||
|
||
For WavTokenizer-40token q=1 (single quantizer). The codes are u16 even
|
||
though the codebook fits in 12 bits — keeping the WAI wire format
|
||
uniform across audio codecs has more value than the 4-bit-per-code save.
|
||
|
||
Usage:
|
||
python tools/wai_wavtokenizer_encode.py <input.wav> \\
|
||
--wavtokenizer-src /tmp/wavtokenizer \\
|
||
--config /tmp/wavtokenizer/configs/wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn.yaml \\
|
||
--ckpt /tmp/wavtokenizer_ckpt/large_unify_40.ckpt \\
|
||
-o <output.wai>
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import struct
|
||
import subprocess
|
||
import sys
|
||
import wave
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
WAI_MAGIC = b"WAI1"
|
||
SR = 24_000
|
||
|
||
|
||
def encode(wav_path: str, out_path: str, src: str, cfg: str, ckpt: str) -> None:
|
||
sys.path.insert(0, src)
|
||
import torch
|
||
from decoder.pretrained import WavTokenizer
|
||
from encoder.utils import convert_audio
|
||
|
||
with wave.open(wav_path, "rb") as w:
|
||
sr_in = w.getframerate()
|
||
raw_bytes = w.getnframes() * 2 * w.getnchannels()
|
||
|
||
pcm = subprocess.run(
|
||
["ffmpeg", "-v", "error", "-i", wav_path, "-ac", "1",
|
||
"-ar", str(SR), "-f", "f32le", "-"],
|
||
check=True, capture_output=True).stdout
|
||
raw = np.frombuffer(pcm, np.float32)
|
||
|
||
print(f"loading WavTokenizer from {Path(ckpt).name}…")
|
||
model = WavTokenizer.from_pretrained0802(cfg, ckpt).eval()
|
||
|
||
# Reproduce infer.py's pipeline.
|
||
wav = torch.from_numpy(raw).reshape(1, -1)
|
||
wav = convert_audio(wav, SR, SR, 1)
|
||
bandwidth_id = torch.tensor([0])
|
||
with torch.no_grad():
|
||
features, discrete_codes = model.encode_infer(wav, bandwidth_id=bandwidth_id)
|
||
# discrete_codes: [n_q, batch=1, T]. Squeeze to [n_q, T].
|
||
codes = discrete_codes.squeeze(1).numpy().astype(np.uint16)
|
||
q, t = codes.shape
|
||
|
||
payload = struct.pack("<III", q, t, len(raw)) + codes.tobytes()
|
||
|
||
manifest = {
|
||
"wai": "1.0", "media": "audio", "intent": "replicate",
|
||
"model_requirement": {
|
||
"capability": "wai.neural.wavtokenizer",
|
||
"fallback": "wai.audio.opus",
|
||
},
|
||
"conditioning": {"kind": "wavtokenizer_tokens"},
|
||
"target": {"sr": SR, "n_samples": int(len(raw)),
|
||
"codebooks": int(q), "frames": int(t)},
|
||
}
|
||
mb = json.dumps(manifest, separators=(",", ":")).encode()
|
||
with open(out_path, "wb") as f:
|
||
f.write(WAI_MAGIC)
|
||
f.write(struct.pack("<I", len(mb))); f.write(mb)
|
||
f.write(struct.pack("<I", len(payload))); f.write(payload)
|
||
|
||
envelope_size = Path(out_path).stat().st_size
|
||
dur = len(raw) / SR
|
||
kbps = envelope_size * 8 / dur / 1000
|
||
|
||
print(f"input audio: {dur:.2f} s mono @ {sr_in} Hz "
|
||
f"({raw_bytes:>10,} B as int16 PCM)")
|
||
print(f"WavTokenizer: {q} codebook × {t} frames = {q*t} codes")
|
||
print(f"WAI envelope: {envelope_size:>10,} B "
|
||
f"({raw_bytes/envelope_size:6.1f}× compressed, "
|
||
f"{kbps:5.1f} kbps)")
|
||
print(f"model on the wire: 0 (WavTokenizer weights stay at the sink)")
|
||
print(f"capability declared:{manifest['model_requirement']['capability']:>30}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("input")
|
||
p.add_argument("-o", "--output", required=True)
|
||
p.add_argument("--wavtokenizer-src", default="/tmp/wavtokenizer",
|
||
help="Path to a clone of the WavTokenizer source repository.") # allow-name: open-source project the tool imports
|
||
p.add_argument("--config", required=True, help="YAML model config")
|
||
p.add_argument("--ckpt", required=True, help=".ckpt checkpoint")
|
||
a = p.parse_args()
|
||
encode(a.input, a.output, a.wavtokenizer_src, a.config, a.ckpt)
|