open-standards/wai/tools/wai_dac_encode.py
Transaction Science fb103ac5ae wai: one registry table, candidate capabilities, public-text checks
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]>
2026-09-30 22:38:09 -04:00

98 lines
3.4 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""WAI source-side neural encoder for `wai.neural.dac`.
DAC (improved RVQGAN, arXiv 2306.06546) at 44.1 kHz, 9 codebooks. The model is the
sink-installed decoder; only the quantized codes ship in the WAI
envelope: 9 codebooks of 10 bits at about 86 frames a second, about 8 kbps.
Wire format (uniform across all wai.neural.<audio> capabilities):
<III> q t n_samples (little-endian u32 × 3)
q*t × u16 codes (row-major, LE)
For DAC q=9 codebooks. n_samples is the original sample count so the
sink can trim trailing-silence padding the model emits.
Usage:
python tools/wai_dac_encode.py <input.wav> -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
MODEL_ID = "descript/dac_44khz" # allow-name: model-hub id the tool defaults to
WAI_MAGIC = b"WAI1"
SR = 44_100
def encode(wav_path: str, out_path: str) -> None:
# Read original audio length (for the round-trip kbps math).
with wave.open(wav_path, "rb") as w:
sr_in = w.getframerate()
raw_bytes = w.getnframes() * 2 * w.getnchannels() # int16 PCM size
# ffmpeg → mono float32 @ 44.1 kHz
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)
import torch
from transformers import DacModel, AutoProcessor
model = DacModel.from_pretrained(MODEL_ID).eval()
proc = AutoProcessor.from_pretrained(MODEL_ID)
inp = proc(raw_audio=raw, sampling_rate=SR, return_tensors="pt")
with torch.no_grad():
enc = model.encode(inp["input_values"])
# DacEncoderOutput.audio_codes: [batch=1, n_codebooks=9, T]
codes = enc.audio_codes.squeeze(0).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.dac",
"fallback": "wai.audio.opus",
},
"conditioning": {"kind": "dac_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"DAC tokens: {q} codebooks × {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 (DAC 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)
a = p.parse_args()
encode(a.input, a.output)