"""Latent Collider — a magenta-style node field that morphs a FLUX image in real time. Architecture mirrors the Magenta RT Image Collider: a long-lived ZeroGPU grant runs a generation loop that watches a per-session "slot" (prompts + inverse-distance weights written by the browser), re-renders with FLUX.1-schnell whenever the blend changes, and streams JPEG frames to the client over a WebSocket. No music — the audio frames become image frames. Set MOCK=1 to run the UI/transport without a GPU (renders a color swatch from the blend). """ import os import sys os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") os.environ["GRADIO_SSR_MODE"] = "false" # serve via FastAPI directly so /ws works import json import time import struct import asyncio import threading import queue as _q import contextvars import gradio as gr from gradio import Server from gradio.context import LocalContext from fastapi import Request, WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse MOCK = os.environ.get("MOCK", "") not in ("", "0", "false", "False") HERE = os.path.dirname(os.path.abspath(__file__)) # --------------------------------------------------------------------------- # Model load (skipped in MOCK). On ZeroGPU the `spaces` shim defers the real # .to("cuda") until inside an @spaces.GPU call, so loading at module top is fine. # --------------------------------------------------------------------------- if not MOCK: import torch import spaces from diffusers import FluxPipeline, AutoencoderTiny from flux_blender import FluxBlender, jpeg_bytes BASE_MODEL = os.environ.get("FLUX_MODEL", "black-forest-labs/FLUX.1-schnell") taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to("cuda") pipe = FluxPipeline.from_pretrained(BASE_MODEL, vae=taef1, torch_dtype=torch.bfloat16) pipe.transformer.to(memory_format=torch.channels_last) pipe.to("cuda") blender = FluxBlender(pipe, device="cuda") else: import io import colorsys import hashlib from PIL import Image, ImageDraw def jpeg_bytes(img, quality=85): buf = io.BytesIO(); img.save(buf, format="JPEG", quality=quality); return buf.getvalue() def _mock_image(prompts, weights, w, h): pairs = [(p, float(x)) for p, x in zip(prompts, weights) if float(x) > 0 and (p or "").strip()] key = "|".join(f"{p}:{round(x,3)}" for p, x in pairs) or "empty" hh = int(hashlib.md5(key.encode()).hexdigest(), 16) r, g, b = [int(c * 255) for c in colorsys.hsv_to_rgb((hh % 360) / 360.0, 0.55, 0.9)] img = Image.new("RGB", (w, h), (r, g, b)) d = ImageDraw.Draw(img) top = max(pairs, key=lambda t: t[1])[0] if pairs else "(empty)" d.text((16, 16), f"MOCK\n{top}\n{len(pairs)} concepts", fill=(20, 20, 20)) return img # --------------------------------------------------------------------------- # Per-session slot (atomic JSON file) — written by /set, read by the GPU loop. # --------------------------------------------------------------------------- SESSION_DIR = "/tmp/latent_collider_sessions" os.makedirs(SESSION_DIR, exist_ok=True) def _slot(sid): return os.path.join(SESSION_DIR, f"{os.path.basename(sid)}.json") def write_slot(sid, d): tmp = _slot(sid) + "." + str(time.time_ns()) with open(tmp, "w") as f: json.dump(d, f) os.replace(tmp, _slot(sid)) def read_slot(sid): try: with open(_slot(sid)) as f: return json.load(f) except Exception: return None app = Server() @app.post("/set") async def set_collider(request: Request): body = await request.json() sid = body["session_id"] write_slot(sid, { "prompts": body.get("prompts") or [], "weights": body.get("weights") or [], "seed": int(body.get("seed", 42)), "steps": max(1, min(8, int(body.get("steps", 4)))), "guidance": float(body.get("guidance", 0.0)), "width": max(256, min(1024, int(body.get("width", 512)))), "height": max(256, min(1024, int(body.get("height", 512)))), "ts": time.time(), }) return {"ok": True} def _gpu_stream_impl(session_id): """Continuous render loop: re-generate whenever the blend signature changes.""" deadline = time.time() + 8.0 while read_slot(session_id) is None and time.time() < deadline: time.sleep(0.05) if not MOCK: # one-time warm-up so the first real frame isn't cold try: blender.generate(["a photograph"], [1.0], seed=42, steps=4, width=512, height=512) print("[warmup] done", flush=True) except Exception as e: print("[warmup]", repr(e), flush=True) t0 = time.time() cap = 1e9 if MOCK else 55.0 cur_sig = None while time.time() - t0 < cap: c = read_slot(session_id) if c is None: time.sleep(0.03); continue prompts = c.get("prompts") or [] weights = c.get("weights") or [] seed = int(c.get("seed", 42)); steps = int(c.get("steps", 4)) guidance = float(c.get("guidance", 0.0)) w = int(c.get("width", 512)); h = int(c.get("height", 512)) sig = (tuple(prompts), tuple(round(float(x), 4) for x in weights), seed, steps, round(guidance, 3), w, h) if sig == cur_sig: time.sleep(0.03); continue cur_sig = sig gen_t = time.time() if MOCK: time.sleep(0.12) # fake some latency img = _mock_image(prompts, weights, w, h) else: img = blender.generate(prompts, weights, seed, steps, guidance, w, h) gen_ms = (time.time() - gen_t) * 1000.0 yield (jpeg_bytes(img), gen_ms) if MOCK: gpu_stream = _gpu_stream_impl else: gpu_stream = spaces.GPU(duration=90)(_gpu_stream_impl) # --------------------------------------------------------------------------- # Worker / session plumbing (re-grants the GPU when a grant expires). # --------------------------------------------------------------------------- _sessions = {} _slock = threading.Lock() def _worker(session_id, frame_q, stop_ev): try: while not stop_ev.is_set(): try: for jb, fm in gpu_stream(session_id): if stop_ev.is_set(): break if frame_q.full(): try: frame_q.get_nowait() except _q.Empty: pass try: frame_q.put_nowait((jb, fm)) except _q.Full: pass except Exception as e: if "abort" in str(e).lower() or "duration" in str(e).lower(): continue # grant expired -> re-grant print("[worker]", repr(e), flush=True); break finally: stop_ev.set() @app.api(name="start") def start(session_id: str = "", request: gr.Request = None) -> str: if not session_id: return "" req = request or LocalContext.request.get(None) # carries X-IP-Token -> ZeroGPU bills the USER with _slock: prev = _sessions.pop(session_id, None) if prev: prev["stop"].set() frame_q = _q.Queue(maxsize=4); stop_ev = threading.Event() def run(): if req is not None: LocalContext.request.set(req) _worker(session_id, frame_q, stop_ev) ctx = contextvars.copy_context() t = threading.Thread(target=ctx.run, args=(run,), daemon=True) with _slock: _sessions[session_id] = {"frames": frame_q, "stop": stop_ev} t.start() return session_id @app.websocket("/ws") async def frame_ws(websocket: WebSocket, session_id: str = ""): await websocket.accept() sess = _sessions.get(session_id) if not sess: await websocket.close(code=1008); return frame_q, stop_ev = sess["frames"], sess["stop"] loop = asyncio.get_event_loop() try: while True: if stop_ev.is_set(): try: await websocket.send_json({"type": "ended"}) except Exception: pass break try: jb, fm = await loop.run_in_executor(None, lambda: frame_q.get(timeout=0.5)) except Exception: continue await websocket.send_bytes(struct.pack("