a2f-mlx 0.1.0

# a2f-rs

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

| flavor | crate | runtime | weights |
|---|---|---|---|
| `onnx` | `a2f-onnx` | ONNX Runtime (`ort`, load-dynamic) | **official `network.onnx`, as-is** |
| `mlx`  | `a2f-mlx`  | Apple MLX via `mlx-rs`, hand-written graph (Metal) | exported from the **same official ONNX** by `tools/export_mlx_weights.py` |

Provenance stays NVIDIA end to end — no third-party weight repacks. The graph
mirror reference is [soniqo/speech-swift](https://github.com/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

```bash
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)

```bash
# 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.

```bash
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):

| check | maxAbs | meanAbs |
|---|---|---|
| `parity --backend mlx` vs ONNX ground truth | 0.028908 | 0.000447 |
| `parity --backend onnx` vs ground truth | 0.000011 | 0.000001 |
| mlx vs onnx, real 8 s zh speech, 242 frames | 0.023552 | 0.000113 |

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

| flavor | real zh 8.03 s → 242 frames | RTF |
|---|---|---|
| `mlx` (Metal) | 2.8 s wall | **0.35 — 2.8× realtime** |
| `onnx` (CPU EP) | 18.8 s wall | 2.34 — 0.4× realtime |

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.