Model into Docker

Dockerfile

@@ -1,17 +1,30 @@
 FROM python:3.10-slim
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
 
-# Install ffmpeg for audio processing
+# 1.  system dependencies
 RUN apt-get update && apt-get install -y ffmpeg git build-essential && rm -rf /var/lib/apt/lists/*
 
 WORKDIR /app
 
-# Install dependencies using uv
+# 2.  hf_hub
+# This ensures changing your code doesn't trigger a re-download
+RUN uv pip install --system huggingface_hub
+
+# 3. Download weights
+RUN mkdir -p /app/vampnet/models && \
+    python3 -c "from huggingface_hub import hf_hub_download; \
+    repo = 'ProjectCETI/wham'; \
+    [hf_hub_download(repo_id=repo, filename=f, local_dir='/app/vampnet/models') \
+    for f in ['codec.pth', 'coarse.pth', 'c2f.pth', 'wavebeat.pth']]"
+
+# 4. Install project dependencies
 COPY pyproject.toml .
 RUN uv pip install --system .
 
-# Copy your code and run the app
+# 5. copy code
 COPY . .
+
 EXPOSE 7860
 ENV GRADIO_SERVER_NAME="0.0.0.0"
+# Ensure we run from the root so imports like 'from vampnet' work
 CMD ["python", "vampnet/app.py"]

vampnet/app.py

@@ -1,3 +1,4 @@
+import os
 import sys
 import uuid
 from pathlib import Path
@@ -8,40 +9,39 @@ import gradio as gr
 import numpy as np
 import torch
 import yaml
+from huggingface_hub import hf_hub_download
+from torch_pitch_shift import pitch_shift
 
 from vampnet import mask as pmask
 from vampnet.interface import Interface
 
-SCRIPT_DIR = Path(__file__).parent
+# 1. Setup paths and WorkDir
+SCRIPT_DIR = Path(__file__).parent.absolute()
+# This ensures relative paths like 'conf/interface.yml' work correctly
+os.chdir(SCRIPT_DIR)
 
-# 2. Define the models directory
 MODEL_DIR = SCRIPT_DIR / "models"
-MODEL_DIR.mkdir(parents=True, exist_ok=True)
 
 
 def ensure_models_exist():
-    """Downloads weights from HF Hub if they aren't in the models/ folder."""
+    """Fallback check for weights. In Docker, these are already baked in."""
     repo_id = "ProjectCETI/wham"
-    # List all the .pth files your app needs
-    files_to_download = ["codec.pth", "coarse.pth", "c2f.pth", "wavebeat.pth"]
-    print(f"Checking for model weights in {MODEL_DIR}...")
-    for filename in files_to_download:
-        target_file = MODEL_DIR / filename
-        if not target_file.exists():
-            print(f"Downloading {filename} from {repo_id}...")
+    files = ["codec.pth", "coarse.pth", "c2f.pth", "wavebeat.pth"]
+
+    for filename in files:
+        if not (MODEL_DIR / filename).exists():
+            print(f"Weight {filename} missing, downloading...")
             hf_hub_download(
-                repo_id=repo_id,
-                filename=filename,
-                local_dir=str(MODEL_DIR),
-                local_dir_use_symlinks=False,
+                repo_id=repo_id, filename=filename, local_dir=str(MODEL_DIR)
             )
-        else:
-            print(f"✓ {filename} found.")
 
 
+# Run the check
 ensure_models_exist()
 
+# 2. Hardware Setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
+# Update sys.argv so argbind finds the config file correctly
 sys.argv = ["app.py", "--args.load", "conf/interface.yml", "--Interface.device", device]
 
 
@@ -50,9 +50,6 @@ Interface = argbind.bind(Interface)
 conf = argbind.parse_args()
 
 
-from torch_pitch_shift import pitch_shift
-
-
 def shift_pitch(signal, interval: int):
     signal.samples = pitch_shift(
         signal.samples, shift=interval, sample_rate=signal.sample_rate
@@ -232,7 +229,6 @@ def _vamp(
 
 
 def _extract_and_call_vamp(data, return_mask):
-    """Extract plain values from Gradio data dict so only picklable args cross the ZeroGPU boundary."""
     return _vamp(
         _input_audio=data[input_audio],
         _num_steps=data[num_steps],