Upload folder using huggingface_hub

README.md

@@ -19,8 +19,9 @@ sliders together. The query embedding is edited through sparse inversion in a
 trained BatchTopK SAE (MuQ-MuLan / music4all), and nearest neighbours are retrieved
 live after each slider move. No audio is hosted — playback is via Spotify embeds.
 
-Runs on **ZeroGPU**: the corpus + MuQ weights are fetched at startup; the model is built
-on the GPU lazily inside the first request and cached.
+Runs on **ZeroGPU**: the corpus + MuQ weights are fetched at startup; the text tower is
+used on GPU only while a new slider is instantiated. Sparse inversion, slider edits,
+and retrieval run on CPU from cached slider masks.
 
 ## Configuration (Space variables / secrets)
 

app.py

@@ -4,9 +4,9 @@ Search for a seed track or draw one at random, make one or more concept sliders,
 and retrieve live as the sliders move. No audio is hosted — playback is via
 Spotify embeds.
 
-ZeroGPU: the corpus + MuQ weights are fetched at startup (CPU); the model is built
-on the GPU lazily inside the first @spaces.GPU call and cached for subsequent
-requests.
+ZeroGPU: the corpus + MuQ weights are fetched at startup (CPU); the text tower is
+used on GPU only while a new slider is made. Live slider edits reuse cached CPU
+masks.
 """
 import html
 import os
@@ -259,7 +259,6 @@ def make_slider(seed_track_id, concept, concepts, *alphas):
     return concepts, note, results_html(results), *_slider_components(concepts, values)
 
 
-@spaces.GPU(duration=120)
 def live_retrieve(seed_track_id, concepts, *alphas):
     concepts = list(concepts or [])
     if not seed_track_id:

core.py

@@ -1,9 +1,9 @@
 """Shared engine for the demo Space and website example generation.
 
 Loads a trained BatchTopK SAE + MuQ text tower and a music4all corpus, then steers a
-seed track toward a free-form concept and retrieves nearest neighbours. Corpus load
-(CPU, cheap) is separated from model load (GPU, MuQ) so the Space can fetch the corpus at
-startup and build the model lazily inside a ZeroGPU call.
+seed track toward free-form concepts and retrieves nearest neighbours. The SAE,
+inversion, edit masks, and retrieval corpus stay on CPU; ZeroGPU is used only while
+embedding a new text concept for slider instantiation.
 
 Configuration (env):
   SSR_CHECKPOINT   local .ckpt path, or a HF model repo id (default: local L0=20 run)
@@ -16,14 +16,13 @@ from __future__ import annotations
 
 import json
 import os
-from functools import lru_cache
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 
 from steerable_retrieval.steer import Slider
 from steerable_retrieval.steer.loading import load_steerable_sae
-from steerable_retrieval.steer.retrieval import topk_cosine_neighbors
 from steerable_retrieval.steer.steering import decode_normalized
 
 HERE = os.path.dirname(__file__)
@@ -32,6 +31,7 @@ CORPUS_REPO = os.environ.get("SSR_CORPUS_REPO")  # HF dataset repo id, if hosted
 CHECKPOINT = os.environ.get("SSR_CHECKPOINT", os.path.join(HERE, os.pardir, "logs/xps/66a0caf9/checkpoints/last.ckpt"))
 L0_SUBFOLDER = os.environ.get("SSR_L0")  # e.g. "L0-20" when CHECKPOINT is a Hub repo
 CONFIG = os.environ.get("SSR_CONFIG")
+INVERSION_METHOD = os.environ.get("SSR_INVERSION_METHOD", "fista")
 
 
 def device() -> str:
@@ -57,17 +57,49 @@ def load_corpus():
 
 class DemoEngine:
     def __init__(self, corpus=None, dev: str | None = None):
-        self.device = dev or device()
+        # Keep SAE inversion, sparse edits, and retrieval on CPU. ZeroGPU is only
+        # needed while embedding new text concepts with MuQ.
+        self.device = "cpu"
+        self.text_device = dev or device()
         embs, ids, meta = corpus if corpus is not None else load_corpus()
-        self.embeddings = torch.from_numpy(embs).to(self.device)
+        self.embeddings = torch.from_numpy(embs).float()
+        self.gallery = F.normalize(self.embeddings, dim=-1)
         self.track_ids = ids
         self.id_to_idx = {t: i for i, t in enumerate(ids)}
         self.meta = meta
