docs/06-stable-audio-integration.md

# 06 — Stable Audio 3 Integration

## What we're using Stable Audio 3 for

PatternTalk uses SA3 for two things, both explicitly listed in the Stability AI challenge brief:

1. **One-shot sample generation** — "give me a trashy 18-inch china"
2. **Loop preview generation** — 5–15 second previews of patterns with brutal drum styling

Plus one bonus use case:

3. **LoRA fine-tuning** — train a small adapter on brutal drum samples so the model actually understands the genre

The MIDI patterns themselves are NOT generated by SA3. They're hand-coded templates. This is a deliberate split: SA3's strength is sample quality, not structured MIDI generation.

## Why local inference, not API

The Stability AI challenge explicitly rewards *"showing the strengths of local open models."* Using the open weights locally:

- Proves we actually use the open model, not a black-box API
- Eliminates rate limits and surprise outages during the demo
- Lets us fine-tune and ship the LoRA weights as a deliverable
- Makes the inference reproducible and inspectable

We pay for this with slower inference on consumer hardware. Mitigation: cloud GPU for the heavy work.

## Hardware reality

### Vega 56 (your machine) — honest assessment

**Specs:** 8GB HBM2, GCN 5th gen, ROCm-supported but old.

**Inference (SA3 small):**
- ✅ Fits in 8GB VRAM in fp16
- ⚠️ Tight memory headroom — no concurrent inference, no batching
- ⚠️ ROCm + PyTorch + audio models = some setup friction
- ⚠️ Throughput roughly **1/4 to 1/6 of an RTX 4090** on diffusion models
- ⚠️ A 30-second sample at 100 denoising steps: expect 2–10 minutes wall time

**Fine-tuning (LoRA):**
- ❌ Not practical. Gradient buffers + optimizer state blow past 8GB.
- ⏱️ If forced: 10–50× slower than cloud. Full LoRA training: hours to days, not feasible in a hackathon.

**Verdict:** Use Vega for development and testing the inference path. Use cloud GPU for fine-tuning and demo-day inference.

### Cloud GPU (RunPod / Vast.ai / Modal)

**Recommended for the hackathon:**

| Provider | GPU | Cost/hr | Notes |
|---|---|---|---|
| RunPod | RTX 4090 | $0.40 | Easiest UX, instant deploy |
| RunPod | A4000 | $0.30 | Slightly slower than 4090 |
| Vast.ai | RTX 3090 | $0.20 | Cheaper, more setup |
| Modal | A10G | $0.50 | Serverless, pay per second |
| Lambda Labs | A100 | $1.10 | Overkill but fast |

**Recommendation:** RunPod with a 4090 for ~6 hours total = **~$3**. Fine-tune during one session, run demo inference during another.

## Setup

### Step 1: Get the model weights

```bash
# Clone the repo
git clone https://github.com/Stability-AI/stable-audio-3.git
cd stable-audio-3

# Request access on HuggingFace (gated)
# https://huggingface.co/stabilityai/stable-audio-3-small
# https://huggingface.co/stabilityai/stable-audio-3-medium

# Login
huggingface-cli login

# Download weights (use small on Vega, medium on cloud)
python scripts/download_weights.py --model small
```

### Step 2: Install dependencies (local, Vega)

ROCm is finicky. Use a pre-built PyTorch container if possible.

```bash
# Option A: Use the official Stability AI Docker image
docker pull stabilityai/stable-audio-tools:latest

# Option B: Manual install (if Docker fails)
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/rocm5.7
pip install stable-audio-tools
pip install fastapi uvicorn
```

**Reality check:** On Vega, expect 1–3 hours of yak-shaving to get inference running the first time. Budget for this.

### Step 3: Verify inference works

```python
# scripts/smoke_test.py
from stable_audio_tools import get_pretrained_model
from stable_audio_tools.inference.generation import generate_diffusion_cond

model, config = get_pretrained_model("stabilityai/stable-audio-3-small")

# Time this — establishes your baseline
import time
start = time.time()
output = generate_diffusion_cond(
    model,
    steps=100,
    cfg_scale=7,
    conditioning=[{"prompt": "d-beat drum loop, 180 BPM, brutal", "seconds_start": 0, "seconds_total": 8}],
    batch_size=1,
    sample_size=44100 * 8,
    device="cuda",
)
elapsed = time.time() - start
print(f"Generated in {elapsed:.1f}s")
# Save output
import torchaudio
torchaudio.save("smoke_test.wav", output.squeeze().cpu(), 44100)
```

If this takes < 2 minutes on Vega, you're fine. If it takes > 10 minutes, commit to cloud GPU for the demo.

## Fine-tuning: brutal-drum LoRA

### Why fine-tune

Off-the-shelf SA3 doesn't know "brutal" drums. Your testing confirmed it. Fine-tuning on a curated dataset makes the model actually understand the genre vocabulary.

### Training data curation

**Minimum viable dataset (start with this):**

- 20–40 one-shots:
  - 5–10 kicks (tight, clicky, triggered)
  - 5–10 snares (snappy, trashy, mid-range)
  - 3–5 chinas (trashy, dark, bright variants)
  - 3–5 crashes (various sizes/washes)
  - 2–3 hi-hats (closed, open, stack)
  - 2–3 rides (dry, washy, bell-forward)
- 10–20 loops (3–10 seconds each):
  - 3–5 d-beat loops at various tempos
  - 3–5 blast beat loops (traditional, hammer, hyperblast)
  - 2–3 half-time grooves
  - 2–3 punk rock / hardcore loops
  - 2–3 djent polyrhythm loops

**Sources (in order of preference):**

1. **Your own recordings** — best, no licensing issues, your taste
2. **Freesound.org** — filter to CC0 or CC-BY, document the user
3. **Splice / Loopcloud** — if you have a subscription, export with licensing documented
4. **Bandcamp / label sample packs** — check license terms, some are CC-BY
5. **Your band's existing recordings** — if you produced them, you own them

**Discipline:** Every file in `data/training/` has a corresponding entry in `data/training/manifest.yaml` with source, license, duration, BPM (for loops), and tags.

### Pre-processing

All samples normalized to:
- WAV format
- 44100 Hz sample rate (or 48000 — match SA3's expected rate)
- Mono (drums don't need stereo for LoRA training; stereo adds noise)
- Loudness normalized to ~ -14 LUFS
- Trimmed silence at start/end
- Loops: aligned to bar boundaries, ideally with a single downbeat

```python
# services/audio/training/preprocess.py
import torchaudio
import pyloudnorm as pyln

def preprocess(input_path: str, output_path: str):
    waveform, sr = torchaudio.load(input_path)
    # Resample
    if sr != 44100:
        waveform = torchaudio.functional.resample(waveform, sr, 44100)
    # To mono
    waveform = waveform.mean(dim=0, keepdim=True)
    # Loudness normalize
    meter = pyln.Meter(44100)
    loudness = meter.integrated_loudness(waveform.numpy().T)
    normalized = pyln.normalize.loudness(waveform.numpy().T, loudness, -14.0)
    waveform = torch.from_numpy(normalized.T).unsqueeze(0)
    torchaudio.save(output_path, waveform, 44100)
```

### LoRA training script

Stability's repo supports LoRA via `stable_audio_tools.training.lora`. Use their example as a base.

```python
# services/audio/training/train_lora.py
import torch
from stable_audio_tools.models import create_model_from_config
from stable_audio_tools.training.lora import LoRADataset, LoRATrainer
from stable_audio_tools.data.utils import load_training_manifest

# Load base model
model, config = create_model_from_config("model_config.json")
model.load_state_dict(torch.load("stable_audio_3_small.safetensors"))

# Load training manifest
dataset = LoRADataset(
    manifest_path="data/training/manifest.yaml",
    audio_dir="data/training/",
)

# LoRA config
lora_config = {
    "rank": 32,
    "alpha": 32,
    "dropout": 0.05,
    "target_modules": ["to_q", "to_k", "to_v", "to_out.0"],  # SA3 specifics
}

# Train
trainer = LoRATrainer(
    model=model,
    dataset=dataset,
    lora_config=lora_config,
    learning_rate=1e-4,
    batch_size=2,
    gradient_accumulation=4,
    max_steps=1500,
    save_every=500,
    output_dir="loras/brutal-drums",
    device="cuda",
)

trainer.train()
```

**Expected training time on RTX 4090:** 30–90 minutes for 1500 steps with rank 32.

**Output:**
- `loras/brutal-drums/adapter.safetensors` — the LoRA weights (~50MB)
- `loras/brutal-drums/checkpoints/` — intermediate checkpoints
- `loras/brutal-drums/training_log.json` — loss curve

### Publishing the LoRA

After training, push to HuggingFace:

```python
# services/audio/training/publish.py
from huggingface_hub import HfApi

api = HfApi()
api.create_repo("your-username/patterntalk-brutal-drums", repo_type="model")

api.upload_folder(
    folder_path="loras/brutal-drums/",
    repo_id="your-username/patterntalk-brutal-drums",
    commit_message="Initial LoRA trained on brutal drum samples",
)

# Upload model card
api.upload_file(
    path_or_fileobj="loras/brutal-drums/README.md",
    path_in_repo="README.md",
    repo_id="your-username/patterntalk-brutal-drums",
)
```

**Model card must include:**
- Base model reference
- Training data summary (with manifest link)
- License
- Intended use
- Limitations (it's a LoRA, not general-purpose)
- Citation to Stability AI and PatternTalk

## Inference service

### Architecture

```python
# services/audio/sa3/server.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from stable_audio_tools import get_pretrained_model
from stable_audio_tools.inference.generation import generate_diffusion_cond
import torchaudio
import io
import hashlib

app = FastAPI()

# Load model once at startup
model, config = get_pretrained_model("stabilityai/stable-audio-3-small")
model = model.to("cuda")
model.eval()

# Load LoRA adapter
model.load_adapter("loras/brutal-drums/adapter.safetensors", adapter_name="brutal")

# Sample cache (LRU + disk)
cache_dir = Path("cache/")
cache_dir.mkdir(exist_ok=True)

class GenerateRequest(BaseModel):
    prompt: str
    duration_seconds: float = 8.0
    intensity: str = "medium"  # soft | medium | brutal | brutal-max
    use_lora: bool = True
    cfg_scale: float = 7.0
    steps: int = 100

class GenerateResponse(BaseModel):
    audio_url: str  # URL to download the WAV
    prompt: str
    duration: float
    cached: bool

@app.post("/generate", response_model=GenerateResponse)
async def generate(req: GenerateRequest):
    # Cache key
    cache_key = hashlib.sha256(
        f"{req.prompt}|{req.duration_seconds}|{req.intensity}|{req.use_lora}".encode()
    ).hexdigest()[:16]

    cache_path = cache_dir / f"{cache_key}.wav"

    if cache_path.exists():
        return GenerateResponse(
            audio_url=f"/cache/{cache_key}.wav",
            prompt=req.prompt,
            duration=req.duration_seconds,
            cached=True,
        )

    # Build conditioning
    conditioning = [{
        "prompt": req.prompt,
        "seconds_start": 0,
        "seconds_total": req.duration_seconds,
    }]

    # Apply LoRA
    if req.use_lora:
        model.set_adapter("brutal")
    else:
        model.disable_adapters()

    # Generate
    try:
        output = generate_diffusion_cond(
            model,
            steps=req.steps,
            cfg_scale=req.cfg_scale,
            conditioning=conditioning,
            sample_size=int(44100 * req.duration_seconds),
            device="cuda",
        )
    except Exception as e:
        raise HTTPException(500, f"Generation failed: {e}")

    # Save
    torchaudio.save(str(cache_path), output.squeeze().cpu(), 44100)

    return GenerateResponse(
        audio_url=f"/cache/{cache_key}.wav",
        prompt=req.prompt,
        duration=req.duration_seconds,
        cached=False,
    )

@app.get("/cache/{key}.wav")
async def get_cached(key: str):
    path = cache_dir / f"{key}.wav"
    if not path.exists():
        raise HTTPException(404)
    return FileResponse(path)
```

### Intensity → prompt engineering

Map user-friendly intensity to model-friendly prompts:

```python
INTENSITY_PROMPTS = {
    "soft": "clean studio drums, polished, tight, controlled",
    "medium": "live room drums, punchy, present",
    "brutal": "brutal drums, trashy, aggressive, raw, distorted",
    "brutal-max": "ultra-brutal drums, completely destroyed, blown-out, panic-attack intensity",
}

def build_audio_prompt(req: AudioPrompt) -> str:
    parts = [
        req.pattern,
        req.styleHints.join(", "),
        INTENSITY_PROMPTS[req.intensity],
        "drum recording, close-mic'd",
    ]
    if req.limb:
        parts.append(f"single {req.limb} hit")
    return ", ".join(filter(None, parts))
```

## Prompting strategy for SA3

Stable Audio 3 responds well to:
- Genre terms ("d-beat", "black metal", "brutal")
- Tempo ("180 BPM")
- Recording style ("close-mic'd", "room mic", "triggered")
- Specific instrument descriptors ("tight snare", "trashy china")
- Intensity adjectives ("brutal", "raw", "polished")

SA3 responds poorly to:
- Vague aesthetic terms ("cool", "interesting")
- Mixing many genres ("jazz-black-metal-funk")
- Trying to specify timing ("snare on the and of 2")

Our brutal-drum LoRA biases the model toward:
- Recognizing metal subgenres
- Generating physically realistic drum sounds
- Avoiding pop/EDM patterns

## Demo strategy

For the live demo, we want generation to be fast and reliable. Pre-generate a few showcase samples:

```bash
# Pre-generate demo cache
python scripts/pregenerate_demo.py
# Generates:
#   cache/dbeat-180.wav
#   cache/blast-traditional-200.wav
#   cache/china-trashy.wav
#   cache/kick-brutal.wav
#   cache/skank-120.wav
#   cache/variation-1.wav ... variation-4.wav
```

During the demo, hit the live API for the "regenerate" command (to show inference happening), but rely on cache for the initial generation to keep the demo snappy.

## What can go wrong

| Failure | Mitigation |
|---|---|
| Cloud GPU instance dies | Pre-generated cache + fallback to Vega inference |
| SA3 produces generic output | LoRA adapter handles this; fallback is to label output "demo variation" |
| Generation exceeds 30s | Show progress indicator + have pre-generated fallback |
| Audio quality is bad | Iterate on LoRA training data; have multiple variations to pick from |
| ROCm issues on Vega | Skip local dev, use cloud-only |

## Cost summary

| Item | Cost |
|---|---|
| Cloud GPU (RunPod 4090, 6 hours total) | ~$3 |
| HuggingFace Pro (if needed for gated weights) | $0 (free tier sufficient) |
| Total cloud spend | **~$3** |

If your budget is truly $0, do everything on Vega and accept 2–10 minute inference times. The product still works, the demo just needs more buffer time.

## Post-hackathon

- Quantize the model (int8) for faster inference
- Train additional LoRAs: rock, jazz, funk (one per genre)
- Host inference on Modal/Replicate for public use
- Build a "train your own LoRA" UI — the "personal LoRA trainer" use case from the brief