9474a402d8
- int_audio: the direct path's fast dispatch (conv_t_1d over int_kernels) is deleted; the test-only direct decode runs each transposed convolution on its scalar kernel (reference::conv_t_1d, verbatim, itself test-only now). Its guard layer_in_range, and run_layer, live only inside the test-only legacy module: the production audio decode's exactness checks are its graph's check ranges, the PCM step's included. - The kernel_equivalence module (the dispatch against the scalar kernel; retargeted at the interpreter at gate 5) is deleted. Its edge geometry moves into int_audio's geometry tests, through the lowered graph under every policy: a one-sample input upsampled by any stride a WAM1 field holds; several samples whose uncropped length is past the cap (too_large, the fast kernel never entered); padding past the fast kernels' pad factor (the interpreter computes it); an input that could overflow, refused by the check range before any arithmetic. - int_hyper_synth's docs: igdn_check's audio caller and run_net_checked's are test-only now; the conditional codec's checked decode is the last production caller (M4). int_kernels_geometry's and the export tool's notes follow. Co-Authored-By: Claude Opus 5.5 <[email protected]>
501 lines
21 KiB
Python
501 lines
21 KiB
Python
#!/usr/bin/env python3
|
||
"""Train a small LEARNED 1-D audio transform codec and export its SYNTHESIS as an
|
||
INTEGER-exact decoder — the audio analog of `int_hyper`, extending "conformance is
|
||
a hash" to learned audio.
|
||
|
||
A neural audio codec decodes through a FLOAT network, so a ~1e-5 round-off makes the
|
||
reconstructed waveform — and its hash — diverge between machines: never receiptable.
|
||
Here the synthesis is a stack of integer 1-D transposed convolutions + integer IGDN
|
||
(per-out-channel int12 weights, fixed-point activations, all i64), so the
|
||
reconstructed PCM is byte-identical on every machine (NeuralIntegerExact) —
|
||
*sample-equivalence*, a portable PCM hash.
|
||
|
||
**The signal is generated here; the tool reads no audio file.** Both the training
|
||
signal and the test signal (the clip the vector codes) are written by `signal()`
|
||
below from a seed: 16 kHz, mono, signed 16-bit, a sequence of segments, each a tone
|
||
(one to three partials with an attack and a release), a chirp (a linear frequency
|
||
sweep), a noise burst (white noise, low-passed by a moving average or high-passed by
|
||
a difference) or a speech-like formant sweep (a harmonic series on a gliding,
|
||
vibrato-modulated pitch, shaped by three formant resonances that move between two
|
||
vowel targets), sometimes two overlaid, with short gaps. The generator is integer
|
||
arithmetic only — a SplitMix64 counter stream, a 32-bit phase accumulator, a
|
||
4096-entry sine table built by Bhaskara I's rational approximation — so a seed gives
|
||
the same samples on every machine and every Python. The training signal is stream 0
|
||
of the seed and the test signal stream 1: the codec is trained on one signal and
|
||
codes another of the same kind.
|
||
|
||
Pipeline (mirrors wai_sota_int_synth_export, in 1-D):
|
||
1. generate the training and test signals from the seed;
|
||
2. train g_a/g_s (Conv1d+GDN / ConvTranspose1d+IGDN) on the training signal (float;
|
||
this is the encoder, float is fine here — only the DECODE must be integer), with
|
||
torch seeded from the same seed, deterministic algorithms and a fixed thread
|
||
count, so a run on one host and torch build gives the same weights every time;
|
||
3. quantize the test signal's latent per channel (factorized prior, data-fit CDF +
|
||
rANS, the byte-exact match of the Rust RansDecoder);
|
||
4. PTQ g_s → integer (WAM1 binary), run an INTEGER reference decode (numpy,
|
||
mirroring the Rust i64 arithmetic) → int16 PCM → FNV-1a hash.
|
||
|
||
python tools/wai_int_audio_export.py <out-dir> [--seed S] [--channels C]
|
||
[--samples N] [--steps K] [--train-seconds T]
|
||
python tools/wai_int_audio_export.py --signal-only <out.s16> [--seed S] [--stream 0|1]
|
||
[--samples N]
|
||
|
||
`--signal-only` writes one generated signal as raw s16le and needs nothing beyond the
|
||
Python standard library. The committed vector (`wai-rs/tests/vectors/int_audio/`) and
|
||
the site's demo clip are both exports of this tool; their commands are in that
|
||
directory's README.
|
||
"""
|
||
import argparse
|
||
import base64
|
||
import json
|
||
import math
|
||
import os
|
||
import struct
|
||
import sys
|
||
|
||
AFRAC = 12; WBITS = 12; SHIFT = 24; GF = 16 # same fixed-point as int_hyper
|
||
SR = 16000; T = 2048; K = 5; STR = 2; NLAY = 4 # codec shape (16x downsample)
|
||
LR = 1.5e-3; RATE_W = 8e-6
|
||
THREADS = 4 # fixed: CPU reductions depend on it
|
||
|
||
# ---- the generated signal (integer arithmetic only) ----------------------------------
|
||
M64 = (1 << 64) - 1
|
||
GOLDEN = 0x9E3779B97F4A7C15
|
||
TAB_BITS = 12
|
||
TAB_N = 1 << TAB_BITS
|
||
|
||
|
||
def _mix64(x):
|
||
"""SplitMix64's finalizer."""
|
||
x &= M64
|
||
x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & M64
|
||
x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & M64
|
||
return x ^ (x >> 31)
|
||
|
||
|
||
class Rng:
|
||
"""A SplitMix64 stream, keyed by (seed, stream)."""
|
||
|
||
def __init__(self, seed, stream):
|
||
self.s = _mix64(seed * GOLDEN + stream * 0xD1B54A32D192ED03)
|
||
|
||
def next(self):
|
||
self.s = (self.s + GOLDEN) & M64
|
||
return _mix64(self.s)
|
||
|
||
def int(self, lo, hi):
|
||
"""Uniform integer in [lo, hi]."""
|
||
return lo + self.next() % (hi - lo + 1)
|
||
|
||
def s16(self):
|
||
"""Uniform signed 16-bit value."""
|
||
return (self.next() >> 48) - 32768
|
||
|
||
|
||
def _sine_table():
|
||
"""sin(2π i / TAB_N) × 32767, by Bhaskara I's approximation, rounded half up:
|
||
sin(πp) ≈ 16 p (1 − p) / (5 − 4 p (1 − p)), p = u / half."""
|
||
half = TAB_N // 2
|
||
out = []
|
||
for i in range(TAB_N):
|
||
u = i % half
|
||
a = u * (half - u)
|
||
den = 5 * half * half - 4 * a
|
||
v = (2 * 32767 * 16 * a + den) // (2 * den)
|
||
out.append(v if i < half else -v)
|
||
return out
|
||
|
||
|
||
SINE = _sine_table()
|
||
|
||
|
||
def _inc(f_q8):
|
||
"""32-bit phase increment per sample for a frequency in Hz × 256."""
|
||
return (f_q8 << 24) // SR
|
||
|
||
|
||
def _osc(phase):
|
||
return SINE[(phase & 0xFFFFFFFF) >> (32 - TAB_BITS)]
|
||
|
||
|
||
def _envelope(n, attack, release):
|
||
"""Linear attack and release, Q15, length n."""
|
||
attack, release = max(1, min(attack, n // 2)), max(1, min(release, n // 2))
|
||
return [min(32767, (32767 * (i + 1)) // attack, (32767 * (n - i)) // release) for i in range(n)]
|
||
|
||
|
||
def _root64(x):
|
||
"""2^(x/64) in Q16, floored: the largest y with y^64 ≤ 2^(x + 16·64). Exact."""
|
||
lo, hi = 1 << 16, 1 << 17
|
||
while lo < hi:
|
||
mid = (lo + hi + 1) // 2
|
||
if mid ** 64 <= 1 << (x + 1024):
|
||
lo = mid
|
||
else:
|
||
hi = mid - 1
|
||
return lo
|
||
|
||
|
||
ROOT64 = [_root64(x) for x in range(64)]
|
||
|
||
|
||
def _log_freq(r, lo, hi):
|
||
"""A frequency in Hz × 256, log-uniform in [lo, hi] Hz: lo · 2^(k/64) for a random
|
||
k, drawn again while it is above hi."""
|
||
octaves = 0
|
||
while (lo << octaves) < hi:
|
||
octaves += 1
|
||
while True:
|
||
o, j = divmod(r.int(0, 64 * octaves), 64)
|
||
f = ((lo << o) * ROOT64[j]) >> 8
|
||
if f <= hi << 8:
|
||
return f
|
||
|
||
|
||
# Formant targets (F1, F2, F3) in Hz for five vowels.
|
||
VOWELS = [(730, 1090, 2440), (270, 2290, 3010), (300, 870, 2240), (530, 1840, 2480), (570, 840, 2410)]
|
||
|
||
|
||
def _tone(r, n):
|
||
out = [0] * n
|
||
for _ in range(r.int(1, 3)):
|
||
inc, amp, ph = _inc(_log_freq(r, 80, 3500)), r.int(4000, 16000), r.next() & 0xFFFFFFFF
|
||
for i in range(n):
|
||
out[i] += (amp * _osc(ph)) >> 15
|
||
ph += inc
|
||
env = _envelope(n, r.int(SR // 500, SR // 30), r.int(SR // 50, SR // 6))
|
||
return [(v * e) >> 15 for v, e in zip(out, env)]
|
||
|
||
|
||
def _chirp(r, n):
|
||
i0, i1 = _inc(_log_freq(r, 100, 6000)), _inc(_log_freq(r, 100, 6000))
|
||
amp, ph, out = r.int(8000, 20000), r.next() & 0xFFFFFFFF, []
|
||
for i in range(n):
|
||
out.append((amp * _osc(ph)) >> 15)
|
||
ph += i0 + (i1 - i0) * i // n
|
||
env = _envelope(n, r.int(SR // 500, SR // 50), r.int(SR // 100, SR // 20))
|
||
return [(v * e) >> 15 for v, e in zip(out, env)]
|
||
|
||
|
||
def _noise(r, n):
|
||
w = [r.s16() for _ in range(n + 16)]
|
||
kind = r.int(0, 2)
|
||
if kind == 0: # white
|
||
x = w[:n]
|
||
elif kind == 1: # low-passed: moving average
|
||
L = [2, 3, 5, 8, 13][r.int(0, 4)]
|
||
x = [sum(w[i:i + L]) // L for i in range(n)]
|
||
else: # high-passed: first difference
|
||
x = [(w[i + 1] - w[i]) >> 1 for i in range(n)]
|
||
env = _envelope(n, r.int(SR // 1000, SR // 200), r.int(SR // 50, SR // 5))
|
||
amp = r.int(4000, 14000)
|
||
return [(((v * amp) >> 15) * e) >> 15 for v, e in zip(x, env)]
|
||
|
||
|
||
def _formant(r, n):
|
||
"""A harmonic series on a gliding pitch with vibrato, each harmonic weighted by
|
||
three resonances 1 / (1 + ((f − F)/B)^2) whose centres move between two vowels."""
|
||
f0a = r.int(90, 240) << 8
|
||
f0b = (f0a * r.int(80, 125)) // 100
|
||
va, vb = VOWELS[r.int(0, 4)], VOWELS[r.int(0, 4)]
|
||
bw = (r.int(60, 110), r.int(90, 160), r.int(120, 220))
|
||
vib_inc, vib_depth = _inc(r.int(4, 7) << 8), r.int(5, 25) # depth in ‰
|
||
nh = max(1, (4800 << 8) // max(f0a, f0b))
|
||
ph0, vph, out = r.next() & 0xFFFFFFFF, r.next() & 0xFFFFFFFF, []
|
||
amps = [0] * nh
|
||
for i in range(n):
|
||
f0 = f0a + (f0b - f0a) * i // n
|
||
f0 += (f0 * vib_depth * _osc(vph)) // (1000 * 32767)
|
||
vph += vib_inc
|
||
if i % 32 == 0: # the weights at block rate
|
||
F = [va[j] + (vb[j] - va[j]) * i // n for j in range(3)]
|
||
for k in range(nh):
|
||
fk = ((k + 1) * f0) >> 8
|
||
g = 0
|
||
for j in range(3):
|
||
d = ((fk - F[j]) << 10) // bw[j]
|
||
g += (1 << 20) // (1024 + ((d * d) >> 10)) # Q10 resonance gain
|
||
amps[k] = (g * 1024) // (1024 + 256 * k) # spectral tilt
|
||
ph0 += _inc(f0)
|
||
acc = 0
|
||
for k in range(nh):
|
||
acc += amps[k] * _osc((k + 1) * ph0)
|
||
out.append(acc >> 15)
|
||
env = _envelope(n, r.int(SR // 200, SR // 20), r.int(SR // 50, SR // 10))
|
||
return [(v * e) >> 15 for v, e in zip(out, env)]
|
||
|
||
|
||
KINDS = [_tone, _chirp, _noise, _formant]
|
||
|
||
|
||
def signal(seed, stream, n):
|
||
"""n samples of the generated signal (a list of ints in the int16 range)."""
|
||
r = Rng(seed, stream)
|
||
out = []
|
||
while len(out) < n:
|
||
m = r.int(SR * 8 // 100, SR // 2) # a segment: 80–500 ms
|
||
seg = KINDS[r.int(0, 3)](r, m)
|
||
if r.int(0, 3) == 0: # sometimes a second, quieter kind
|
||
over = KINDS[r.int(0, 3)](r, m)
|
||
seg = [a + (b >> 2) for a, b in zip(seg, over)]
|
||
peak = max(1, max(abs(v) for v in seg))
|
||
target = r.int(6000, 28000)
|
||
seg = [(v * target) // peak for v in seg]
|
||
out += seg + [0] * r.int(0, SR // 25) # a gap: 0–40 ms
|
||
return [max(-32768, min(32767, v)) for v in out[:n]]
|
||
|
||
|
||
def s16le(x):
|
||
return struct.pack(f'<{len(x)}h', *x)
|
||
|
||
|
||
# ---- learned model (float; the encoder) -------------------------------------
|
||
def build_codec(C):
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
class GDN1d(nn.Module):
|
||
"""1-D (inverse) generalized divisive normalization with a full C×C gamma.
|
||
Forward divides by sqrt(beta + gamma·x²); inverse multiplies — the exact form
|
||
crate::int_transform::igdn reproduces in integer (out = x·isqrt(norm))."""
|
||
def __init__(self, ch, inverse):
|
||
super().__init__()
|
||
self.inverse = inverse
|
||
self.rg = nn.Parameter(torch.randn(ch, ch) * 0.01) # raw → softplus → γ≥0
|
||
self.rb = nn.Parameter(torch.zeros(ch)) # raw → softplus → β≈1
|
||
|
||
def gb(self):
|
||
return F.softplus(self.rg), F.softplus(self.rb + 1.0)
|
||
|
||
def forward(self, x): # x [N,C,T]
|
||
g, b = self.gb()
|
||
norm = b[None, :, None] + torch.einsum('ij,njt->nit', g, x * x)
|
||
return x * torch.sqrt(norm) if self.inverse else x / torch.sqrt(norm)
|
||
|
||
class Codec(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
ga, gs = [], []
|
||
for i in range(NLAY):
|
||
ga += [nn.Conv1d(1 if i == 0 else C, C, K, STR, K // 2), GDN1d(C, False)]
|
||
for i in range(NLAY):
|
||
last = i == NLAY - 1
|
||
gs += [nn.ConvTranspose1d(C, 1 if last else C, K, STR, K // 2, output_padding=STR - 1)]
|
||
if not last:
|
||
gs += [GDN1d(C, True)]
|
||
self.g_a, self.g_s = nn.Sequential(*ga), nn.Sequential(*gs)
|
||
|
||
def forward(self, x):
|
||
y = self.g_a(x)
|
||
y_hat = y + (torch.rand_like(y) - 0.5) # STE quantization noise
|
||
return self.g_s(y_hat), y
|
||
|
||
return Codec()
|
||
|
||
|
||
# ---- integer reference decode (numpy; mirrors the Rust i64 arithmetic) -------
|
||
def conv_t_1d_int(xi, wi, mult, bias, stride, pad, opad):
|
||
"""Integer 1-D transposed conv + per-out-channel (mult>>SHIFT)+bias — byte-for-
|
||
byte what the Rust decoder's transposed convolution does (integer-operators §13.3,
|
||
wam1/1). xi [cin,T]; wi [cin,cout,k]."""
|
||
import numpy as np
|
||
cin, cout, k = wi.shape; Tin = xi.shape[1]
|
||
To = (Tin - 1) * stride + k + opad - 2 * pad
|
||
Tp = To + 2 * pad
|
||
raw = np.zeros((cout, Tp), dtype=object)
|
||
for kx in range(k):
|
||
contrib = np.tensordot(wi[:, :, kx], xi, axes=([0], [0])) # [cout,Tin]
|
||
raw[:, kx:kx + stride * Tin:stride] += contrib
|
||
out = (raw[:, pad:pad + To] * mult[:, None]) >> SHIFT
|
||
return out + bias[:, None]
|
||
|
||
|
||
def igdn_int_1d(x, gi, bi):
|
||
"""Integer IGDN, the EXACT arithmetic of crate::int_transform::igdn (per-term
|
||
>>GF before the sum, isqrt at 2^AFRAC, (x·s)>>AFRAC). x [C,T] python-int."""
|
||
import numpy as np
|
||
Cc, Tt = x.shape
|
||
x2 = x * x # 2^(2·AFRAC)
|
||
norm = np.tile(np.array(bi, dtype=object)[:, None], (1, Tt)) # 2^(2·AFRAC)
|
||
for j in range(Cc):
|
||
norm = norm + ((gi[:, j:j + 1] * x2[j:j + 1, :]) >> GF)
|
||
s = np.vectorize(lambda v: math.isqrt(int(v)) if v > 0 else 0, otypes=[object])(norm)
|
||
return (x * s) >> AFRAC
|
||
|
||
|
||
def run_g_s_int(xi, layers):
|
||
import numpy as np
|
||
x = xi.astype(object)
|
||
for L in layers:
|
||
if L['type'] == 'convT':
|
||
wi = np.array(L['weight'], dtype=object).reshape(L['in_ch'], L['out_ch'], L['k'])
|
||
x = conv_t_1d_int(x, wi, np.array(L['mult'], dtype=object),
|
||
np.array(L['bias'], dtype=object), L['stride'], L['pad'], L['opad'])
|
||
else: # igdn
|
||
c = L['c']
|
||
gi = np.array(L['gamma'], dtype=object).reshape(c, c)
|
||
x = igdn_int_1d(x, gi, np.array(L['beta'], dtype=object))
|
||
return x # [1,T] @ 2^AFRAC
|
||
|
||
|
||
def to_pcm_i16(x_fixed):
|
||
"""[1,T] @ 2^AFRAC (≈[-1,1]) → int16 PCM, round-half-up, clamp. Matches Rust."""
|
||
import numpy as np
|
||
v = x_fixed[0]
|
||
num = v * 32767 + (1 << (AFRAC - 1))
|
||
pcm = num >> AFRAC # python>> floors (== i64 arith shift)
|
||
return np.array([max(-32768, min(32767, int(p))) for p in pcm], dtype=np.int16)
|
||
|
||
|
||
# ---- PTQ + entropy helpers ---------------------------------------------------
|
||
def q_weight_pc(w): # ConvTranspose1d [in,out,k] → per-OUT int12
|
||
import numpy as np
|
||
s = np.abs(w).max(axis=(0, 2), keepdims=True) / (2 ** (WBITS - 1) - 1) + 1e-12
|
||
wi = np.round(w / s).clip(-(2 ** (WBITS - 1)), 2 ** (WBITS - 1) - 1).astype(np.int64)
|
||
return wi, s.reshape(-1)
|
||
|
||
|
||
def fit_cdf(sym, lo, hi):
|
||
import numpy as np
|
||
n = hi - lo + 1
|
||
hist = np.bincount(np.asarray(sym, np.int64) - lo, minlength=n).astype(np.float64)
|
||
total = 1 << 16
|
||
f = np.maximum((hist / hist.sum() * (total - n)).astype(np.int64), 1)
|
||
f[f.argmax()] += (total - 1) - int(f.sum())
|
||
cdf = [0]
|
||
for v in f:
|
||
cdf.append(cdf[-1] + int(v))
|
||
cdf.append(total)
|
||
return cdf
|
||
|
||
|
||
def fnv1a(buf):
|
||
h = 14695981039346656037
|
||
for b in buf:
|
||
h = ((h ^ b) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
|
||
return h
|
||
|
||
|
||
def main(outdir, seed, C, nsamp, steps, train_seconds):
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from compressai.ans import RansEncoder
|
||
|
||
assert nsamp % T == 0, f'--samples must be a multiple of {T}'
|
||
os.makedirs(outdir, exist_ok=True)
|
||
torch.manual_seed(seed)
|
||
torch.use_deterministic_algorithms(True)
|
||
torch.set_num_threads(THREADS)
|
||
|
||
ntrain = (train_seconds * SR // T) * T
|
||
train = np.array(signal(seed, 0, ntrain), np.float32) / 32768.0
|
||
test_i16 = signal(seed, 1, nsamp)
|
||
raw = np.array(test_i16, np.float32) / 32768.0
|
||
frames = torch.from_numpy(train.reshape(-1, 1, T)) # [Nf,1,T]
|
||
|
||
m = build_codec(C)
|
||
opt = torch.optim.Adam(m.parameters(), LR)
|
||
# Keep the rate term tiny: the latent must stay O(1+) so integer rounding (the
|
||
# quantizer) preserves it — too much rate pressure collapses it below the grid.
|
||
for step in range(steps):
|
||
xr, y = m(frames)
|
||
loss = F.mse_loss(xr, frames) + RATE_W * y.abs().mean()
|
||
opt.zero_grad(); loss.backward(); opt.step()
|
||
m.eval()
|
||
with torch.no_grad():
|
||
xr_f, _ = m(frames)
|
||
snr_f = 10 * np.log10((frames ** 2).sum() / ((xr_f - frames) ** 2).sum() + 1e-12)
|
||
print(f' (float recon SNR on the training signal {float(snr_f):.1f} dB)')
|
||
|
||
# the test signal's latent → per-channel factorized symbols (s = round(y - median))
|
||
with torch.no_grad():
|
||
xfull = torch.from_numpy(raw.reshape(1, 1, nsamp))
|
||
y = m.g_a(xfull)[0].numpy() # [C, L]
|
||
L = y.shape[1]
|
||
median = np.round(np.median(y, axis=1)) # integer-grid median
|
||
sym = np.round(y - median[:, None]).astype(np.int64) # [C, L] symbols
|
||
median_q = np.round(median * (2 ** AFRAC)).astype(np.int64)
|
||
xi = (sym * (2 ** AFRAC) + median_q[:, None]).astype(np.int64) # fixed-point latent
|
||
|
||
# PTQ g_s → WAM1 integer model
|
||
layers = []
|
||
for layer in m.g_s.children():
|
||
t = type(layer).__name__
|
||
if t == 'ConvTranspose1d':
|
||
wi, sc = q_weight_pc(layer.weight.detach().numpy())
|
||
layers.append({'type': 'convT', 'in_ch': int(wi.shape[0]), 'out_ch': int(wi.shape[1]),
|
||
'k': int(wi.shape[2]), 'stride': int(layer.stride[0]), 'pad': int(layer.padding[0]),
|
||
'opad': int(layer.output_padding[0]), 'weight': wi.reshape(-1).tolist(),
|
||
'mult': np.round(sc * (2 ** SHIFT)).astype(np.int64).tolist(),
|
||
'bias': np.round(layer.bias.detach().numpy() * (2 ** AFRAC)).astype(np.int64).tolist()})
|
||
else: # GDN1d (inverse)
|
||
with torch.no_grad():
|
||
g, b = layer.gb()
|
||
gi = np.round(g.numpy() * (2 ** GF)).astype(np.int64)
|
||
bi = np.round(b.numpy() * (2 ** (2 * AFRAC))).astype(np.int64)
|
||
layers.append({'type': 'igdn', 'c': int(b.shape[0]), 'gamma': gi.reshape(-1).tolist(), 'beta': bi.tolist()})
|
||
|
||
# INTEGER reference decode → PCM (what the Rust port reproduces)
|
||
x_fixed = run_g_s_int(xi, layers)
|
||
pcm = to_pcm_i16(x_fixed)
|
||
snr = 10 * np.log10((raw[:len(pcm)] ** 2).sum() / (((pcm.astype(np.float64) / 32768) - raw[:len(pcm)]) ** 2).sum() + 1e-12)
|
||
|
||
# factorized rANS over the symbols (channel-major), per-channel data-fit CDF
|
||
lo, hi = int(sym.min()), int(sym.max())
|
||
cdfs, syms_flat, idx = [], [], []
|
||
for ch in range(C):
|
||
cdfs.append(fit_cdf(sym[ch], lo, hi))
|
||
syms_flat.extend(sym[ch].tolist())
|
||
idx.extend([ch] * L)
|
||
enc = RansEncoder()
|
||
bits = enc.encode_with_indexes(syms_flat, idx, cdfs, [len(cdfs[0])] * C, [lo] * C)
|
||
|
||
# WAM1 binary model
|
||
mb = bytearray(b'WAM1')
|
||
mb += struct.pack('<IIII', AFRAC, GF, SHIFT, len(layers))
|
||
for Lr in layers:
|
||
if Lr['type'] == 'convT':
|
||
mb += struct.pack('<BIIIIII', 0, Lr['in_ch'], Lr['out_ch'], Lr['k'], Lr['stride'], Lr['pad'], Lr['opad'])
|
||
mb += np.array(Lr['weight'], dtype='<i2').tobytes()
|
||
mb += np.array(Lr['mult'], dtype='<i8').tobytes() + np.array(Lr['bias'], dtype='<i8').tobytes()
|
||
else:
|
||
c = Lr['c']
|
||
mb += struct.pack('<BI', 2, c)
|
||
assert np.abs(np.array(Lr['gamma'])).max() < 2 ** 31, 'gamma > i32'
|
||
mb += np.array(Lr['gamma'], dtype='<i4').tobytes() + np.array(Lr['beta'], dtype='<i8').tobytes()
|
||
open(f'{outdir}/audio_model.bin', 'wb').write(mb)
|
||
|
||
json.dump({'n_ch': C, 'lat_len': L, 't_out': int(len(pcm)), 'afrac': AFRAC,
|
||
'lo': lo, 'cdfs': cdfs, 'median_q': median_q.tolist(),
|
||
'rans_b64': base64.b64encode(bytes(bits)).decode()},
|
||
open(f'{outdir}/audio_bitstream.json', 'w'))
|
||
pcm_hash = format(fnv1a(pcm.tobytes()), '016x')
|
||
json.dump({'pcm_fnv1a': pcm_hash, 'n_samples': int(len(pcm)), 'sample_rate': SR},
|
||
open(f'{outdir}/audio.json', 'w'))
|
||
pcm.tofile(f'{outdir}/audio_decoded.s16')
|
||
open(f'{outdir}/audio_source.s16', 'wb').write(s16le(test_i16))
|
||
print(f'learned 1-D audio codec (seed {seed}): {len(pcm)} samples @ {SR}Hz, C={C} L={L}, '
|
||
f'model {len(mb)}B, rANS {len(bits)}B -> integer-decode SNR {snr:.1f} dB, pcm_hash {pcm_hash}')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
ap = argparse.ArgumentParser(description='Train and export the integer learned audio codec on a '
|
||
'generated signal (no audio input).')
|
||
ap.add_argument('outdir', nargs='?', help='where audio_model.bin, audio_bitstream.json and audio.json go')
|
||
ap.add_argument('--seed', type=int, default=1)
|
||
ap.add_argument('--channels', type=int, default=32, help='latent channels C')
|
||
ap.add_argument('--samples', type=int, default=32768, help='test-signal (coded clip) length, a multiple of 2048')
|
||
ap.add_argument('--steps', type=int, default=3500, help='training steps')
|
||
ap.add_argument('--train-seconds', type=int, default=8, help='training-signal length in seconds')
|
||
ap.add_argument('--signal-only', metavar='OUT.s16', help='write one generated signal and exit')
|
||
ap.add_argument('--stream', type=int, default=1, help='with --signal-only: 0 training, 1 test')
|
||
a = ap.parse_args()
|
||
if a.signal_only:
|
||
n = a.samples if a.stream == 1 else (a.train_seconds * SR // T) * T
|
||
open(a.signal_only, 'wb').write(s16le(signal(a.seed, a.stream, n)))
|
||
sys.exit(0)
|
||
if not a.outdir:
|
||
ap.error('outdir is required unless --signal-only is given')
|
||
main(a.outdir, a.seed, a.channels, a.samples, a.steps, a.train_seconds)
|