Upload app.py with huggingface_hub

app.py

@@ -76,18 +76,11 @@ app = Server()
 async def set_collider(request: Request):
     body = await request.json()
     sid = body["session_id"]
-    prompts = body.get("prompts") or ["instrumental music"]
-    weights = body.get("weights") or [1.0] * len(prompts)
-    s = float(sum(weights)) or 1.0
-    emb = np.zeros(style_model.embedding_dim, np.float32)
-    for label, w in zip(prompts, weights):
-        if w <= 0:
-            continue
-        emb += (w / s) * _embed(label)
-    tokens = style_model.tokenize(torch.from_numpy(emb)).tolist()
     prev = read_slot(sid) or {}
     write_slot(sid, {
-        "style_tokens": tokens,
+        "prompts": body.get("prompts") or ["instrumental music"],
+        "weights": body.get("weights") or [1.0],
+        "audio": body.get("audio") or [],
         "temperature": float(body.get("temperature", 1.1)),
         "top_k": int(body.get("top_k", 50)),
         "cfg_musiccoca": float(body.get("cfg", 1.6)),
@@ -102,7 +95,27 @@ async def set_collider(request: Request):
         "reset": int(body.get("reset", 0)),
         "seed": int(body.get("seed", 0)),
     })
-    return {"ok": True, "tokens": tokens}
+    return {"ok": True}
+
+
+@app.post("/audio")
+async def set_audio(request: Request):
+    """Store an uploaded clip (native-rate mono float32) per (session, slot);
+    the GPU stream embeds it with MusicCoCa.embed_audio (resampy -> 16kHz)."""
+    body = await request.json()
+    sid = os.path.basename(body["session_id"]); slot = int(body.get("slot", 0))
+    path = os.path.join(SESSION_DIR, f"{sid}_aud{slot}.bin")
+    samples = body.get("samples")
+    if not samples:
+        try: os.remove(path)
+        except OSError: pass
+        return {"ok": True}
+    raw = base64.b64decode(samples)
+    sr = int(body.get("sample_rate", 48000))
+    with open(path + ".tmp", "wb") as f:
+        f.write(int(sr).to_bytes(4, "little")); f.write(raw)
+    os.replace(path + ".tmp", path)
+    return {"ok": True}
 
 
 @app.api(name="stream")
@@ -127,6 +140,7 @@ def stream(session_id: str) -> str:
     prev_active = set()
     cur_reset = cur_seed = 0
     force_reenc = False
+    txt_cache, aud_cache = {}, {}
     while time.time() - t0 < 55.0:
         c = read_slot(session_id)
         if c is None:
@@ -146,17 +160,48 @@ def stream(session_id: str) -> str:
         toks = []
         for _ in range(10):                        # per-frame slot read -> steering applies ~instantly
             c = read_slot(session_id) or c
-            tokens = c["style_tokens"]
+            prompts = c.get("prompts") or ["instrumental music"]
+            weights = c.get("weights") or [1.0] * len(prompts)
+            aflags = (c.get("audio") or []) + [False] * len(prompts)
             active = set(c.get("notes") or [])
             unmask = int(c.get("unmaskwidth", 0))
             onsetmode = bool(c.get("onsetmode", False))
             drumless = bool(c.get("drumless", False))
-            sig = (tuple(tokens), tuple(sorted(active)), unmask, onsetmode, drumless)
+            sig = (tuple(prompts), tuple(round(float(w), 4) for w in weights),
+                   tuple(bool(a) for a in aflags[:len(prompts)]),
+                   tuple(sorted(active)), unmask, onsetmode, drumless)
             if source is None or force_reenc or (sig != cur_sig and time.time() - last_enc >= 0.2):
                 onsets = active - prev_active
                 prev_active = set(active)
                 cur_sig, last_enc = sig, time.time()
-                force_reenc = onsetmode and bool(onsets)   # onset(2)->continuation(1) on the next frame
+                force_reenc = onsetmode and bool(onsets)
+                # ---- style tokens on GPU: embed (text+audio), weight-blend, tokenize ----
+                emb = torch.zeros(style_model.embedding_dim, device=dev, dtype=torch.float32)
+                tot = float(sum(w for w in weights if w > 0)) or 1.0
+                for i in range(len(prompts)):
+                    w = float(weights[i]) if i < len(weights) else 0.0
+                    if w <= 0:
+                        continue
+                    if i < len(aflags) and aflags[i]:
+                        ap = os.path.join(SESSION_DIR, f"{os.path.basename(session_id)}_aud{i}.bin")
+                        try: mt = os.path.getmtime(ap)
+                        except OSError: continue
+                        ce = aud_cache.get(i)
+                        if ce is None or ce[0] != mt:
+                            data = open(ap, "rb").read()
+                            asr = int.from_bytes(data[:4], "little")
+                            samp = np.frombuffer(data[4:], dtype=np.float32)
+                            ce = (mt, style_model.embed_audio(samp, asr).float()); aud_cache[i] = ce
+                        e = ce[1]
+                    else:
+                        lbl = (prompts[i] or "").strip()
+                        if not lbl:
+                            continue
+                        e = txt_cache.get(lbl)
+                        if e is None:
+                            e = style_model.embed(lbl).float(); txt_cache[lbl] = e
+                    emb = emb + (w / tot) * e
+                tokens = style_model.tokenize(emb).tolist()   # onset(2)->continuation(1) on the next frame
                 nvec = []                                  # per-pitch: -1 masked / 0 off / 1 cont / 2 onset / 3 on
                 for pitch in range(128):
                     if pitch in active: