app.py

"""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("<f", fm) + jb)
    except WebSocketDisconnect:
        pass
    finally:
        stop_ev.set()
        with _slock:
            _sessions.pop(session_id, None)


@app.get("/")
async def index():
    with open(os.path.join(HERE, "index.html")) as f:
        return HTMLResponse(f.read())


app.launch(show_error=True, ssr_mode=False)