# facebookresearch/audiobox-aesthetics architecture

> Architecture diagram of the GitHub repository facebookresearch/audiobox-aesthetics, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-20). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/facebookresearch/audiobox-aesthetics
- Repository: https://github.com/facebookresearch/audiobox-aesthetics

## Overview

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.

Main entry points: [`src/audiobox_aesthetics/model/aes.py`](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/model/aes.py) (AES model), [`src/audiobox_aesthetics/cli.py`](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/cli.py) (CLI), [`src/audiobox_aesthetics/infer.py`](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py) (Batching), [`README.md`](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/README.md) (README)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_g_runtime["Runtime"]
  node_n_cli["CLI<br/>entrypoint<br/>[cli.py]"]
  node_n_api["Python API<br/>entrypoint<br/>[infer.py]"]
  node_n_batch["Batching<br/>inference runner<br/>[infer.py]"]
  node_n_remote["SLURM mode<br/>cluster execution<br/>[cli.py]"]
  node_n_input["Input adapter<br/>audio ingest<br/>[infer.py]"]
  node_n_crop["Segment crop<br/>preprocess<br/>[infer.py]"]
  node_n_output["JSONL scores<br/>output schema<br/>[infer.py]"]
end

subgraph group_g_model["Model"]
  node_n_aes["AES model<br/>multi-output model<br/>[aes.py]"]
  node_n_wavlm["WavLM backbone<br/>audio encoder<br/>[wavlm.py]"]
  node_n_mutils["Model utils<br/>model helpers<br/>[utils.py]"]
  node_n_utils["Shared utils<br/>helpers<br/>[utils.py]"]
  node_n_ckpt[("Pretrained weights<br/>external artifact")]
end

subgraph group_g_data["Benchmarks"]
  node_n_eval["Eval data<br/>benchmark set"]
  node_n_track2["Track 2 data<br/>benchmark set"]
end

subgraph group_g_docs["Docs"]
  node_n_readme["README<br/>package docs<br/>[README.md]"]
end

node_n_cli -->|"runs"| node_n_batch
node_n_cli -->|"optional"| node_n_remote
node_n_api -->|"calls"| node_n_batch
node_n_batch -->|"reads"| node_n_input
node_n_input -->|"normalizes"| node_n_crop
node_n_crop -->|"feeds"| node_n_aes
node_n_aes -->|"encodes"| node_n_wavlm
node_n_aes -->|"loads"| node_n_ckpt
node_n_aes -->|"uses"| node_n_mutils
node_n_batch -->|"emits"| node_n_output
node_n_output -->|"streams"| node_n_cli
node_n_output -->|"returns"| node_n_api
node_n_utils -->|"supports"| node_n_input
node_n_eval -.->|"documents"| node_n_readme
node_n_track2 -.->|"documents"| node_n_readme
node_n_readme -.->|"explains"| node_n_cli

click node_n_cli "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/cli.py"
click node_n_api "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py"
click node_n_batch "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py"
click node_n_remote "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/cli.py"
click node_n_input "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py"
click node_n_crop "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py"
click node_n_output "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py"
click node_n_aes "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/model/aes.py"
click node_n_wavlm "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/model/wavlm.py"
click node_n_mutils "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/model/utils.py"
click node_n_utils "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/utils.py"
click node_n_eval "https://github.com/facebookresearch/audiobox-aesthetics/tree/main/evaluation_data"
click node_n_track2 "https://github.com/facebookresearch/audiobox-aesthetics/tree/main/audiomos2025_track2"
click node_n_readme "https://github.com/facebookresearch/audiobox-aesthetics/blob/main/README.md"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_n_cli,node_n_api,node_n_batch,node_n_remote,node_n_input,node_n_crop,node_n_output toneBlue
class node_n_aes,node_n_wavlm,node_n_mutils,node_n_utils,node_n_ckpt toneAmber
class node_n_eval,node_n_track2 toneMint
class node_n_readme toneRose
```

## Components

### Runtime

User-facing inference entry points and execution modes

- [CLI](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/cli.py): entrypoint · audio-aes command-line interface · `src/audiobox_aesthetics/cli.py`
- [Python API](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py): entrypoint · Library entry for predictor initialization and inference · `src/audiobox_aesthetics/infer.py`
- [Batching](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py): inference runner · Batched forward path over JSONL inputs · `src/audiobox_aesthetics/infer.py`
- [SLURM mode](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/cli.py): cluster execution · Chunked remote execution for managed compute · `src/audiobox_aesthetics/cli.py`
- [Input adapter](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py): audio ingest · Accepts file paths or waveform tensors with sample rate metadata · `src/audiobox_aesthetics/infer.py`
- [Segment crop](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py): preprocess · Optional start/end time slicing before scoring · `src/audiobox_aesthetics/infer.py`
- [JSONL scores](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/infer.py): output schema · Per-item predictions for CE, CU, PC, and PQ · `src/audiobox_aesthetics/infer.py`

### Model

Pretrained model, backbone, and supporting code

- [AES model](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/model/aes.py): multi-output model · High-level pretrained aesthetics predictor · `src/audiobox_aesthetics/model/aes.py`
- [WavLM backbone](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/model/wavlm.py): audio encoder · Feature extractor backbone adapted from WavLM lineage · `src/audiobox_aesthetics/model/wavlm.py`
- [Model utils](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/model/utils.py): model helpers · Supporting utilities for model and data handling · `src/audiobox_aesthetics/model/utils.py`
- [Shared utils](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/src/audiobox_aesthetics/utils.py): helpers · General package utilities · `src/audiobox_aesthetics/utils.py`
- **Pretrained weights**: external artifact · Default checkpoint fetched automatically when absent

### Benchmarks

Packaged evaluation and task data

- [Eval data](https://github.com/facebookresearch/audiobox-aesthetics/tree/main/evaluation_data): benchmark set · JSONL annotation sets for speech, music, natural sound, and PAM · `evaluation_data`
- [Track 2 data](https://github.com/facebookresearch/audiobox-aesthetics/tree/main/audiomos2025_track2): benchmark set · AudioMOS 2025 Track 2 lists and human annotations · `audiomos2025_track2`

### Docs

Repository documentation and release material

- [README](https://github.com/facebookresearch/audiobox-aesthetics/blob/main/README.md): package docs · Usage, model, and inference documentation · `README.md`

## Connections

- CLI → Batching: runs · CLI dispatches local batch inference
- CLI → SLURM mode: optional · CLI can switch into remote SLURM execution
- Python API → Batching: calls · Python API initializes and executes predictor flow
- Batching → Input adapter: reads · Consumes JSONL rows or in-memory wave inputs
- Input adapter → Segment crop: normalizes · Prepares and optionally slices audio segments
- Segment crop → AES model: feeds · Cleaned examples are scored by the model
- AES model → WavLM backbone: encodes · Model uses WavLM-based audio features
- AES model → Pretrained weights: loads · Pretrained weights are loaded from external artifacts
- AES model → Model utils: uses · Model-specific helper functions support inference
- Batching → JSONL scores: emits · Produces one JSON object per input line
- JSONL scores → CLI: streams · Output is line-aligned for CLI piping
- JSONL scores → Python API: returns · API exposes the same prediction schema
- Shared utils → Input adapter: supports · General helpers assist the ingest path
- Eval data → README: documents · Benchmark datasets are described in repo docs
- Track 2 data → README: documents · Task-specific data layout is documented
- README → CLI: explains · README describes CLI usage and flags