README.md

---
license: mit
pipeline_tag: text-to-audio
language:
- en
tags:
- audio
- text-to-audio
- diffusion
- rectified-flow
- spectrogram
- clap
- custom-code
- tiny
---

# AudioModel v1 🔊

<p>
<audio controls src="https://huggingface.co/bench-labs/AudioModel-v1/resolve/main/samples/rain_falling_on_a_roof.wav"></audio> "rain falling on a roof"
</p>
<p>
<audio controls src="https://huggingface.co/bench-labs/AudioModel-v1/resolve/main/samples/footsteps_on_gravel.wav"></audio> "footsteps on gravel"
</p>
<p>
<audio controls src="https://huggingface.co/bench-labs/AudioModel-v1/resolve/main/samples/a_dog_barking.wav"></audio> "a dog barking"
</p>
<p>
<audio controls src="https://huggingface.co/bench-labs/AudioModel-v1/resolve/main/samples/a_car_engine_starting.wav"></audio> "a car engine starting"
</p>

A tiny text-to-audio diffusion model. You give it a sentence, it gives you 17.8
seconds of sound. It is 40,296,844 trained parameters and it trained in under four
hours on a single H100.

This is the PixelModel line moving into a third modality. v1 through v6 were tiny
latent diffusion for images, VoxelModel did 3D occupancy grids, and this one treats
a mel spectrogram as a grayscale image and runs the same DiT plus rectified flow
recipe on it. Same argument every time: low resolution as a deliberate aesthetic,
applied to a new kind of output.

## Why this combination didn't already exist

Riffusion proved that treating a spectrogram as an image and running image
diffusion on it works, including inverting it back to audio. But Riffusion is a
fine-tune of Stable Diffusion v1.5, not trained from scratch, and not tiny.
tiny-audio-diffusion is from scratch and genuinely small, but it diffuses raw
waveforms, not spectrograms, and isn't built for free-text conditioning.
AudioLDM and Make-An-Audio are the real text-to-audio research lineage, but both
are full-scale systems with a VAE-compressed latent space and a separately trained
vocoder, not something you train from scratch on a rented GPU in an afternoon.

Nothing combined from-scratch, hobbyist scale, spectrogram-as-image, rectified
flow, and free-text CLAP conditioning in one place. That's the gap this fills.

## The one honest result

FAD (Frechet Audio Distance, PANN embeddings) on 300 held-out Clotho evaluation
clips, comparing three things against real audio:

![FAD comparison](audiomodel-v1-fad.png)

Text conditioning clearly helps: 67.0 with a real prompt against 87.1 with no
prompt at all. And the model is not that far from the ceiling a Griffin-Lim
vocoder allows in the first place: taking real audio, converting it to our exact
mel spectrogram, and inverting it right back with Griffin-Lim already costs 60.5
FAD against the original file, before the model has generated anything. Most of
the remaining distance between the model and real audio is the vocoder, not the
model.

Worth knowing if you use this number: PANN embeddings are 2048-dimensional and our
full held-out evaluation split is only 1,045 clips, so the covariance estimate
behind FAD is structurally rank-deficient here. Scipy warns about it
(`LinAlgWarning: Matrix is singular`) and it would still warn even using every
single held-out clip we have, since 1,045 is still under 2,048. The relative
ordering of the three numbers is the trustworthy part, not their precision to one
decimal place.

We also found a real bug while computing this: the `frechet_audio_distance`
library's PANN path concatenates each clip's flat 2048-length embedding with
`np.concatenate` instead of stacking them, which silently collapses 300 separate
embeddings into one 614,400-length vector and breaks the statistics downstream.
Their own README admits PANN isn't covered by CI. `eval_fad.py` reshapes the
embeddings back to `(n, 2048)` before computing FAD, which is the whole fix.

## Training

![training loss](audiomodel-v1-loss.png)

| | |
|---|---|
| clips | 3,839 train (development split), 1,045 held out (validation split) |
| pairs | 19,195, since each clip carries 5 independent captions |
| steps | 90,000 at batch 256 |
| hardware | one H100 SXM 80GB, 6.51 steps/s, 28.3 GiB |
| wall clock | 3.85 hours |
| optimizer | AdamW 2e-4, betas (0.9, 0.99), grad clip 1.0 |
| schedule | cosine to 1e-6, 500 step warmup |
| EMA | 0.9999 |
| CFG dropout | 10% |

