| """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)) |
|
|
