a2f-rs

NVIDIA Audio2Face-3D v2.3.1 (James) — audio → 169 face coefficients, in Rust, with two flavors behind one CLI:

Provenance stays NVIDIA end to end — no third-party weight repacks. The graph mirror reference is soniqo/speech-swift (Audio2Face3DMLXRuntime.swift, Apache-2.0); the MLX graph in this repo was parity-validated against the official ONNX ground truth.

Model: https://huggingface.co/nvidia/Audio2Face-3D-v2.3.1-James (NVIDIA Open Model License).

Layout

crates/a2f-core   constants, wav IO, sliding-window framing, emotion, npy out
crates/a2f-onnx   ort backend (load-dynamic; needs ORT_DYLIB_PATH)
crates/a2f-mlx    mlx-rs backend, hand-written forward graph
crates/a2f-cli    `a2f` binary: run / parity, --backend onnx|mlx
tools/            export_mlx_weights.py (official ONNX → mlx bundle)
tests/fixtures    parity fixtures (seeded noise + anger=0.9 → 169 coefs, f32 raw)

Build & run

cargo build --release

# MLX flavor (Metal, no extra deps)
cargo run -r -p a2f-cli -- run --backend mlx \
  --bundle <mlx-bundle> --wav speech.wav --out coefs.npy

# ONNX flavor (needs a libonnxruntime dylib)
ORT_DYLIB_PATH=$(python3 -c 'import onnxruntime,os;print(os.path.join(os.path.dirname(onnxruntime.__file__),"capi","libonnxruntime.1.31.0.dylib"))') \
cargo run -r -p a2f-cli -- run --backend onnx \
  --bundle <nvidia-bundle> --wav speech.wav
# or: brew install onnxruntime → /opt/homebrew/opt/onnxruntime/lib/

# emotion control (10 explicit channels)
... run --backend mlx --emotion joy=0.8 --emotion anger=0.2 ...

Bundle setup (one-time)

# 1. official NVIDIA bundle (file downloads via hf-mirror endpoint if you like)
huggingface-cli download nvidia/Audio2Face-3D-v2.3.1-James --local-dir nvidia-bundle

# 2. mlx flavor weights, extracted from that same official ONNX
python3 tools/export_mlx_weights.py nvidia-bundle -o mlx-bundle

mlx-bundle = audio2face3d.safetensors + export_manifest.json + copied configs. nvidia-bundle is used directly by the onnx flavor.

Parity

Fixtures: seeded rng(42) noise window [1,1,8320] + emotion with anger=0.9; ground truth from onnxruntime 1.31 CPU on the official ONNX.

cargo run -r -p a2f-cli -- parity --backend mlx  --bundle tests/fixture-bundle --fixtures tests/fixtures
cargo run -r -p a2f-cli -- parity --backend onnx --bundle <nvidia-bundle>    --fixtures tests/fixtures
# PASS bar (same as speech-swift): maxAbs < 0.05 && meanAbs < 0.005

Measured on M1 Max (this repo, release build):

Residual is pure float-accumulation order: the ONNX graph's LeakyRelu are all α=0 (≡ ReLU) and MLX computes on Metal GPU.

Throughput (same machine, release):

The mlx flavor is the production path; the onnx flavor is the reference / ground-truth path (also where a CoreML/CUDA EP would slot in).

Output contract

169 f32 per frame @30 fps (hop 533, window 8320 = ±260 ms around the label):

[0..140)   skin    (regression geometry space — NOT normalized ARKit blendshapes)
[140..150) tongue
[150..164) jaw
[164..169) eye

Output 1:1 with the official ONNX result — downstream consumers (mesh eval via shapes_mean_skin/shapes_matrix_skin, retargets, renderers) don't care which flavor produced the coefficients.

Notes / traps

• MLX tensor-name variants: the ONNX initializer set contains two graph families ("2230" / "2232"); "2230" is the live one for this bundle — decided by parity, not by name. tools/export_mlx_weights.py only exports the resolved variant.
• rust-toolchain.toml pins 1.95 (ort rc.14 needs ≥1.92).
• mlx-rs builds Apple MLX from source on first compile (cmake + Metal); be patient once.
• If your global git has core.autocrlf=true (this machine does), the FetchContent git clone inside mlx-sys checks out MLX's *.sh helpers with CRLF and the cmake build dies on $'\r': command not found. Handled automatically here: .cargo/config.toml sets GIT_CONFIG_GLOBAL=no-crlf.gitconfig (repo-local, autocrlf=false) for all cargo-invoked builds, so no manual env is needed.