wanstreamer/serve/server.py

"""LiveWan browser demo server.

    livewan-serve --port 17070   # or: python -m wanstreamer.serve.server

Serves a single steerable stream: pick a world (or generate a new one from text),
then steer it with the 96-prompt bank or free text while it runs.
"""

import argparse
import asyncio
import queue
import threading
import time
from contextlib import asynccontextmanager
from pathlib import Path

import torch
import uvicorn
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, Response
from pydantic import BaseModel

from . import paths
from .auth import TokenAuth
from .conditioning import theme_of
from .engine import WORLDS, Engine
import cv2

INDEX = Path(__file__).resolve().parent / "web" / "index.html"

# Populated by main(); the endpoints read it rather than module-level globals built at
# import time, so importing this module has no side effects.
_S = {"engine": None, "boot": {"messages": [], "done": False, "error": None},
      "thumbs": {}, "clients": set()}


def E() -> Engine:
    return _S["engine"]


def build_parser():
    p = argparse.ArgumentParser(prog="livewan-serve", description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--assets", default=paths.ASSETS, help="the LiveWan weights/data dir")
    p.add_argument("--wan-repo", default=paths.WAN_REPO, help="Wan2.1 reference checkout")
    p.add_argument("--base-dir", default=paths.BASE_DIR, help="VAE / umt5 / base model")
    p.add_argument("--weights", default=None, help="defaults to checkpoints/t14b_b64/latest.pt")
    p.add_argument("--worlds-dir", default=paths.WORLDS_DIR, help="where generated worlds go")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=17070)
    p.add_argument("--jpeg-quality", type=int, default=85)
    p.add_argument("--buffer-seconds", type=float, default=1.5)
    p.add_argument("--no-compile", action="store_true",
                   help="skip the one-off VAE decoder compile (~50 ms/block slower)")
    p.add_argument("--no-worldgen", action="store_true",
                   help="disable text->world generation (saves 2.6 GB VRAM)")
    p.add_argument("--token", default="1234",
                   help="shared token for ?token=/Bearer/cookie auth; empty string disables")
    return p



class StartReq(BaseModel):
    world: int
    prompt_idx: int | None = None
    prompt_text: str | None = None
    seed: int = 0
    steps: int = 2
    window: int = 6
    latent_norm: float = 1.0


class SteerReq(BaseModel):
    prompt_idx: int | None = None
    prompt_text: str | None = None
    crossfade: int = 0


class WorldReq(BaseModel):
    prompt_text: str | None = None
    prompt_idx: int | None = None
    steps: int = 30
    seed: int = 0
    guide: float = 5.0



def _make_thumb(wid):
    """One representative frame per world, cached in memory for the picker."""
    

    entry = E().worlds.get(int(wid))
    if not entry:
        return
    d = torch.load(entry["path"], map_location="cpu", weights_only=False)
    frame = d["pixels"][40].numpy()[:, :, ::-1]
    _S["thumbs"][int(wid)] = cv2.imencode(
        ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 82])[1].tobytes()


def _boot_load():
    try:
        E().load(progress=lambda m: _S["boot"]["messages"].append(m))
        for w in list(E().worlds):
            _make_thumb(w)
        _S["boot"]["done"] = True
    except Exception as e:
        _S["boot"]["error"] = f"{type(e).__name__}: {e}"


@asynccontextmanager
async def _lifespan(_app):
    task = asyncio.create_task(_pump())
    yield
    task.cancel()


app = FastAPI(title="LiveWan", lifespan=_lifespan)



@app.get("/", response_class=HTMLResponse)
def index():
    return INDEX.read_text()


@app.get("/api/boot")
def boot():
    b = _S["boot"]
    return {"done": b["done"], "error": b["error"], "messages": b["messages"][-6:],
            "detail": E().status.detail if E() else ""}


@app.get("/api/info")
def info():
    if not _S["boot"]["done"]:
        raise HTTPException(503, "still loading")
    d = E().info()
    d["world_themes"] = {w: theme_of(w) for w in WORLDS}
    return d


@app.get("/api/status")
def status():
    e = E()
    s = e.status
    return {
        "state": s.state, "detail": s.detail, "world": s.world,
        "prompt": s.prompt, "prompt_source": s.prompt_source,
        "stats": s.stats, "error": s.error, "progress": s.progress,
        "buffered": e.frames.qsize(),
        "encoder_loaded": e.encoder.loaded if e.encoder else False,
    }


