QaDiT — text-to-audio latent Diffusion Transformer

A ~159M-parameter latent Diffusion Transformer that turns a text caption into 10.24 s of 16 kHz mono audio: FLAN-T5 conditioning → DiT denoising of AudioLDM KL-VAE latents → VAE decode → HiFi-GAN vocoder.

Piece	Choice
Backbone	DiT-B — depth 12, width 768, 12 heads, MLP ratio 4.0 (~159M)
Latent grid	`[8, 256, 16]` (channels × time × freq)
Patchify	2×2 → 1024 tokens, fixed 2-D sincos positions
Text	FLAN-T5-large cross-attention every block + pooled text in adaLN-Zero
Train target	v-prediction
Noise schedule	cosine ᾱ, T = 1000
Timestep sampling (train)	logit-normal
CFG	p_uncond = 0.1 train; default guidance 4.0 at sample
Sampler	DDIM, default 50 steps, η = 0
Aux loss	REPA vs frozen AST features (train only)
Decode stack	`cvssp/audioldm-s-full-v2` VAE + HiFi-GAN

1. Big picture — three pipelines

2. Offline pre-compute (frozen models, run once)

Heavy frozen models run once; the training loop never loads T5, the VAE, or the REPA encoder.

3. Training step (what is actually optimized)

Only the DiT and its small glue layers receive gradients.

4. Inside one DiT block

5. Inference / sampling (caption → waveform)

6. Component ownership

Training objective (DDIM + v-prediction)

Forward process

$z_t = \sqrt{\bar{\alpha}_t}\,z_0 + \sqrt{1-\bar{\alpha}_t}\,\varepsilon$

Network target

$v = \sqrt{\bar{\alpha}_t}\,\varepsilon - \sqrt{1-\bar{\alpha}_t}\,z_0$

At sample time the DiT predicts v; we recover \hat{z}_0 and \hat{\varepsilon}, then step with DDIM (\eta = 0). CFG is applied in v-space with default scale (s = 4.0). After DDIM, latents are divided by latent_scale ≈ 0.95035 before VAE decode — that whole chain is what model.generate() runs.

Dataset


Source	OpenSound/AudioCaps
Split	`train` · 45,178 clips after precompute
Clip length	10.24 s @ 16 kHz
Cached fields	VAE latents, FLAN-T5 embeddings + mask, AST REPA targets
`latent_scale`	`0.9503493000009796` (baked into `config.json`)

AudioCaps is captioned environmental / everyday sound — not speech or music. Those domains are out of distribution for this checkpoint.

Training run (this checkpoint)


Optimizer	AdamW, lr 1e-4, weight decay 0
Steps	23,999 (EMA exported)
Global batch	256 (2 GPUs × microbatch 16 × grad accum 8)
EMA decay	0.9999
REPA	weight 0.5, decayed over 15k steps
AMP	yes

Training curves

Put W&B / TensorBoard screenshots (or exports) under assets/ using the filenames below. Until then the images show as broken links on the Hub — that is intentional so the slots are obvious.

Diffusion / total loss

REPA loss

Usage

pip install transformers diffusers soundfile sentencepiece

import soundfile as sf
import torch
from transformers import AutoModel

model = AutoModel.from_pretrained("QuarkML/QaDiT", trust_remote_code=True)
model = model.to("cuda" if torch.cuda.is_available() else "cpu").eval()

out = model.generate(
    "A small waterfall flows through a forest while insects buzz and birds sing.",
    num_inference_steps=200,
    guidance_scale=16.0,
    seed=0,
)
sf.write("sample.wav", out.audios[0], out.sampling_rate)

First generate downloads the frozen helpers this run was trained with: google/flan-t5-large and the VAE + vocoder from cvssp/audioldm-s-full-v2.

Output types

`output_type`	Field	Content
`"np"` (default)	`audios`	list of float32 numpy waveforms in `[-1, 1]`
`"pt"`	`audio_values`	`[B, num_samples]` tensor
`"latent"`	`latents`	`[B, 8, 256, 16]` scaled latents (skips VAE/vocoder)

Precomputed T5 states

out = model.generate(
    encoder_hidden_states=text_emb,      # [B, 64, 1024]
    encoder_attention_mask=text_mask,    # [B, 64]
)

Single denoising step

v = model(latents, timesteps, encoder_hidden_states, encoder_attention_mask).sample

Precision and devices

Runs on CPU and CUDA. With dtype=torch.float16 or torch.bfloat16 the DiT backbone runs in half precision; DDIM schedule math stays in float32. Keep T5 / VAE / vocoder in float32 (FLAN-T5 overflows easily in fp16).

Important details

config.latent_scale (0.9503493) must match training precompute. generate divides by it before VAE decode.
Every sample is fixed length: 10.24 s @ 16 kHz.
repa_layer exists for REPA fine-tuning; inference ignores it.
Sampling always uses the EMA weights packaged here.

Limitations and intended use

Intended use: education, reproduction of a small latent DiT audio stack, ablations, and a starting checkpoint for longer / wider training.

Not intended for: production SFX libraries, speech synthesis, music generation, or safety-critical audio.

Known limits of this checkpoint

~24k steps on ~45k AudioCaps clips — undertrained vs public SOTA systems
Weak on speech, music, and densely described scenes
Inherits caption biases and coverage holes of AudioCaps
Prefer the default 50 DDIM steps for demos; low step counts sound coarse

Research artifact — how to improve this

This release is a research artifact, not a production host model. The architecture and sampling path are solid enough to build on; the ceiling is mostly data and compute:

Train longer — continue past 24k steps with the same recipe (or lower LR).
Scale the dataset — mix in larger captioned audio corpora beyond AudioCaps.
Retune sampling — CFG scale, DDIM step count, and prompt formatting.
Keep measuring — log diffusion loss, REPA loss, and listening tests.

Those levers will move quality more than inventing a new backbone for this size of model. Contributions and longer runs are welcome; treat this Hub page as a reproducible baseline, not a finished product.

Citation

@misc{qadit2026,
  title        = {QaDiT: A Text-to-Audio Latent Diffusion Transformer},
  author       = {Sidharth GN},
  year         = {2026},
  note         = {Research artifact. Weights and transformers remote-code loading.}
}

Resource
Dataset	OpenSound/AudioCaps
VAE / vocoder	cvssp/audioldm-s-full-v2
Text encoder	google/flan-t5-large

Downloads last month: -

Safetensors

Model size

0.2B params

Tensor type

F32