Add BigVGAN vocoder with CUDA support

- Auto-detects GPU at startup, loads BigVGAN from HuggingFace Hub
- Falls back to Griffin-LIM on CPU
- Loads fine-tuned BigVGAN weights from models/vocoder/ if present

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Dockerfile

@@ -17,11 +17,10 @@ RUN useradd -m -u 1000 user
 # Install uv for faster pip installs
 RUN pip install --no-cache-dir uv
 
-# Install CPU-only PyTorch (avoids ~3GB of CUDA bloat)
+# Install PyTorch with CUDA support (GPU Space)
 RUN uv pip install --system --no-cache-dir \
     torch \
     torchaudio \
-    --index-url https://download.pytorch.org/whl/cpu \
     && uv cache clean
 
 # Install app dependencies

main.py

@@ -1,9 +1,10 @@
 """Kicks API — VAE kick drum synthesizer deployed on Hugging Face Spaces.
 
-Uses CPU-only PyTorch with Griffin-LIM vocoder (no GPU needed).
+Uses BigVGAN vocoder on GPU (falls back to Griffin-LIM on CPU).
 PCA parameters are pre-computed and loaded from pca_params.npz.
 """
 
+import glob
 import io
 import os
 
@@ -34,6 +35,9 @@ LOG_MEL_MAX = 2.5         # headroom above observed max
 N_PCS = 5
 LATENT_DIM = 128
 
+VOCODER_DIR = "models/vocoder"  # fine-tuned BigVGAN weights
+BIGVGAN_MODEL = "nvidia/bigvgan_v2_44khz_128band_256x"
+
 
 # ──────────────────────────────────────────────
 # VAE model (identical architecture to kicks/model.py)
@@ -45,7 +49,6 @@ class VAE(nn.Module):
         super().__init__()
         self.latent_dim = latent_dim
 
-        # Encoder: (B, 1, 128, 256) → (B, 256, 8, 16)
         self.encoder = nn.Sequential(
             nn.Conv2d(1, 32, 3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
             nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
@@ -54,8 +57,6 @@ class VAE(nn.Module):
         )
         self.fc_mu = nn.Linear(256 * 8 * 16, latent_dim)
         self.fc_logvar = nn.Linear(256 * 8 * 16, latent_dim)
-
-        # Decoder: z → (B, 1, 128, 256)
         self.fc_decode = nn.Linear(latent_dim, 256 * 8 * 16)
         self.decoder = nn.Sequential(
             nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1),
@@ -75,10 +76,44 @@ class VAE(nn.Module):
 
 
 # ──────────────────────────────────────────────
-# Griffin-LIM vocoder (runs on CPU, no GPU needed)
+# BigVGAN vocoder
+# ──────────────────────────────────────────────
+def load_bigvgan(device: torch.device):
+    """Load BigVGAN from HF Hub, optionally with fine-tuned weights."""
+    import bigvgan as _bigvgan
+
+    # Patch for huggingface_hub >= 1.0 compatibility
+    original = _bigvgan.BigVGAN._from_pretrained.__func__
+
+    @classmethod
+    def _patched(cls, **kwargs):
+        kwargs.setdefault("proxies", None)
+        kwargs.setdefault("resume_download", False)
+        return original(cls, **kwargs)
+
+    _bigvgan.BigVGAN._from_pretrained = _patched
+
+    model = _bigvgan.BigVGAN.from_pretrained(BIGVGAN_MODEL, use_cuda_kernel=False)
+
+    # Load fine-tuned weights if available
+    pth_files = sorted(glob.glob(os.path.join(VOCODER_DIR, "*.pth")))
+    if pth_files:
+        pth_path = pth_files[0]
+        checkpoint = torch.load(pth_path, map_location=device, weights_only=False)
+        model.load_state_dict(checkpoint["generator"])
+        print(f"Loaded fine-tuned BigVGAN from {pth_path}")
+    else:
+        print("Using pretrained BigVGAN from HuggingFace Hub")
+
+    model.remove_weight_norm()
+    return model.eval().to(device)
+
+
+# ──────────────────────────────────────────────
+# Griffin-LIM fallback vocoder
 # ──────────────────────────────────────────────
 class GriffinLimVocoder:
-    """Mel spectrogram → audio via pseudo-inverse mel filterbank + Griffin-LIM."""
+    """Mel spectrogram → audio via pseudo-inverse + Griffin-LIM."""
 
     def __init__(self, n_iter: int = 64):
         n_stft = N_FFT // 2 + 1
@@ -89,22 +124,16 @@ class GriffinLimVocoder:
             n_mels=N_MELS,
             sample_rate=SAMPLE_RATE,
         )