@app.get("/api/thumb/{world}")
def thumb(world: int):
    if world not in _S["thumbs"]:
        raise HTTPException(404, "no thumbnail")
    return Response(_S["thumbs"][world], media_type="image/jpeg")


@app.post("/api/start")
def start(r: StartReq):
    if not _S["boot"]["done"]:
        raise HTTPException(503, "still loading")
    try:
        E().start(r.world, r.prompt_idx, r.prompt_text, seed=r.seed, steps=r.steps,
                  window=r.window, latent_norm=r.latent_norm)
    except Exception as e:
        raise HTTPException(400, f"{type(e).__name__}: {e}")
    return status()


@app.post("/api/steer")
def steer(r: SteerReq):
    try:
        E().steer(r.prompt_idx, r.prompt_text, crossfade=r.crossfade)
    except Exception as e:
        raise HTTPException(400, f"{type(e).__name__}: {e}")
    return status()


@app.post("/api/world/new")
def world_new(r: WorldReq):
    """Generate a brand-new opening world from text. Slow: ~52 s at 30 steps."""
    if not _S["boot"]["done"]:
        raise HTTPException(503, "still loading")
    try:
        wid = E().generate_world(r.prompt_text, r.prompt_idx, steps=r.steps,
                                 seed=r.seed, guide=r.guide)
        _make_thumb(wid)
    except Exception as e:
        raise HTTPException(400, f"{type(e).__name__}: {e}")
    return {"world": wid, **E().info()}


@app.delete("/api/world/{wid}")
def world_delete(wid: int):
    try:
        E().delete_world(wid)
        _S["thumbs"].pop(wid, None)
    except Exception as e:
        raise HTTPException(400, f"{type(e).__name__}: {e}")
    return E().info()


@app.post("/api/stop")
def stop():
    E().stop()
    return status()



def _next_frame(timeout=0.1):
    try:
        return E().frames.get(True, timeout)
    except queue.Empty:
        return None


async def _pump():
    """Single reader of the frame queue, broadcasting to every viewer.

    One consumer only. If each socket drained the queue itself, two viewers would
    each get half the frames.
    """
    loop = asyncio.get_running_loop()
    clients = _S["clients"]
    last = 0.0
    while True:
        if not clients or E() is None:
            await asyncio.sleep(0.1)
            continue
        frame = await loop.run_in_executor(None, _next_frame, 0.1)
        dead = []
        if frame is not None:
            for c in list(clients):
                try:
                    await c.send_bytes(frame)
                except Exception:
                    dead.append(c)
        if time.time() - last > 0.5:
            last = time.time()
            payload = {"t": "status", **status()}
            for c in list(clients):
                try:
                    await c.send_json(payload)
                except Exception:
                    dead.append(c)
        for c in dead:
            clients.discard(c)


@app.websocket("/ws")
async def ws(sock: WebSocket):
    await sock.accept()
    _S["clients"].add(sock)
    try:
        await sock.send_json({"t": "status", **status()})
        while True:
            await sock.receive_text()  # client keepalive; ignored
    except (WebSocketDisconnect, RuntimeError):
        pass
    finally:
        _S["clients"].discard(sock)



def main(argv=None):
    args = build_parser().parse_args(argv)
    _S["engine"] = Engine(
        args.assets, args.wan_repo, args.base_dir, args.weights,
        jpeg_quality=args.jpeg_quality, buffer_seconds=args.buffer_seconds,
        worlds_dir=args.worlds_dir, allow_worldgen=not args.no_worldgen,
        compile_vae=not args.no_compile,
    )
    threading.Thread(target=_boot_load, daemon=True, name="livewan-boot").start()

    served = TokenAuth(app, args.token) if args.token else app
    if args.token:
        print(f"token auth on; open  http://{args.host}:{args.port}/?token={args.token}")
    else:
        print(f"no auth; open  http://{args.host}:{args.port}/")
    uvicorn.run(served, host=args.host, port=args.port, log_level="warning",
                ws_max_size=16 * 2**20)


if __name__ == "__main__":
    main()