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:

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():

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:

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

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.