open-standards/wai/tools/wai_mimi_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

93 lines
3.1 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.mimi`.
Mimi (arXiv 2410.00037) at 24 kHz, 12.5 Hz frame rate, 8 codebooks of 2048.
Designed for real-time speech use. The token rate is 8 codebooks of 11 bits at
12.5 Hz, about 1.1 kbps on the wire.
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)
Usage:
python tools/wai_mimi_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 = "kyutai/mimi" # allow-name: model-hub id the tool defaults to
WAI_MAGIC = b"WAI1"
SR = 24_000
def encode(wav_path: str, out_path: str) -> None:
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)
import torch
from transformers import MimiModel, AutoFeatureExtractor
model = MimiModel.from_pretrained(MODEL_ID).eval()
proc = AutoFeatureExtractor.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"], inp.get("padding_mask"))
# MimiEncoderOutput.audio_codes: [batch=1, n_codebooks, 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.mimi",
"fallback": "wai.audio.opus",
},
"conditioning": {"kind": "mimi_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"Mimi 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 (Mimi 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)