-        # fb shape: (n_stft, n_mels) → fb.T shape: (n_mels, n_stft)
-        # pinv(fb.T) shape: (n_stft, n_mels)
         self.fb_pinv = torch.linalg.pinv(fb.T)
         self.griffin_lim = torchaudio.transforms.GriffinLim(
-            n_fft=N_FFT,
-            win_length=WIN_SIZE,
-            hop_length=HOP_LENGTH,
-            n_iter=n_iter,
+            n_fft=N_FFT, win_length=WIN_SIZE,
+            hop_length=HOP_LENGTH, n_iter=n_iter,
         )
 
     def __call__(self, log_mel: torch.Tensor) -> torch.Tensor:
-        """Input: (B, n_mels, T). Output: (B, T)."""
         mel_linear = torch.exp(log_mel).cpu()
-        # Linear spectrogram via pseudo-inverse
         linear_spec = torch.clamp(self.fb_pinv @ mel_linear, min=0.0)
-        linear_spec = linear_spec + 1e-4  # noise floor to avoid phase artifacts
+        linear_spec = linear_spec + 1e-4
         return self.griffin_lim(linear_spec)
 
 
@@ -122,9 +151,10 @@ def denormalize(spec: torch.Tensor) -> torch.Tensor:
 class AppState:
     device: torch.device
     model: VAE
-    vocoder: GriffinLimVocoder
-    pca_components: np.ndarray  # shape (N_PCS, LATENT_DIM)
-    pca_mean: np.ndarray        # shape (LATENT_DIM,)
+    vocoder: object
+    vocoder_type: str
+    pca_components: np.ndarray
+    pca_mean: np.ndarray
     pc_names: list[str]
     pc_mins: list[float]
     pc_maxs: list[float]
