A dusty GTX 1660 and a weekend are all you need. This tutorial walks through training a working kick drum diffusion model on 6GB of VRAM, from dataset prep to first generated sample.
You don't need an H100 to train a generative audio model. A dusty Linux tower with a GTX 1660 and 6GB of VRAM can train a working kick drum diffusion model over a weekend, and the samples that come out are surprisingly usable.
This tutorial walks through the full pipeline: dataset prep, model choice, training config, and sampling. It's based on the zhinit.dev writeup that ran the same experiment on modest hardware. If you produce music, tinker with ML, or just want a low-stakes intro to audio diffusion, this is a pretty solid weekend project. If you'd rather generate finished songs than train your own model, our Suno AI music guide covers the no-code path.
The goal: a tiny denoising diffusion probabilistic model (DDPM) that generates 1-second mono kick drum samples at 22050 Hz. The model outputs mel-spectrograms which then get inverted back to waveforms with a vocoder or Griffin-Lim.
Total parameters: around 25M. Training time on a 6GB card: 6 to 12 hours depending on dataset size. Output quality: not Splice-tier, but genuinely usable in a beat.
Yes. A small U-Net based DDPM with roughly 25M parameters fits comfortably on a 6GB card when you combine fp16 mixed precision, gradient checkpointing, and a modest batch size of 8 to 12. Training completes in 6 to 15 hours on consumer hardware.
Before you touch a single line of code, make sure you've got the following.
Hardware
Software
Dataset
And yes, more data helps. But diminishing returns kick in around 8k samples for something this simple.
Create a clean virtual environment. This matters because PyTorch and CUDA versions get finicky fast.
python3.11 -m venv kickdiff
source kickdiff/bin/activate
pip install --upgrade pip
Install PyTorch with CUDA 12.4 support:
pip install torch==2.5.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
Then grab the rest of the stack:
pip install diffusers==0.30.0 accelerate librosa soundfile numpy tqdm einops
Verify CUDA sees your GPU:
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"
If that prints True and your GPU name, you're set. If not, your driver is probably too old. The nvidia-driver-550 package from the graphics-drivers PPA fixes 90% of cases.
Kick samples come in wildly different formats: 44.1 k Hz, 48 k Hz, stereo, weird lengths. Normalize them all first.
import librosa
import soundfile as sf
from pathlib import Path
SR = 22050
TARGET_LEN = SR # 1 second
def normalize_kick(input_path, output_path):
y, _ = librosa.load(input_path, sr=SR, mono=True)
if len(y) > TARGET_LEN:
y = y[:TARGET_LEN]
else:
y = librosa.util.fix_length(y, size=TARGET_LEN)
peak = max(abs(y.max()), abs(y.min()))
if peak > 0:
y = y / peak * 0.95
sf.write(output_path, y, SR)
src = Path("raw_kicks")
dst = Path("kicks_norm")
dst.mkdir(exist_ok=True)
for f in src.glob("*.wav"):
normalize_kick(f, dst / f.name)
A few real-world gotchas to expect. Silent files sneak in from bad packs. Some samples have DC offset. And a bunch of "kicks" from generic drum packs are actually 808s with a long tail, which will confuse the model. Filter those out by hand if you can.
Next, convert everything to mel-spectrograms. A 128-bin mel at hop 256 gives you an 86-frame representation per kick, which is small enough for a 6GB card to chew through.
The temptation is to grab a giant pretrained audio diffusion model. Don't. (If you want a proper walkthrough on adapting a large model, see our LLM fine-tuning guide — the same VRAM constraints apply.) On 6GB you can't fine-tune Stable Audio or AudioLDM without offloading half the model to CPU, which drags training to a crawl.
Instead, use a small U-Net from scratch. The diffusers library ships with UNet2DModel which you can size down to fit. A config like this trains fine:
from diffusers import UNet2DModel
model = UNet2DModel(
sample_size=(128, 96),
in_channels=1,
out_channels=1,
layers_per_block=2,
block_out_channels=(64, 128, 256, 256),
down_block_types=("DownBlock2D",) * 4,
up_block_types=("UpBlock2D",) * 4,
)
That gives you roughly 24M parameters. In fp16 with gradient checkpointing, a batch size of 12 sits comfortably under 6GB. If you're on a 1660 (which lacks tensor cores entirely), fp16 still runs at roughly 2x FP32 throughput on Turing CUDA cores, so keep mixed precision on but drop to batch 8. Don't reach for TF32 here — it's an Ampere-only mode and the 1660 can't use it.
This is where most beginners burn hours. The trick is stacking three VRAM-savers together.
1. Mixed precision (fp16 or bf16)
from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp16")
2. Gradient checkpointing
model.enable_gradient_checkpointing()
This trades ~20% training speed for a ~35% VRAM reduction. On a 6GB card that's the difference between OOM and success.
3. Gradient accumulation
If you still hit OOM, drop the physical batch to 4 and accumulate 3 steps for an effective batch of 12.
accelerator = Accelerator(
mixed_precision="fp16",
gradient_accumulation_steps=3,
)
Set the DDPM scheduler with 1000 timesteps for training. You'll sample with fewer steps later.
from diffusers import DDPMScheduler
noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
Nothing fancy. Just noise, predict, backprop, repeat.
import torch
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
EPOCHS = 300
for epoch in range(EPOCHS):
for batch in dataloader:
with accelerator.accumulate(model):
clean = batch["mel"]
noise = torch.randn_like(clean)
timesteps = torch.randint(0, 1000, (clean.shape[0],), device=clean.device)
noisy = noise_scheduler.add_noise(clean, noise, timesteps)
pred = model(noisy, timesteps).sample
loss = torch.nn.functional.mse_loss(pred, noise)
accelerator.backward(loss)
optimizer.step()
optimizer.zero_grad()
if epoch % 20 == 0:
accelerator.save_state(f"ckpt/epoch_{epoch}")
On a 5,000-sample dataset with batch 12, one epoch takes about 90 seconds on a 3060 and closer to 3 minutes on a 1660. Full 300-epoch runs land between 7 and 15 hours.
Loss should drop from around 1.0 to 0.02 within the first 50 epochs. If it plateaus above 0.1 by epoch 100, your dataset probably has garbage in it. Go back to Step 2.
Load the checkpoint and run the scheduler in inference mode. For sampling, switch to DDIM with 50 steps for speed without much quality loss.
from diffusers import DDIMScheduler
scheduler = DDIMScheduler(num_train_timesteps=1000)
scheduler.set_timesteps(50)
sample = torch.randn(1, 1, 128, 96).cuda()
for t in scheduler.timesteps:
with torch.no_grad():
pred = model(sample, t).sample
sample = scheduler.step(pred, t, sample).prev_sample
Then invert the mel back to audio. Griffin-Lim is the cheap option and it sounds fine for percussive content because you don't hear phase artifacts much on transients.
mel = sample.squeeze().cpu().numpy()
audio = librosa.feature.inverse.mel_to_audio(mel, sr=22050, hop_length=256)
sf.write("generated_kick.wav", audio, 22050)
A few things that will trip you up.
Generate 100 samples. If more than 60 sound like recognizable kicks with clear transient and body, the model is trained enough. If most are muddy or noisy, either train longer or clean the dataset.
For an objective metric, compute the Fréchet Audio Distance against a held-out validation set. The official FAD toolkit works out of the box. Anything under 5.0 is respectable for this scale of model.
Once you've got a working baseline, some directions worth exploring.
sample_size to (128, 192) for 2-second kicks that include a tail.And if this weekend project bites you, congrats: you're now doing generative audio ML on hardware most people wrote off years ago.
Not comfortably. You can technically train with batch size 2 and aggressive gradient accumulation, but the 1050 Ti lacks proper fp16 acceleration so each epoch takes 6 to 8 minutes on a 5k dataset. A full 300-epoch run stretches past 30 hours. If you have any way to rent an RTX 3060 on Vast.ai for around $0.20/hour, that's a much better use of time.
ROCm on AMD RX 6000/7000 series works if you install the pytorch-rocm wheel instead of CUDA. Intel Arc requires the Intel Extension for PyTorch (IPEX). Both paths work but expect roughly 40% slower training and occasional kernel bugs. Stick with NVIDIA if you can, especially for a first project.
As of mid-2026, an RTX 3090 on Vast.ai or RunPod runs around $0.25 to $0.35 per hour. A full training run costs $3 to $5 total. Lambda Labs and Paperspace are pricier but more reliable. Google Colab's free T4 tier works for testing but will time out before a full run completes.
The model itself is yours to license however you want, but the dataset matters. If you trained on samples from a paid pack, the pack's license controls derivative use. Most free packs allow commercial use of generated derivatives, but always check the individual license. Recording your own kicks avoids the question entirely.
Both models are too large to fine-tune on 6GB even with LoRA and offloading. Stable Audio Open is around 1.2B parameters and needs at least 16GB VRAM for LoRA training. For a narrow task like one-shot kicks, a 25M-parameter model from scratch actually converges faster and produces cleaner results than a poorly fine-tuned giant.
