Setting Up SonicMaster on a 6GB GPU: Pitfalls, Tricks, and Hard-Won Lessons

SonicMaster (paper, GitHub) is a diffusion-based model from AMAAI Lab at SUTD that handles denoising, de-clipping, de-reverberation, and bandwidth extension — all guided by text prompts. You feed it a degraded audio file and a prompt like “enhance audio quality, reduce noise, add clarity and brightness”, and it produces a restored version.

Sounds amazing on paper. Getting it running on consumer hardware? That’s where the fun begins. The repo assumes an A100 with 40GB+ VRAM. Here’s every pitfall I hit on my RTX 2060 6GB running Arch Linux — and the workarounds that actually worked.

Pitfall #1: Python 3.13 Requirement (But You’re on 3.14)

The README says python==3.13. My Arch system ships Python 3.14. The requirements pin torch==2.4.0, diffusers==0.30.0, and friends — all from mid-2024. They don’t build cleanly on 3.14.

Fix: Create a separate venv with Python 3.13:

cd /mnt/ssd/home/rio/documents/projects/SonicMaster
uv venv --python 3.13 venv
source venv/bin/activate
pip install -r requirements_sonic.txt

If uv isn’t your thing, pyenv install 3.13 works too. The key point: don’t fight the pinned versions on a newer Python. Just use 3.13. I also tested a venv311 directory with Python 3.11 — diffusers==0.30.0 is more stable there. Either works.

Pitfall #2: The VAE Download Fails Silently

SonicMaster depends on the VAE from stabilityai/stable-audio-open-1.0 (subfolder vae). The from_pretrained() call fails with:

OSError: stabilityai/stable-audio-open-1.0 does not appear to have a file...

Two reasons:

1. The repo has a gated license agreement — you need to accept it on HuggingFace before downloading.
2. Even after accepting, you need HF_TOKEN for authenticated access.

Fix:

1. Go to https://huggingface.co/stabilityai/stable-audio-open-1.0 and accept the license.
2. Create a HuggingFace access token at https://huggingface.co/settings/tokens.
3. Download the VAE manually:

export HF_HOME=/path/to/SonicMaster/models
export HF_TOKEN=hf_xxxxx

python -c "
from huggingface_hub import snapshot_download
snapshot_download(
    'stabilityai/stable-audio-open-1.0',
    local_dir='/path/to/SonicMaster/models/stable-audio-vae'
)
"

The flan-t5-large text encoder also auto-downloads from HF Hub. Set HF_HOME to keep it in your project directory — especially useful on an SSD where space matters.

The inference scripts read hf_token from environment variables (HF_TOKEN, HUGGINGFACE_TOKEN, or HUGGINGFACEHUB_API_TOKEN), so make sure one of those is set.

Pitfall #3: 3.2GB Model + VAE = OOM on 6GB Cards

The model checkpoint (model.safetensors) is 3.2GB. The VAE adds ~700MB. The text encoder (flan-t5-large) is another ~1.3GB. That’s 5.2GB before you’ve even started inference — on a 6GB card, that’s instant OOM.

The original infer_single.py loads everything on GPU at once with vae_batch_size=10. Game over.

Fix: I wrote inference_lowmem.py and inference_lowmem2.py with three key tricks:

Trick 1: Model in float16, Text Encoder on CPU

model.half().to(device).eval()

# Keep text encoder on CPU in float32 for stability
model.text_encoder = model.text_encoder.cpu().float()
for p in model.text_encoder.parameters():
    p.requires_grad = False

The model runs fine in float16. The text encoder only runs once per chunk to encode the prompt — it doesn’t need GPU. Keep it in float32 on CPU for quality.

Trick 2: VAE Shuttling (inference_lowmem2.py)

Instead of keeping the VAE on GPU permanently, load it on CPU and move to GPU only for encode/decode ops:

# VAE starts on CPU
vae = AutoencoderOobleck.from_pretrained(...).cpu().float()

# Before encoding/decoding, move to GPU
vae.to(device).half()
z = vae.encode(batch).latent_dist.mode()
# Immediately move back to CPU
vae.cpu().float()
torch.cuda.empty_cache()

This frees ~700MB of VRAM between VAE operations. The shuttle cost is negligible — it’s just parameter copying, not computation.

Trick 3: Sequential Chunk Processing with VRAM Cleanup

for i in range(degraded_latents.shape[0]):
    # Move text encoder to GPU just for this chunk
    model.text_encoder = model.text_encoder.to(device).half()
    
    result_latent = model.inference_flow(...)
    
    # Move text encoder back to CPU immediately
    model.text_encoder = model.text_encoder.cpu().float()
    del result_latent
    torch.cuda.empty_cache()