@@ -147,13 +177,19 @@ app.add_middleware(
 
 @app.on_event("startup")
 async def startup():
-    state.device = torch.device("cpu")
-
-    # Load PCA parameters pre-computed from the training dataset
+    # Auto-detect device
+    if torch.cuda.is_available():
+        state.device = torch.device("cuda")
+        print("Using CUDA GPU")
+    else:
+        state.device = torch.device("cpu")
+        print("CUDA not available — falling back to CPU")
+
+    # Load PCA parameters
     pca_path = os.environ.get("PCA_PARAMS_PATH", "pca_params.npz")
     pca_data = np.load(pca_path, allow_pickle=True)
-    state.pca_components = pca_data["components"]       # (N_PCS, LATENT_DIM)
-    state.pca_mean = pca_data["mean"]                    # (LATENT_DIM,)
+    state.pca_components = pca_data["components"]
+    state.pca_mean = pca_data["mean"]
     state.pc_names = list(pca_data["pc_names"])
     state.pc_mins = list(pca_data["pc_mins"])
     state.pc_maxs = list(pca_data["pc_maxs"])
@@ -161,21 +197,32 @@ async def startup():
 
     # Load VAE checkpoint
     ckpt_path = os.environ.get("VAE_CHECKPOINT_PATH", "vae_best.pth")
-    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    checkpoint = torch.load(ckpt_path, map_location=state.device, weights_only=False)
     latent_dim = checkpoint.get("latent_dim", LATENT_DIM)
     state.model = VAE(latent_dim=latent_dim)
     state.model.load_state_dict(checkpoint["model"])
+    state.model.to(state.device)
     state.model.eval()
     print(f"VAE loaded (latent_dim={latent_dim})")
 
-    # Initialize Griffin-LIM vocoder
-    state.vocoder = GriffinLimVocoder()
-    print("Griffin-LIM vocoder ready")
+    # Load vocoder — try BigVGAN with CUDA, fall back to Griffin-LIM
+    if state.device.type == "cuda":
+        try:
+            state.vocoder = load_bigvgan(state.device)
+            state.vocoder_type = "bigvgan"
+            print("BigVGAN vocoder ready")
+        except Exception as e:
+            print(f"BigVGAN loading failed ({e}), falling back to Griffin-LIM")
+            state.vocoder = GriffinLimVocoder()
+            state.vocoder_type = "griffinlim"
+    else:
+        state.vocoder = GriffinLimVocoder()
+        state.vocoder_type = "griffinlim"
+
     print("API ready")
 
 
 def _parse_pc_values(request: Request) -> list[float]:
-    """Map slider positions [0, 1] to PC-space values."""
     values = []
     for i in range(N_PCS):
         raw = float(request.query_params.get(f"pc{i + 1}", "0.5"))
@@ -185,17 +232,26 @@ def _parse_pc_values(request: Request) -> list[float]:
 
 
 def _pc_to_latent(pc_values: list[float]) -> torch.Tensor:
-    """Transform 5D PC slider values → 128D latent vector via PCA inverse."""
     z_np = (np.array(pc_values) @ state.pca_components) + state.pca_mean
-    return torch.tensor(z_np, dtype=torch.float32).unsqueeze(0)
+    return torch.tensor(z_np, dtype=torch.float32).to(state.device).unsqueeze(0)
 
 
-def _decode_to_audio(spec: torch.Tensor) -> np.ndarray:
-    """Decode normalized spectrogram → waveform via Griffin-LIM."""
+def _spec_to_audio(spec: torch.Tensor) -> np.ndarray:
+    """Convert normalized spectrogram to waveform using the loaded vocoder."""
     log_mel = denormalize(spec)
     log_mel = log_mel.squeeze(1)  # (1, 128, 256)
-    waveform = state.vocoder(log_mel)  # (1, T)
-    waveform = waveform / (waveform.abs().max() + 1e-8)
+
+    if state.vocoder_type == "bigvgan":
+        with torch.no_grad():
+            waveform = state.vocoder(log_mel.to(state.device))  # (1, 1, T)
+        waveform = waveform.squeeze(1).cpu()  # (1, T)
+        waveform = torchaudio.functional.highpass_biquad(waveform, SAMPLE_RATE, cutoff_freq=25.0)
+        waveform = torchaudio.functional.lowpass_biquad(waveform, SAMPLE_RATE, cutoff_freq=20000.0)
+        waveform = waveform / (waveform.abs().max() + 1e-8)
+    else:
+        waveform = state.vocoder(log_mel)  # (1, T)
+        waveform = waveform / (waveform.abs().max() + 1e-8)
+
     return waveform.squeeze(0).numpy()
 
 
@@ -210,14 +266,12 @@ async def get_config():
             {
                 "id": i + 1,
                 "name": state.pc_names[i],
-                "min": 0.0,
-                "max": 1.0,
-                "default": 0.5,
-                "step": 0.01,
+                "min": 0.0, "max": 1.0,
+                "default": 0.5, "step": 0.01,
             }
             for i in range(N_PCS)
         ],
-        "vocoder": "griffinlim",
+        "vocoder": state.vocoder_type,
     }
 
 
@@ -227,8 +281,8 @@ async def generate(request: Request):
     z = _pc_to_latent(pc_values)
 
     with torch.no_grad():
-        spec = state.model.decode(z)  # (1, 1, 128, 256)
-        audio = _decode_to_audio(spec)
+        spec = state.model.decode(z)
+        audio = _spec_to_audio(spec)
 
     buf = io.BytesIO()
     sf.write(buf, audio, SAMPLE_RATE, format="WAV")
@@ -250,4 +304,4 @@ async def spectrogram_data(request: Request):
 
 @app.get("/health")
 async def health():
-    return {"status": "ok"}
+    return {"status": "ok", "vocoder": state.vocoder_type, "device": str(state.device)}

requirements.txt

@@ -3,3 +3,4 @@ scikit-learn
 soundfile>=0.12
 fastapi>=0.115.0
 uvicorn[standard]>=0.34.0
+bigvgan>=2.4.1