facebookresearch/audiobox-aesthetics

1. What this repository is

- audiobox-aesthetics is a small inference-focused Python package for unified audio quality estimation across speech, music, and sound.

- The architecture is centered on a pretrained model + lightweight CLI/Python wrapper, not on training infrastructure or a larger application stack.

2. Primary user-facing entry points

- CLI: exposed as audio-aes, driven by src/audiobox_aesthetics/cli.py.

- Python API: src/audiobox_aesthetics/infer.py exposes initialize_predictor() and a .forward(...) path for batch inference.

- Hugging Face model loading: src/audiobox_aesthetics/model/aes.py supports from_pretrained(...) for fine-tuning or reuse.

3. Core inference flow

- The input contract is a JSONL stream, where each row identifies an audio source and optionally a time span (start_time, end_time).

- Inputs can be either file paths or already-loaded wave tensors with sample rate metadata.

- The predictor runs the model and emits one JSON object per input, containing four prediction axes: CE, CU, PC, PQ.

4. Model subsystem

- src/audiobox_aesthetics/model/aes.py is the high-level model definition, likely the multi-output prediction head used for all axes.

- src/audiobox_aesthetics/model/wavlm.py provides the WavLM-based backbone; the README notes this code is partly adapted from Microsoft’s unilm/wavlm.

- src/audiobox_aesthetics/model/utils.py and src/audiobox_aesthetics/utils.py contain supporting model/data utilities and general helpers.

5. Data preprocessing boundary

- The repo handles audio ingest at the edge: load from file, or accept a tensor directly.

- The optional start_time / end_time fields indicate the inference path can crop temporal segments before scoring.

- This makes the architecture suitable for both full-file scoring and segment-level scoring without changing the model interface.

6. Batching and throughput

- The CLI is designed for large-batch inference (--batch-size), suggesting the core predictor operates on batched examples rather than single files.

- Output remains line-aligned with input, so the batch execution layer is a throughput optimization only; it does not change the data contract.

7. Distributed / remote execution

- The CLI has a SLURM mode (--remote, --array, --job-dir, --chunk, plus GPU/CPU settings).

- This is a clear architectural boundary: local inference is one path, and cluster-scheduled chunked execution is another.

- The repository therefore supports both workstation use and large-scale batch scoring on managed compute.

8. Checkpoint and model distribution

- A pretrained checkpoint is available from an S3-hosted download and mirrored on Hugging Face.

- If no checkpoint path is supplied, the tooling will auto-fetch the default model, so model weights are treated as an external artifact rather than source code.

- Fine-tuning workflows are expected to go through the Hugging Face model class.

9. Output schema

- The system produces a compact JSON-per-line output with four continuous scores:

- CE — Content Enjoyment

- CU — Content Usefulness

- PC — Production Complexity

- PQ — Production Quality

- This output format is stable and designed for downstream processing with tools like jq or custom pipelines.

10. Evaluation data packaged in the repo

- evaluation_data/ contains JSONL annotation sets for several audio domains: speech, music, natural sound, and PAM.

- audiomos2025_track2/ contains task-specific train/dev lists and human annotations for the AudioMOS 2025 Track 2 setup.

- These datasets are adjacent to the code but function as benchmark inputs, not runtime dependencies.

11. Repository layout boundaries

- src/ holds the installable library code.

- assets/ contains documentation imagery only.

- evaluation_data/ and audiomos2025_track2/ are data/benchmark artifacts.

- The root docs (README.md, CHANGELOG.md, contribution files) support package usage and release management rather than application logic.

12. Technology profile

- The package targets Python 3.9+ and PyTorch 2.2+.

- Audio handling is implied to depend on the PyTorch audio stack for waveform loading and tensor-based inference.

- The model architecture is research-model oriented, with the WavLM backbone being the main external dependency lineage.

13. How to think about the architecture diagram

- Draw a simple pipeline: JSONL or tensor input → preprocessing/cropping/loading → WavLM-based multi-output model → 4-axis JSONL predictions.

- Surround that with two operational shells: local CLI/Python API and SLURM chunked execution.

- Treat pretrained checkpoints and benchmark datasets as external assets feeding or validating the pipeline, not as core runtime services.