README.md

---
license: apache-2.0
library_name: transformers
tags:
  - text-to-audio
  - diffusion
  - dit
  - audio
  - educational
  - research
pipeline_tag: text-to-audio
---

<p align="center">
  <a href="https://github.com/sidharth72/QaDiT">
    <img src="https://img.shields.io/badge/GitHub-Repository-181717?logo=github&logoColor=white" alt="GitHub">
  </a>
  <a href="https://www.quarkml.com">
    <img src="https://img.shields.io/badge/Website-QuarkML-blue?logo=googlechrome&logoColor=white" alt="Website">
  </a>
  <a href="https://huggingface.co/QuarkML">
    <img src="https://img.shields.io/badge/Hugging%20Face-Models-yellow?logo=huggingface&logoColor=black" alt="Hugging Face">
  </a>
</p>

# 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.

> Example:
> Prompt: `A small waterfall flows through a forest while insects buzz and birds sing.`
>
> Output:
> 
> <audio controls src="https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/_RGLUAPxdImTitbAAD0g6.wav"></audio> 

---


| 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

![image](https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/W4tXcMyotNTC01hEYTz6q.png)

---

## 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.

![image](https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/ok-iYdQGBXNXgnBzKtXdS.png)

---

## 3. Training step (what is actually optimized)

Only the **DiT** and its small glue layers receive gradients.


![image](https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/vpUCimXElF6yv-M9UBI3Z.png)

---

## 4. Inside one DiT block

![image](https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/mmTgnW0w7BAyCKLC7gHiA.png)

---

## 5. Inference / sampling (caption → waveform)


![image](https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/-4i-SA1wIcQ589NKJOPAe.png)

---

## 6. Component ownership

![image](https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/d7AJNoRcB47CztYuZt-jC.png)

---

## 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](https://huggingface.co/datasets/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/`](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


![image](https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/JXuluU4tNmP6z5tuirIkb.png)


### REPA loss


![image](https://cdn-uploads.huggingface.co/production/uploads/64054e5e0ab5e22719fc179f/pGbo0f5lliOenjMsGmFP6.png)

---

## Usage

```bash
pip install transformers diffusers soundfile sentencepiece
```

```python
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

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

### Single denoising step

```python
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:

1. **Train longer** — continue past 24k steps with the same recipe (or lower LR).
2. **Scale the dataset** — mix in larger captioned audio corpora beyond AudioCaps.
3. **Retune sampling** — CFG scale, DDIM step count, and prompt formatting.
4. **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

```bibtex
@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](https://huggingface.co/datasets/OpenSound/AudioCaps) |
| VAE / vocoder | [cvssp/audioldm-s-full-v2](https://huggingface.co/cvssp/audioldm-s-full-v2) |
| Text encoder | [google/flan-t5-large](https://huggingface.co/google/flan-t5-large) |