Upload app.py with huggingface_hub

app.py

@@ -95,6 +95,12 @@ async def set_collider(request: Request):
         "cfg_drums": float(body.get("cfg_drums", 4.0)),
         "model": body.get("model", prev.get("model", "mrt2_base")),
         "buffer": int(body.get("buffer", 0)),
+        "notes": list(body.get("notes", [])),
+        "unmaskwidth": int(body.get("unmaskwidth", 0)),
+        "drumless": bool(body.get("drumless", False)),
+        "onsetmode": bool(body.get("onsetmode", False)),
+        "reset": int(body.get("reset", 0)),
+        "seed": int(body.get("seed", 0)),
     })
     return {"ok": True, "tokens": tokens}
 
@@ -117,32 +123,52 @@ def stream(session_id: str) -> str:
     t0 = time.time()
     buf_log, buf_t = [], t0
     last_enc = 0.0
+    cur_sig = None
+    prev_active = set()
+    cur_reset = cur_seed = 0
     while time.time() - t0 < 55.0:
         c = read_slot(session_id)
         if c is None:
             time.sleep(0.02)
             continue
         mname = c.get("model", "mrt2_base")
-        if mname != cur_name:                      # live model switch -> re-init state
-            cur_name, model = mname, MODELS.get(mname, mrt_base)
+        reset = int(c.get("reset", 0))
+        if mname != cur_name or reset != cur_reset:   # model switch / state reset -> re-init
+            cur_name, model, cur_reset = mname, MODELS.get(mname, mrt_base), reset
             dstate = model.model.decoder.init_streaming_f(1, dev, dt)
             history = torch.zeros((1, 0, model.cfg.num_codebooks), dtype=torch.long, device=dev)
-            emitted, cur_tokens, cur_notes, source = 0, None, None, None
+            emitted, source, cur_sig, prev_active = 0, None, None, set()
+        seed = int(c.get("seed", 0))
+        if seed != cur_seed:
+            cur_seed = seed
+            gen = torch.Generator(device=dev).manual_seed(seed)
         toks = []
         for _ in range(10):                        # per-frame slot read -> steering applies ~instantly
             c = read_slot(session_id) or c
             tokens = c["style_tokens"]
-            active = c.get("notes") or []
-            if source is None or ((tokens != cur_tokens or active != cur_notes) and time.time() - last_enc >= 0.2):
-                cur_tokens, cur_notes = tokens, active
-                last_enc = time.time()
-                aset = set(active)
-                notes = [8 if i in aset else -1 for i in range(128)]
-                cfgs = [discretize_cfg(c.get("cfg_musiccoca", 3.0), 0.2, 40),
+            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)
+            if source is None 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()
+                nvec = []                                  # per-pitch: -1 masked / 0 off / 1 cont / 2 onset / 3 on
+                for pitch in range(128):
+                    if pitch in active:
+                        nvec.append((2 if pitch in onsets else 1) if onsetmode else 3)
+                    elif active and (unmask >= 127 or any(abs(pitch - h) <= unmask for h in active)):
+                        nvec.append(0)                     # solo / unmasked-off
+                    else:
+                        nvec.append(-1)                    # masked (model free)
+                drm = [0] if drumless else [-1]            # 0 no-drum / -1 masked
+                cfgs = [discretize_cfg(c.get("cfg_musiccoca", 1.6), 0.2, 40),
                         discretize_cfg(c.get("cfg_notes", 2.4), 0.2, 40),
                         discretize_cfg(c.get("cfg_drums", 4.0), 1.0, 8)]
                 cond = model._conditioning((list(tokens) + [-1] * model.num_musiccoca)[:model.num_musiccoca],
-                                           notes, drums, cfgs)
+                                           nvec, drm, cfgs)
                 source = model.model.encode(cond).to(dt)
             sampler = make_sampler(c.get("temperature", 1.3), c.get("top_k", 40), gen)
             toks.append(model.model.decoder.step_f(dstate, source, sampler=sampler,