README.md

metadata

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

AudioModel v1 🔊

"rain falling on a roof"

"footsteps on gravel"

"a dog barking"

"a car engine starting"

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:

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

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

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