World Model Accelerator (WMA)

from typing import TypedDict
from fal.wma import RealtimeApp, BatchedFnTrack


class SessionParams(TypedDict, total=False):
    prompt: str


class InfiniteWorlds(RealtimeApp):
    async def on_connect(self, event_handler, session_params: SessionParams):
        @event_handler.on("track")
        def on_track(track):
            if track.kind != "video":
                return

            event_handler.add_track(
                BatchedFnTrack(
                    track,
                    batch_size=4,
                    fn=lambda frames: do_inference(frames, session_params),
                )
            )

BatchedFnTrack(
    source_track,        # the inbound track (e.g. the browser webcam)
    batch_size=4,        # number of frames to group before running fn
    fn=run_inference,    # callable over the batched frames
)

class SessionParams(TypedDict, total=False):
    prompt: str
    guidance_scale: float

import asyncio
import json

import fal
from pydantic import BaseModel
from fastapi.responses import StreamingResponse


class OfferRequest(BaseModel):
    sdp: str
    type: str
    session_id: str


class GrayscaleTrack:
    """Wraps a video track and converts each frame to grayscale."""

    kind = "video"

    def __init__(self, track):
        self._track = track
        self.id = track.id

    async def recv(self):
        import av
        frame = await self._track.recv()
        img = frame.to_ndarray(format="bgr24")
        import numpy as np
        gray = np.mean(img, axis=2, keepdims=True).astype(np.uint8)
        img_gray = np.broadcast_to(gray, img.shape).copy()
        new_frame = av.VideoFrame.from_ndarray(img_gray, format="bgr24")
        new_frame.pts = frame.pts
        new_frame.time_base = frame.time_base
        return new_frame

    def stop(self):
        self._track.stop()


class GrayscaleDemo(fal.App, name="grayscale-demo"):
    requirements = ["aiortc", "numpy"]

    max_multiplexing = 10
    keepalive = 3600
    machine_type = "M"

    @fal.endpoint("/start-session")
    async def start_session(self, request: OfferRequest) -> StreamingResponse:
        from aiortc import RTCPeerConnection, RTCSessionDescription

        pc = RTCPeerConnection()
        dc = pc.createDataChannel("ping")

        @dc.on("message")
        def on_dc_message(message):
            d = json.loads(message)
            if d.get("type") == "ping":
                dc.send(json.dumps({"type": "pong", "client_ts": d["ts"]}))

        @pc.on("track")
        def on_track(track):
            if track.kind == "video":
                pc.addTrack(GrayscaleTrack(track))

        await pc.setRemoteDescription(
            RTCSessionDescription(sdp=request.sdp, type=request.type)
        )
        answer = await pc.createAnswer()
        await pc.setLocalDescription(answer)

        closed = asyncio.Event()

        @pc.on("connectionstatechange")
        def _on_state_change():
            if pc.connectionState in ("closed", "failed", "disconnected"):
                closed.set()

        async def event_stream():
            try:
                # First event: hand the SDP answer back to the client.
                yield "data: " + json.dumps({
                    "sdp": pc.localDescription.sdp,
                    "type": pc.localDescription.type,
                    "session_id": request.session_id,
                }) + "\n\n"

                # Hold the request open while the peer connection is alive.
                # Emit an SSE comment every 15s so intermediaries don't
                # time the connection out.
                while not closed.is_set():
                    try:
                        await asyncio.wait_for(closed.wait(), timeout=15)
                    except asyncio.TimeoutError:
                        yield ": keepalive\n\n"
            finally:
                await pc.close()

        return StreamingResponse(event_stream(), media_type="text/event-stream")

{
  "app_id": "your-username/your-app",
  "sdp": "<offer SDP from RTCPeerConnection.createOffer()>",
  "type": "offer"
}

{
  "session_id": "<uuid>",
  "sdp": "<answer SDP>",
  "type": "answer"
}

{
  "session_id": "<session_id from /session>"
}

const FAL_KEY = "your-fal-key";
const APP_ID = "your-username/your-app";
const WMA_URL = "https://wma.fal.run";

// 1. Create a peer connection and add the camera.
const pc = new RTCPeerConnection({
  iceServers: [{ urls: "stun:stun.l.google.com:19302" }],
});

const stream = await navigator.mediaDevices.getUserMedia({ video: true });
stream.getTracks().forEach((track) => pc.addTrack(track, stream));

// Receive the processed video from the server.
pc.ontrack = (ev) => {
  remoteVideo.srcObject = ev.streams[0] ?? new MediaStream([ev.track]);
};

// The server may open a data channel for application-level messages
// (e.g. ping/pong latency probes or session parameter updates).
pc.ondatachannel = (ev) => {
  const dc = ev.channel;
  dc.onmessage = (msg) => {
    const data = JSON.parse(msg.data);
    // handle messages from the server
  };
};

// 2. Create the offer and wait for ICE candidates to be gathered.
const offer = await pc.createOffer();
await pc.setLocalDescription(offer);

await new Promise((resolve) => {
  if (pc.iceGatheringState === "complete") return resolve();
  const timeout = setTimeout(resolve, 5000);
  pc.onicegatheringstatechange = () => {
    if (pc.iceGatheringState === "complete") {
      clearTimeout(timeout);
      resolve();
    }
  };
});

// 3. Send the offer to the WMA bridge and apply the answer.
const resp = await fetch(WMA_URL + "/session", {
  method: "POST",
  headers: {
    "Authorization": "Key " + FAL_KEY,
    "Content-Type": "application/json",
  },
  body: JSON.stringify({
    app_id: APP_ID,
    sdp: pc.localDescription.sdp,
    type: pc.localDescription.type,
  }),
});

const { session_id, sdp, type } = await resp.json();
await pc.setRemoteDescription(new RTCSessionDescription({ sdp, type }));

// 4. Keep the session alive with heartbeats.
const heartbeat = setInterval(() =>
  fetch(WMA_URL + "/session/heartbeat", {
    method: "POST",
    headers: {
      "Authorization": "Key " + FAL_KEY,
      "Content-Type": "application/json",
    },
    body: JSON.stringify({ session_id }),
  }),
  5000
);

// 5. Clean up when the connection ends.
function cleanup() {
  clearInterval(heartbeat);
  pc.close();
  stream.getTracks().forEach((t) => t.stop());
}

pc.onconnectionstatechange = () => {
  if (pc.connectionState === "failed" || pc.connectionState === "closed") {
    cleanup();
  }
};