Held-out loss bottomed at **0.06961 at step 54,000** and drifted up to 0.07078 by
step 90,000. That's a real overfitting inflection, and it arrives earlier than it
did on VoxelModel (step 64,000 of 90,000 there), which makes sense: Clotho's
development split has 3,839 unique clips against VoxelModel's 28,415 unique
meshes, so the same step budget means a lot more repetition per unique example
here. The shipped weights are the step 54,000 EMA checkpoint, not the step 90,000
final one, because this run saves the best validation checkpoint separately as it
trains. VoxelModel v1 didn't do this and lost its own best checkpoint to an
overwrite bug. That was worth fixing before this run, not after.

## No VAE, but not zero baggage either

AudioModel diffuses directly on the mel spectrogram image. There's no autoencoder
compressing it first, which is the same "no asterisk" choice VoxelModel made for
voxels instead of copying AudioLDM's VAE step. Total trained footprint is 40.30M.

Conditioning is frozen CLAP (`laion/clap-htsat-unfused`), cross-attention on the
RoBERTa text tower's last hidden state (768 dim, 32 tokens) plus adaLN-zero from
the projected pooled embedding (512 dim). Loading `ClapModel` pulls in the whole
153.49M-parameter checkpoint, audio tower included, but only the text tower and
projection, 125.30M of it, ever runs a forward pass here. The HTSAT audio encoder
just sits there unused. If you're adapting this code, swapping in a text-only CLAP
checkpoint would trim that dead weight; we didn't bother for a single training run.

## Architecture

`AudioDiT` patchifies a 384x256 mel spectrogram at patch size 16, giving a 24x16
grid, 384 tokens, dim 384, depth 12, heads 6. Positional encoding is a 2D sincos
over that grid. Objective is rectified flow with `target = x1 - x0` and
logit-normal timestep sampling, `t = sigmoid(randn)`, identical to VoxelModel's
objective on a different kind of image.

## Data

Clotho v2.1 (Drossos, Lipping and Virtanen, DCASE 2020), from Zenodo, not from an
HF mirror since none of the official kind exists. Development split for training,
validation split for the held-out loss during training, evaluation split reserved
untouched for the FAD numbers above. Each clip carries 5 independent crowdsourced
captions, and all 5 get used as separate training pairs against the same
spectrogram rather than picking one, which is a free 5x on the text side without
touching the audio side.

Spectrograms come from `diffusers.pipelines.deprecated.audio_diffusion.mel.Mel`
at `x_res=384, hop_length=1024`, both bumped from the library's defaults
(`x_res=256, hop_length=512`) to cover 17.8 seconds instead of 5.9. `y_res=256`
mel bins is the library default, left alone. The `Mel` class already pads short
clips with silence and only reads the first slice of longer ones, so every clip
in the 15 to 30 second Clotho range ends up as a fixed 17.8 second window starting
at the beginning of the clip, no custom windowing logic needed.

Griffin-Lim, 32 iterations, is the inversion back to audio. It's the zero-training
default and it's known to sound phasey and metallic compared to a learned vocoder
like BigVGAN-v2. We didn't swap one in for v1: the FAD numbers above suggest the
vocoder actually is the dominant source of remaining distance from real audio, so
that's the correct place to spend effort on a v2, not on the diffusion model
itself.

## Usage

```bash
python sample.py "a dog barking" "rain falling on a roof" --cfg 4.0 --steps 50
```

Writes 17.8 second WAVs at 22,050 Hz to `--out-dir`.

## What v1 does not do

**A learned vocoder.** Griffin-Lim is what's shipped. The FAD breakdown above is
the argument for why that's a defensible v1 choice and not just laziness.

**More than 300 clips of FAD signal.** The evaluation split only has 1,045 clips
total and PANN embeddings are 2048-dimensional, so the covariance estimate behind
every FAD number in this card is on the edge of what's estimable at all. Treat the
ordering as real and the decimal places as decoration.

**Augmentation.** 90,000 steps over 3,839 unique clips is a lot of repetition with
nothing to break it up. Time-shift or pitch-shift augmentation on the waveform
before spectrogram conversion is the obvious next lever, same shape of fix as the
yaw-and-mirror augmentation flagged as unbuilt on VoxelModel v1.

## Files

| file | what |
|---|---|
| `model.safetensors` | EMA weights at step 90,000 |
| `model_best.safetensors` | EMA weights at step 54,000, the shipped checkpoint |
| `audio_dit.py` | the model |
| `train.py` | training loop |
| `sample.py` | inference |
| `eval_fad.py` | FAD evaluation against held-out real audio |
| `fetch_clotho.py` | Zenodo download and extraction |
| `build_data.py` | mel spectrogram and CLAP embedding precompute |
| `benchmark_throughput.py` | the throughput check this run's batch size came from |
| `config.json` | architecture and training config |
| `fad_results.json` | the raw numbers behind the FAD chart |
| `requirements.txt` | what you need installed to run any of the above |
| `samples/` | the four clips embedded at the top of this card |