QaDiT: A 160M text-to-audio DiT trained on cosumer GPU's

sidharthGN1 pts0 comments

QuarkML/QaDiT · Hugging Face

Log In<br>Sign Up

QaDiT — text-to-audio latent Diffusion Transformer

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

Piece<br>Choice

Backbone<br>DiT-B — depth 12, width 768, 12 heads, MLP ratio 4.0 (~159M)

Latent grid<br>[8, 256, 16] (channels × time × freq)

Patchify<br>2×2 → 1024 tokens, fixed 2-D sincos positions

Text<br>FLAN-T5-large cross-attention every block + pooled text in adaLN-Zero

Train target<br>v-prediction

Noise schedule<br>cosine ᾱ, T = 1000

Timestep sampling (train)<br>logit-normal

CFG<br>p_uncond = 0.1 train; default guidance 4.0 at sample

Sampler<br>DDIM , default 50 steps, η = 0

Aux loss<br>REPA vs frozen AST features (train only)

Decode stack<br>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,<br>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

zt=αˉt z0+1−αˉt ε<br>z_t = \sqrt{\bar{\alpha}_t}\,z_0<br>+ \sqrt{1-\bar{\alpha}_t}\,\varepsilon<br>zt​=αˉt​​z0​+1−αˉt​​ε

Network target

v=αˉt ε−1−αˉt z0<br>v = \sqrt{\bar{\alpha}_t}\,\varepsilon<br>- \sqrt{1-\bar{\alpha}_t}\,z_0<br>v=αˉt​​ε−1−αˉt​​z0​

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

Dataset

Source<br>OpenSound/AudioCaps

Split<br>train · 45,178 clips after precompute

Clip length<br>10.24 s @ 16 kHz

Cached fields<br>VAE latents, FLAN-T5 embeddings + mask, AST REPA targets

latent_scale<br>0.9503493000009796 (baked into config.json)

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

Training run (this checkpoint)

Optimizer<br>AdamW, lr 1e-4, weight decay 0

Steps<br>23,999 (EMA exported)

Global batch<br>256 (2 GPUs × microbatch 16 × grad accum 8)

EMA decay<br>0.9999

REPA<br>weight 0.5, decayed over 15k steps

AMP<br>yes

Training curves

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

Diffusion / total loss

REPA loss

Usage

pip install transformers diffusers soundfile sentencepiece

import soundfile as sf<br>import torch<br>from transformers import AutoModel

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

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

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

Output types

output_type<br>Field<br>Content

"np" (default)<br>audios<br>list of float32 numpy waveforms in [-1, 1]

"pt"<br>audio_values<br>[B, num_samples] tensor

"latent"<br>latents<br>[B, 8, 256, 16] scaled latents (skips VAE/vocoder)

Precomputed T5 states

out = model.generate(<br>encoder_hidden_states=text_emb, # [B, 64, 1024]<br>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<br>DiT backbone runs in half precision; DDIM schedule math stays in float32.<br>Keep T5 / VAE / vocoder in float32 (FLAN-T5 overflows easily in fp16).

Important details

config.latent_scale (0.9503493 ) must match training precompute.<br>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,<br>ablations, and a starting checkpoint for longer / wider training.

Not intended for: production SFX libraries, speech synthesis, music<br>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...

training model latents ddim repa text

Related Articles