File size: 8,922 Bytes
50611ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 | # User Guide
## What This Model Does
Generates 2-second kick drum audio samples (WAV, 44.1kHz) from noise using a latent diffusion model. Supports optional text conditioning with keywords like "808", "punchy", "deep", etc.
## Files Required for Inference
To integrate into another project, you need these files:
### Model code (4 files)
```
models/
├── autoencoder.py # KickVAE (encoder/decoder)
├── diffusion.py # LatentUNet, NoiseScheduler
├── text_encoder.py # KeywordEncoder
└── vocoder.py # HiFiGANGenerator
```
### Inference code (2 files)
```
inference/
├── __init__.py
└── generate.py # Generation pipeline + CLI
```
### Weights (3 files)
```
weightss/
├── vae_epoch_100.pt # VAE decoder weights
├── diffusion_step_100000.pt # Diffusion U-Net + text encoder + vocab
└── vocoder_epoch_50.pt # HiFi-GAN generator weights
```
The diffusion checkpoint bundles everything needed for text conditioning: the U-Net weights (EMA), text encoder weights, vocabulary list, and model config.
### Python dependencies
```
torch>=2.10.0
torchaudio>=2.10.0
numpy>=2.4.2
scipy>=1.14.0
```
`scipy` is only used for writing WAV files (`scipy.io.wavfile.write`). If your project already handles WAV output differently, you can skip it and handle the waveform tensor directly.
## Python API
Import and call the `generate()` function directly:
```python
from pathlib import Path
from inference.generate import generate
output_path = generate(
diffusion_checkpoint=Path("checkpoints/diffusion_step_100000.pt"),
vae_checkpoint=Path("checkpoints/vae_epoch_100.pt"),
vocoder_checkpoint=Path("checkpoints/vocoder_epoch_50.pt"),
prompt="808",
cfg_scale=3.0,
ddim_steps=50,
output_path=Path("output/my_kick.wav"),
seed=42,
)
# output_path is the Path to the generated .wav file
```
### Parameters
| Parameter | Type | Default | Description |
| ---------------------- | -------------- | -------- | ----------------------------------------------------------------- |
| `diffusion_checkpoint` | `Path` | required | Path to diffusion model checkpoint |
| `vae_checkpoint` | `Path` | required | Path to VAE checkpoint |
| `vocoder_checkpoint` | `Path \| None` | `None` | Path to HiFi-GAN checkpoint. `None` falls back to Griffin-Lim |
| `prompt` | `str` | `""` | Space/comma-separated keywords (e.g. `"punchy 808"`) |
| `cfg_scale` | `float` | `3.0` | Classifier-free guidance strength. Higher = more prompt adherence |
| `ddim_steps` | `int` | `50` | DDIM sampling steps. More steps = higher quality, slower |
| `output_path` | `Path \| None` | `None` | Output WAV path. `None` auto-generates in `generations/` folder |
| `seed` | `int \| None` | `None` | Random seed for reproducible output |
### Return value
Returns a `Path` to the generated WAV file.
### Output format
- WAV, 44,100 Hz, mono, float32
- 2 seconds (88,200 samples)
- Normalized to 0.95 peak amplitude
### Working with the waveform tensor directly
If you need the raw tensor instead of a file (e.g. for streaming or further processing), you can replicate the pipeline steps from `generate()`:
```python
import torch
from pathlib import Path
from models.autoencoder import KickVAE
from models.diffusion import LatentUNet, NoiseScheduler
from models.text_encoder import KeywordEncoder
from models.vocoder import HiFiGANGenerator
from inference.generate import DDIMSampler, parse_prompt
device = torch.device("cpu")
# Load diffusion model
diff_ckpt = torch.load("checkpoints/diffusion_step_100000.pt", weights_only=False, map_location=device)
vocab = diff_ckpt["vocab"]
cfg = diff_ckpt["config"]
model = LatentUNet(latent_dim=cfg.latent_dim, base_channels=cfg.base_channels, cond_dim=cfg.cond_dim).to(device)
model.load_state_dict(diff_ckpt["ema_state_dict"])
model.eval()
text_enc = KeywordEncoder(vocab_size=len(vocab), embed_dim=cfg.text_embed_dim, cond_dim=cfg.cond_dim).to(device)
text_enc.load_state_dict(diff_ckpt["text_enc_state_dict"])
text_enc.eval()
scheduler = NoiseScheduler(cfg.timesteps, cfg.beta_start, cfg.beta_end).to(device)
# Encode prompt
token_ids = parse_prompt("808", vocab)
cond = text_enc([token_ids], device)
uncond = text_enc([[]], device)
# Sample latent
sampler = DDIMSampler(scheduler, num_steps=50)
latent = sampler.sample(model, shape=(1, cfg.latent_dim, 8, 11), cond=cond, uncond=uncond, cfg_scale=3.0, device=device)
# Decode with VAE
vae_ckpt = torch.load("checkpoints/vae_epoch_100.pt", weights_only=False, map_location=device)
vae = KickVAE(latent_dim=cfg.latent_dim).to(device)
vae.load_state_dict(vae_ckpt["model_state_dict"])
vae.eval()
with torch.no_grad():
log_mel = vae.decode(latent).squeeze(0) # (1, 128, 173)
# Vocoder
vocoder = HiFiGANGenerator(in_channels=128).to(device)
voc_ckpt = torch.load("checkpoints/vocoder_epoch_50.pt", weights_only=False, map_location=device)
vocoder.load_state_dict(voc_ckpt["generator"])
vocoder.eval()
vocoder.remove_weight_norm()
with torch.no_grad():
waveform = vocoder(log_mel) # (1, 1, ~88576)
waveform = waveform.squeeze()[:88200] # trim to 2 seconds
# waveform is a 1D tensor of float32 audio samples at 44100 Hz
```
## CLI Usage
Run from the `pytorch/` directory:
```bash
# Basic generation (unconditional)
uv run inference/generate.py --vocoder-ckpt checkpoints/vocoder_epoch_50.pt
# With text prompt
uv run inference/generate.py --vocoder-ckpt checkpoints/vocoder_epoch_50.pt --prompt "808"
# With custom settings
uv run inference/generate.py \
--vocoder-ckpt checkpoints/vocoder_epoch_50.pt \
--prompt "punchy analog" \
--cfg-scale 5.0 \
--steps 50 \
--seed 42
# Specify output path
uv run inference/generate.py \
--vocoder-ckpt checkpoints/vocoder_epoch_50.pt \
--output my_kick.wav
# Griffin-Lim fallback (no vocoder needed, lower quality)
uv run inference/generate.py --no-vocoder
```
### CLI Arguments
| Argument | Default | Description |
| ------------------ | -------------------------------------- | --------------------------------------------------------------- |
| `--prompt` | `""` | Text prompt keywords |
| `--cfg-scale` | `3.0` | Guidance scale |
| `--steps` | `50` | DDIM steps |
| `--seed` | None | Random seed |
| `--output` | auto | Output path (default: `generations/kick_<keywords>_<hash>.wav`) |
| `--diffusion-ckpt` | `checkpoints/diffusion_step_100000.pt` | Diffusion checkpoint |
| `--vae-ckpt` | `checkpoints/vae_epoch_100.pt` | VAE checkpoint |
| `--vocoder-ckpt` | `checkpoints/vocoder.pt` | Vocoder checkpoint |
| `--no-vocoder` | false | Use Griffin-Lim instead of HiFi-GAN |
## Text Conditioning
The model supports keyword-based text prompts. Prompts are matched against a vocabulary built from the training data filenames. Keywords that don't match any vocabulary entry are silently ignored.
To see what keywords are available, they are stored in the diffusion checkpoint:
```python
ckpt = torch.load("checkpoints/diffusion_step_100000.pt", weights_only=False)
print(ckpt["vocab"]) # list of valid keyword strings
```
When no prompt is given (or no keywords match), the model generates unconditionally.
**CFG scale** controls how strongly the output follows the prompt:
- `1.0` = no guidance (unconditional)
- `3.0` = default, moderate adherence
- `5.0+` = stronger adherence, may reduce variety
## Vocoder Options
**HiFi-GAN** (recommended): Higher quality output. Requires the vocoder checkpoint.
**Griffin-Lim** (fallback): Lower quality, no checkpoint needed. Uses iterative phase estimation to invert the mel spectrogram. Useful for quick testing.
## Device Support
The generation pipeline auto-detects the best available device: CUDA > MPS > CPU. All checkpoints are loaded with `map_location` so they work on any device regardless of where training was done.
|