docs/modelingProcess.md

# Modeling Process

## Overview

This project trains a **latent diffusion model** that generates kick drum audio samples. It uses a three-component architecture:

1. **VAE (Variational Autoencoder)** -- compresses mel spectrograms into a compact latent space
2. **Diffusion Model (U-Net)** -- generates new latents from noise, with optional text conditioning
3. **Vocoder (HiFi-GAN)** -- converts mel spectrograms back to audio waveforms

The diffusion process operates in the VAE's latent space rather than on raw audio, which makes both training and inference much faster while maintaining quality.

## Audio Parameters

These constants are shared across all components:

| Parameter | Value |
|-----------|-------|
| Sample rate | 44,100 Hz |
| Duration | 2 seconds (88,200 samples) |
| n_fft | 2048 |
| hop_length | 512 |
| n_mels | 128 |
| Mel shape | (1, 128, 173) |
| Latent shape | (4, 8, 11) -- 352 floats, ~63x compression |

## Data Pipeline

### Step 1: Aggregate samples (`scripts/aggregate_data.py`)

Collects kick drum samples from a source directory of nested folders into a single flat `data/raw/` folder. Filters files by:
- Filename contains "kick"
- Not a loop or BPM-tagged file
- File size between 5KB and 1MB

Each file is renamed with a hash prefix to avoid collisions (e.g. `a3f2e1_KICK3.WAV`). A `data/metadata.csv` is generated with columns: `filename`, `original_path`, `keywords`.

Keywords are extracted from filenames by splitting on separators (`-`, `_`, spaces, dots), lowercasing, and removing pure numbers and tokens shorter than 2 characters.

**Result:** ~13,600 kick samples in `data/raw/`, with `data/metadata.csv`.

### Step 2: Preprocess audio (`scripts/preprocess.py`)

Converts each raw audio file to a log-mel spectrogram tensor:

1. Load audio via `soundfile`
2. Convert to mono
3. Resample to 44,100 Hz
4. Pad or trim to exactly 2 seconds
5. Normalize to -1 dB peak
6. Apply 0.2s linear fade-out
7. Compute mel spectrogram via `torchaudio.transforms.MelSpectrogram`
8. Convert to log scale: `log(mel.clamp(min=1e-5))`
9. Save as `.pt` tensor in `data/processed/`

**Result:** 13,613 tensors of shape `(1, 128, 173)` in `data/processed/`.

## VAE Training

### Architecture

**Encoder:** 4 downsampling stages using strided Conv2d + ResBlocks (GroupNorm + SiLU):
```
(1, 128, 173) -> (32, 64, 87) -> (64, 32, 44) -> (128, 16, 22) -> (256, 8, 11)
```
Then 1x1 convolutions produce `mu` and `logvar`, each of shape `(4, 8, 11)`.

**Decoder:** Mirror of encoder using ConvTranspose2d, outputs are cropped to `(1, 128, 173)`.

**Reparameterization:** `z = mu + std * epsilon` during training; inference uses `mu` directly (no sampling).

### Loss

```
loss = MSE(recon, target) + L1(recon, target) + kl_weight * KL_divergence
```

KL weight anneals linearly from 0.0001 to 0.001 over the first 20 epochs to prevent posterior collapse.

### Training Config

| Parameter | Value |
|-----------|-------|
| Batch size | 32 |
| Learning rate | 1e-4 (AdamW) |
| Epochs | 100 |
| Mixed precision | Yes (fp16 on CUDA) |
| Checkpoint every | 10 epochs |

**Checkpoint:** `checkpoints/vae_epoch_100.pt` (contains model state, optimizer state, config).

### Pre-encoding latents

Before diffusion training, all mel spectrograms are encoded to latents using the frozen VAE encoder (`mu` only, no sampling). Saved as `.pt` files in `data/latents/`, shape `(4, 8, 11)` each.

## Diffusion Model Training

### Architecture

**LatentUNet** operates on `(4, 8, 11)` latent tensors:

- **Timestep embedding:** Sinusoidal positional encoding (dim=64) -> MLP -> (256,)
- **Text embedding:** Keyword embeddings averaged -> linear projection -> (256,)
- **Conditioning:** timestep + text embeddings added together

**Down path:**
```
Conv2d(4, 64) -> CondResBlock(64)              -> (64, 8, 11)
Conv2d(64, 128, stride=2) -> CondResBlock(128) -> (128, 4, 6)
Conv2d(128, 256, stride=2)                     -> (256, 2, 3)
```

**Middle:** CondResBlock(256) -> SelfAttention2d(256) -> CondResBlock(256)

**Up path** with skip connections (concat + 1x1 reduce):
```
ConvTranspose2d + skip + CondResBlock -> (128, 4, 6)
ConvTranspose2d + skip + CondResBlock -> (64, 8, 11)
```

**Output:** GroupNorm + SiLU + Conv2d -> `(4, 8, 11)` predicted noise.

### Text Encoder

`KeywordEncoder` maps keyword token IDs to a conditioning vector:
- Embedding layer: `(vocab_size, 64)`
- Average keyword embeddings -> Linear projection to `(256,)`
- Empty keywords use a learnable `null_embedding` (enables unconditional generation)

Vocabulary is built from `metadata.csv` keywords with `min_count=5`.

### Noise Schedule

Linear beta schedule from 0.0001 to 0.02 over 1000 timesteps.

### Classifier-Free Guidance (CFG)

During training, 15% of samples have their keywords dropped (replaced with empty list), training the model to handle both conditional and unconditional generation. At inference, CFG interpolates between conditional and unconditional predictions.

### Training Config

| Parameter | Value |
|-----------|-------|
| Batch size | 16 |
| Learning rate | 1e-4 (AdamW) |
| Iterations | 100,000 |
| Gradient accumulation | 2 |
| Mixed precision | Yes (fp16 on CUDA) |
| EMA decay | 0.9999 |
| CFG dropout | 0.15 |
| Checkpoint every | 5,000 iterations |

**Checkpoint:** `checkpoints/diffusion_step_100000.pt` (contains model state, EMA state, text encoder state, vocab list, config).

## Vocoder Training

### Architecture

**HiFi-GAN v1 Generator** (sized for 6GB VRAM):
```
Input: (128, T) mel spectrogram
-> Conv1d(128, 256, 7)
-> Upsample 8x + 3x ResBlock1  -> (128, T*8)
-> Upsample 8x + 3x ResBlock1  -> (64, T*64)
-> Upsample 2x + 3x ResBlock1  -> (32, T*128)
-> Upsample 2x + 3x ResBlock1  -> (16, T*256)
-> Upsample 2x + 3x ResBlock1  -> (8, T*512)
-> Conv1d(8, 1, 7) + tanh
Output: (1, T*512) waveform
```

Total upsample factor: 8 * 8 * 2 * 2 * 2 = 512 (matches hop_length).

**Discriminators:**
- Multi-Period Discriminator (periods: 2, 3, 5, 7, 11)
- Multi-Scale Discriminator (3 scales with average pooling)

### Loss

```
loss_g = adversarial_loss + 2.0 * feature_matching_loss + 45.0 * mel_reconstruction_loss
loss_d = LS-GAN discriminator loss (real vs fake)
```

### Training Config

| Parameter | Value |
|-----------|-------|
| Batch size | 4-8 |
| Learning rate | 2e-4 (AdamW, betas=0.8/0.99) |
| Epochs | 50 |
| Segment size | 8,192 samples (random crop per batch item) |
| Mixed precision | No (fp32 for GAN stability) |
| LR scheduler | ExponentialLR(gamma=0.999) |
| Checkpoint every | 10 epochs |

The vocoder dataset lazy-loads raw audio paired with pre-computed mel spectrograms. Only file paths are stored in memory.

**Checkpoint:** `checkpoints/vocoder_epoch_50.pt` (contains generator, discriminator, optimizer, and scheduler states).

## Training Hardware

- **Training:** Linux desktop with 6GB VRAM GPU
- **Inference:** CPU-only (designed for Railway / Django deployment)

## Dependencies

```
torch>=2.10.0
torchaudio>=2.10.0
numpy>=2.4.2
pandas>=3.0.0
soundfile>=0.13.1
scipy>=1.14.0
tensorboard>=2.20.0
tqdm>=4.67.2
```