jam-buddy / docs /06-stable-audio-integration.md
salgadev's picture
Sync from GitHub 6feaf31d
b2e4883 verified
|
Raw
History Blame Contribute Delete
14.7 kB

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:

  1. 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

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

# 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

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

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

# 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

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

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:

# 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