File size: 6,245 Bytes
80b78f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """Reproduce Core AI export and validate every output over independent state trajectories."""
import argparse
import asyncio
import json
import os
from pathlib import Path
import numpy as np
import torch
os.environ.setdefault("USE_OS_COREAI", "1")
from coreai.runtime import AIModel, AIModelAssetMetadata, NDArray, SpecializationOptions
from coreai_torch import TorchConverter, get_decomp_table
from models import ROOT, load_model
def psnr(ref, actual):
a, b = np.asarray(ref, dtype=np.float64), np.asarray(actual, dtype=np.float64)
mse = np.mean((a - b) ** 2)
if mse == 0:
return 300.0
return float(10 * np.log10(max(np.max(np.abs(a)), 1e-12) ** 2 / mse))
def export(variant, dtype):
model, args, ins, outs, _checkpoint, _ = load_model(variant, dtype)
ep = torch.export.export(model, args=args, strict=False).run_decompositions(
get_decomp_table()
)
program = (
TorchConverter()
.add_exported_program(ep, input_names=ins, output_names=outs)
.to_coreai()
)
path = ROOT / "exports" / f"{variant}_{str(dtype).split('.')[-1]}_streaming.aimodel"
if path.exists():
raise FileExistsError(path)
meta = AIModelAssetMetadata()
meta.author = (
"Xiaobin Rong and contributors"
if variant == "ulunas_dns3"
else "Jean-Marc Valin and RNNoise contributors"
)
meta.license = (
"MIT"
if variant == "ulunas_dns3"
else "BSD-3-Clause; see vendor/rnnoise/LICENSE and source-file notices"
)
meta.model_description = f"{variant} pretrained streaming neural network; Core AI conversion by Max Farrell. Audio DSP is external. See README and provenance.json."
program.save_asset(path, meta)
return path
async def validate(variant, dtype, path, count=32, wav=None):
model, args, ins, outs, _, source = load_model(variant, dtype)
runtime = await AIModel.load(path, SpecializationOptions.cpu_only())
fn = runtime.load_function("main")
assert list(fn.desc.input_names) == ins
assert list(fn.desc.output_names) == outs
ts, cs = args[1:], args[1:]
gen = torch.Generator().manual_seed(1234)
frames = torch.randn(
(1, 257, count, 2) if variant == "ulunas_dns3" else (1, count, 65),
generator=gen,
).to(dtype)
if variant == "ulunas_dns3":
import soundfile as sf
if wav:
audio, sr = sf.read(wav, dtype="float32")
else:
sr = 16000
t = np.arange(16000, dtype=np.float32) / sr
audio = (
0.15 * np.sin(2 * np.pi * 180 * t)
+ 0.05 * np.random.default_rng(1234).standard_normal(16000)
).astype(np.float32)
assert sr == 16000 and audio.ndim == 1
z = torch.stft(
torch.tensor(audio)[None],
512,
256,
512,
torch.hann_window(512),
return_complex=True,
)
frames = torch.view_as_real(z)[:, :, :count].to(dtype)
with torch.no_grad():
spec = frames.permute(0, 3, 2, 1)
feat = source.erb.bm(
torch.log10(torch.norm(spec, dim=1, keepdim=True).clamp(1e-12))
)
feat, skips = source.encoder(feat)
mask = source.erb.bs(source.decoder(source.dpgrnn(feat), skips))
offline = (spec * mask).permute(0, 3, 2, 1)
scores = {name: [] for name in outs}
streaming = []
ref_scores = []
with torch.no_grad():
for i in range(count):
x = (
frames[:, :, i : i + 1, :]
if variant == "ulunas_dns3"
else frames[:, i : i + 1, :]
)
pre_state = ts
ref = model(x, *ts)
runtime_out = await fn(
{
name: NDArray(data=t.contiguous())
for name, t in zip(ins, (x, *cs), strict=True)
}
)
values = [runtime_out[name].numpy() for name in outs]
ts = tuple(t.detach() for t in ref[-len(ts) :])
cs = tuple(torch.from_numpy(t.copy()) for t in values[-len(cs) :])
for name, r, v in zip(outs, ref, values, strict=True):
assert np.isfinite(v).all(), name
scores[name].append(psnr(r.numpy(), v))
if variant == "ulunas_dns3":
streaming.append(ref[0])
elif i >= 4:
direct = source(frames[:, i - 4 : i + 1, :], list(pre_state[2:]))
ref_scores.append(
min(
psnr(direct[0].numpy(), ref[0].numpy()),
psnr(direct[1].numpy(), ref[1].numpy()),
)
)
if streaming:
ref_scores = [psnr(offline.numpy(), torch.cat(streaming, dim=2).numpy())]
result = {
"variant": variant,
"precision": str(dtype),
"device": "CPU",
"frames": count,
"input_fixture": str(wav) if wav else "deterministic synthetic input",
"minimum_psnr_db": {k: min(v) for k, v in scores.items()},
"source_graph_psnr_db": min(ref_scores),
}
print(json.dumps(result, indent=2))
assert min(ref_scores) >= 70, "Streaming wrapper differs from upstream graph"
assert min(min(v) for v in scores.values()) >= 40, "Core AI parity failed"
(
ROOT / "docs" / f"{variant}_{str(dtype).split('.')[-1]}_validation.json"
).write_text(json.dumps(result, indent=2) + "\n")
if __name__ == "__main__":
p = argparse.ArgumentParser(description=__doc__)
p.add_argument(
"variant", choices=["ulunas_dns3", "rnnoise10Ga_12", "rnnoise10Gb_15"]
)
p.add_argument("--dtype", choices=["float16", "float32"], default="float32")
p.add_argument("--validate-only", action="store_true")
p.add_argument(
"--wav", type=Path, help="Optional 16 kHz mono WAV for UL-UNAS parity"
)
a = p.parse_args()
dtype = getattr(torch, a.dtype)
path = (
ROOT / "exports" / f"{a.variant}_{a.dtype}_streaming.aimodel"
if a.validate_only
else export(a.variant, dtype)
)
asyncio.run(validate(a.variant, dtype, path, wav=a.wav))
|