Each chunk processes one at a time, with aggressive cleanup after each.

Results

inference_lowmem2.py is what run.sh uses — it’s the safest option for 6GB cards.

Pitfall #4: Default Overlap is Too Aggressive

The original infer_single.py defaults to overlap_duration=10 seconds with chunk_duration=30. On 6GB, the carry-over conditioning latent is large.

Better defaults for low VRAM:

--chunk_duration 30
--overlap_duration 5   (reduced from 10)
--vae_batch_size 2     (reduced from 10)

Smaller overlap means less carry context but significantly less VRAM pressure.

Pitfall #5: The externally-managed-environment Error on Arch

If you try pip install -r requirements_sonic.txt on Arch without a venv:

error: externally-managed-environment
× This environment is externally managed

Fix: Always use a venv. The uv venv approach above handles this. Don’t use --break-system-packages.

Pitfall #6: Hardcoded Training Paths in Config

The tangoflux_config.yaml has hardcoded absolute paths for training data:

paths:
  train_file: "/mastering/FINAL_DATA2/trainset_pt.jsonl"
  val_file: "/mastering/FINAL_DATA2/valset_pt.jsonl"

These are irrelevant for inference — the model only needs the model section of the config. Don’t bother editing them.

Full Setup Walkthrough

# 1. Clone the repo
git clone https://github.com/AMAAI-Lab/SonicMaster.git
cd SonicMaster

# 2. Create venv with Python 3.13
uv venv --python 3.13 venv
source venv/bin/activate

# 3. Install requirements
pip install -r requirements_sonic.txt

# 4. Set environment variables
export HF_HOME=$(pwd)/models
export HF_TOKEN=hf_xxxxx

# 5. Download the model checkpoint (~3.2GB)
python -c "
from huggingface_hub import snapshot_download
snapshot_download('amaai-lab/SonicMaster', local_dir='models/sonicmaster-checkpoint')
"

# 6. Pre-download the VAE and text encoder
python -c "
from diffusers import AutoencoderOobleck
from transformers import AutoModelForSeq2SeqLM

# VAE
vae = AutoencoderOobleck.from_pretrained(
    'stabilityai/stable-audio-open-1.0', subfolder='vae',
    use_auth_token='hf_xxxxx'
)
vae.save_pretrained('models/stable-audio-vae')

# Text encoder
te = AutoModelForSeq2SeqLM.from_pretrained('google/flan-t5-large')
te.save_pretrained('models/flan-t5-large')
"

# 7. Run inference (low-VRAM mode)
python inference_lowmem2.py \
    --ckpt models/sonicmaster-checkpoint/model.safetensors \
    --input your_audio.wav \
    --prompt "enhance audio quality, reduce noise, add clarity and brightness" \
    --output enhanced.wav \
    --config configs/tangoflux_config.yaml \
    --chunk_duration 30 \
    --overlap_duration 5 \
    --num_inference_steps 10 \
    --guidance_scale 1.0 \
    --seed 0

Or use the run.sh wrapper:

INPUT=your_audio.wav PROMPT="enhance audio quality, reduce noise" ./run.sh

Tips

1. Set HF_HOME to your project’s models/ directory — keeps HuggingFace caches on your SSD and avoids filling up root.

2. Delete optimizer.bin — it’s 3.9GB and only needed for training, not inference:

   rm models/sonicmaster-checkpoint/optimizer.bin

3. Guidance scale of 1.0 works best. Higher values cause artifacts. Be descriptive in prompts: “reduce background noise, enhance vocals, add warmth and presence”.

4. Crossfade stitching between chunks is handled automatically in the lowmem scripts. The overlap creates smooth transitions.

5. Monitor VRAM during inference — the scripts print VRAM used at each step. If you see it creeping past 5.5GB, reduce vae_batch_size to 1.

6. First run is slow because it downloads flan-t5-large (~1.3GB) and the VAE (~700MB). After that, they’re cached in HF_HOME.

The Honest Take

SonicMaster is impressive research software that produces genuinely good audio restoration. But it’s exactly that — research software. The repo assumes an A100 with 40GB+ VRAM. Making it work on 6GB required writing custom inference scripts, shuttling models between CPU and GPU, and accepting slower processing.

If you’ve got 8GB+ VRAM, the original infer_single.py probably works fine. If you’re on 6GB like me, inference_lowmem2.py is your friend. And if you have no GPU at all, inference_cpu.py will get you there — just pack a snack while you wait.

The output quality, though? Worth it. The text-prompt conditioning gives you real control over the restoration style, which sets it apart from pure enhancement models.