Upload app.py with huggingface_hub

app.py

@@ -94,6 +94,7 @@ async def set_collider(request: Request):
         "onsetmode": bool(body.get("onsetmode", False)),
         "reset": int(body.get("reset", 0)),
         "seed": int(body.get("seed", 0)),
+        "bank_op": body.get("bank_op", prev.get("bank_op")),
     })
     return {"ok": True}
 
@@ -118,6 +119,12 @@ async def set_audio(request: Request):
     return {"ok": True}
 
 
+@app.get("/banks")
+async def banks(session_id: str = ""):
+    sid = os.path.basename(session_id)
+    return {"bankStatus": [os.path.exists(os.path.join(SESSION_DIR, f"{sid}_bank{i}.pt")) for i in range(3)]}
+
+
 @app.api(name="stream")
 @spaces.GPU(duration=90)
 def stream(session_id: str) -> str:
@@ -141,6 +148,7 @@ def stream(session_id: str) -> str:
     cur_reset = cur_seed = 0
     force_reenc = False
     txt_cache, aud_cache = {}, {}
+    cur_bank_ver = 0
     while time.time() - t0 < 55.0:
         c = read_slot(session_id)
         if c is None:
@@ -157,6 +165,19 @@ def stream(session_id: str) -> str:
         if seed != cur_seed:
             cur_seed = seed
             gen = torch.Generator(device=dev).manual_seed(seed)
+        bop = c.get("bank_op")
+        if bop and int(bop.get("ver", 0)) != cur_bank_ver:    # save/recall generation state
+            cur_bank_ver = int(bop.get("ver", 0))
+            bpath = os.path.join(SESSION_DIR, f"{os.path.basename(session_id)}_bank{int(bop.get('idx', 0))}.pt")
+            try:
+                if bop.get("action") == "save":
+                    torch.save({"dstate": dstate, "history": history, "emitted": emitted}, bpath)
+                elif bop.get("action") == "load" and os.path.exists(bpath):
+                    d = torch.load(bpath, map_location=dev)
+                    dstate, history, emitted = d["dstate"], d["history"].to(dev), int(d["emitted"])
+                    source, cur_sig = None, None
+            except Exception as e:
+                print("[bank] error:", repr(e), flush=True)
         toks = []
         for _ in range(10):                        # per-frame slot read -> steering applies ~instantly
             c = read_slot(session_id) or c