-        self.model = load_steerable_sae(CHECKPOINT, device=self.device, subfolder=L0_SUBFOLDER, config_path=CONFIG)
+        self.model = load_steerable_sae(CHECKPOINT, device="cpu", subfolder=L0_SUBFOLDER, config_path=CONFIG)
+        self._slider_cache: dict[str, Slider] = {}
+        self._text_encoder_device = "cpu"
+
+    def _slider_key(self, concept: str) -> str:
+        return " ".join(concept.strip().lower().split())
+
+    def has_slider(self, concept: str) -> bool:
+        return self._slider_key(concept) in self._slider_cache
+
+    def _set_text_encoder_device(self, target: str) -> None:
+        if self._text_encoder_device == target:
+            return
+        text_encoder = self.model.text_encoder
+        if text_encoder is None:
+            return
+        text_encoder.to(target)
+        if hasattr(text_encoder, "device"):
+            text_encoder.device = target
+        self._text_encoder_device = target
 
-    @lru_cache(maxsize=64)
     def _slider(self, concept: str) -> Slider:
-        return Slider(concept, model=self.model, method="adam")
+        key = self._slider_key(concept)
+        if key not in self._slider_cache:
+            target = self.text_device if str(self.text_device).startswith("cuda") and torch.cuda.is_available() else "cpu"
+            self._set_text_encoder_device(target)
+            try:
+                self._slider_cache[key] = Slider(concept.strip(), model=self.model, method=INVERSION_METHOD, device="cpu")
+            finally:
+                # The cached slider stores its text embedding/mask; live edits no
+                # longer need the text tower or any GPU memory.
+                self._set_text_encoder_device("cpu")
+        return self._slider_cache[key]
 
     def track_meta(self, track_id: str) -> dict:
         m = dict(self.meta.get(track_id, {}))
@@ -77,11 +109,21 @@ class DemoEngine:
     def seed_embedding(self, track_id: str) -> torch.Tensor:
         return self.embeddings[self.id_to_idx[track_id]]
 
+    def _retrieve_dense(self, query: torch.Tensor, k: int, *, exclude_idx: int | None = None):
+        query = F.normalize(query.reshape(-1), dim=0)
+        sims = self.gallery @ query
+        if exclude_idx is not None:
+            sims[exclude_idx] = -1e9
+        n_avail = self.gallery.shape[0] - (1 if exclude_idx is not None else 0)
+        vals, idx = torch.topk(sims, k=min(int(k), int(n_avail)), largest=True, sorted=True)
+        return idx.detach().cpu(), vals.detach().cpu()
+
     def steer_and_retrieve(self, seed_track_id: str, concept: str, alpha: float = 1.0, k: int = 8) -> list[dict]:
         slider = self._slider(concept)
         z = self.seed_embedding(seed_track_id)
         seed_idx = self.id_to_idx[seed_track_id]
-        idx, scores = slider.retrieve(z, self.embeddings, alpha=alpha, k=k, exclude_idx=seed_idx)
+        edited = slider.steer(z, alpha=alpha)
+        idx, scores = self._retrieve_dense(edited, k, exclude_idx=seed_idx)
         return self._format_results(idx, scores)
 
     def multi_steer_and_retrieve(self, seed_track_id: str, sliders: list[tuple[str, float]], k: int = 8) -> list[dict]:
@@ -93,19 +135,23 @@ class DemoEngine:
         """
         z = self.seed_embedding(seed_track_id).reshape(1, -1)
         seed_idx = self.id_to_idx[seed_track_id]
+        slider_masks = []
+        for concept, alpha in sliders:
+            concept = concept.strip()
+            alpha = float(alpha)
+            if not concept or abs(alpha) < 1e-6:
+                continue
+            slider_masks.append((self._slider(concept).mask, alpha))
+
         with torch.inference_mode():
             _, sparse, _, _ = self.model.inference(z)
             edited_sparse = sparse.clone()
-            for concept, alpha in sliders:
-                concept = concept.strip()
-                alpha = float(alpha)
-                if not concept or abs(alpha) < 1e-6:
-                    continue
-                mask = self._slider(concept).mask.to(device=self.device, dtype=edited_sparse.dtype)
+            for mask, alpha in slider_masks:
+                mask = mask.to(device=self.device, dtype=edited_sparse.dtype)
                 edited_sparse = edited_sparse + alpha * mask.reshape(1, -1)
             edited_sparse = edited_sparse.clamp_min(0.0)
             edited = decode_normalized(self.model, edited_sparse)
-        idx, scores = topk_cosine_neighbors(edited, self.embeddings, k, exclude_idx=seed_idx)
+        idx, scores = self._retrieve_dense(edited, k, exclude_idx=seed_idx)
         return self._format_results(idx, scores)
 
     def _format_results(self, idx, scores) -> list[dict]: