Deploy NVIDIA Cosmos Robot World Model with WMA

git clone https://github.com/fal-ai-community/fal-demos.git
cd fal-demos
uv venv --python 3.13
source .venv/bin/activate
uv pip install 'fal[wma]==1.82.1'
fal auth login

DOCKERFILE = f"""
FROM nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y --no-install-recommends \\
    ca-certificates curl ffmpeg git libgl1 libglib2.0-0 && rm -rf /var/lib/apt/lists/*
# ...
RUN uv pip install 'fal[wma]==1.82.1'
"""

class RobotWorldModel(fal.wma.App):
    machine_type = "GPU-H100"
    image = ContainerImage.from_dockerfile_str(DOCKERFILE)
    min_concurrency = 0
    max_concurrency = 1
    max_multiplexing = 1
    keep_alive = 300
    startup_timeout = 1800

    realtime_contract = fal.wma.RealtimeContract(
        media=fal.wma.MediaContract(receive=(fal.wma.Track(kind="video"),)),
        client_messages=ClientMessage,
        server_messages=ServerMessage,
    # ...
        interaction=profile.ConfiguredSession(
            setup=ConfigureMessage,
            ready=profile.Event(ConfiguredMessage, correlation="/instruction_version"),
            sequence=profile.Sequence(field="/instruction_version"),

    def on_instruction(self, raw):
        if not self.configured:
            return self.fail("not_configured", "Send configure first.")
        try:
            message = InstructionMessage.model_validate(raw)
        except ValidationError as exc:
            version = raw.get("instruction_version")
            if isinstance(version, int):
                return self.send(
                    InstructionRejectedMessage(
                        instruction_version=version, reason="invalid_instruction"
                    )
                )
            return self.invalid(raw, exc)
        newest = max(self.version, self.pending[0] if self.pending else 0)
        if message.instruction_version <= newest:
            return self.send(
                InstructionRejectedMessage(
                    instruction_version=message.instruction_version,
                    reason="stale_instruction_version",
                )
            )
        # The newest text wins. It reaches the model at the start of the next chunk.
        self.pending = (message.instruction_version, message.instruction)
        self.send(InstructionPendingMessage(instruction_version=message.instruction_version))

        frames = asyncio.Queue(maxsize=FRAME_QUEUE_SIZE)
        robot = RobotSession(
            session, self.service, self.executor, self.observation, frames
        )
        session.on_message("configure", robot.on_configure)
        session.on_message("instruction", robot.on_instruction)
        session.on_message("step", robot.on_step)
        session.defer(robot.shutdown)

        async def expire():
            await asyncio.sleep(SESSION_SECONDS)
            await session.close()

        session.create_task(expire())
        track = fal.wma.make_video_queue_track(frames, FPS)
        return fal.wma.AiortcPeer(
            session,
            lambda pc: pc.addTrack(track),
            rtc_configuration=RTCConfiguration(iceServers=servers),
            disconnected_grace_seconds=10,
        )

fal run fal_demos/video/wma/robot_world_model.py::RobotWorldModel

fal deploy fal_demos/video/wma/robot_world_model.py::RobotWorldModel --auth private

npm install @fal-ai/client@1.11.0-alpha.5 @fal-ai/server-proxy@1.2.1

import { fal } from "@fal-ai/client";
import { wma } from "@fal-ai/client/realtime/wma";

fal.config({ proxyUrl: "/api/fal/proxy" });

// The element needs the muted and playsinline attributes for autoplay: <video id="output" muted playsinline>
const video = document.querySelector<HTMLVideoElement>("#output")!;
let stepId = 0;
let ended = false;

// Your own code. It runs one chunk on the robot and resolves when the robot finishes.
declare function executeActions(actions: number[][]): Promise<void>;

function endSession(message: string) {
  // A failure reaches both onError and onState, so run the cleanup once.
  if (ended) return;
  ended = true;
  video.srcObject = null;
  console.log(message);
}

function sendStep() {
  if (ended) return;
  stepId += 1;
  // Add joint_position, gripper_position, and image_url from your robot to start from its real state.
  session.send({ type: "step", step_id: stepId });
}

const session = fal.realtime.open(wma("your-account/<app-name>"), {
  receive: ["video"],
  onMedia(stream) {
    video.srcObject = stream;
    void video.play().catch(console.error);
  },
  onData: (raw) => {
    const msg = JSON.parse(raw);
    switch (msg.type) {
      case "configured":
        sendStep();
        break;
      case "action_chunk":
        // Match the actions to the step with msg.step_id. Treat the video as a preview.
        executeActions(msg.actions)
          .then(sendStep)
          .catch((error) => {
            console.error(error);
            endSession("Robot action failed.");
            void session.close();
          });
        break;
      case "error":
        console.error(msg.code, msg.error);
        break;
    }
  },
  onState(state) {
    if (state === "closed" || state === "failed") endSession("Session ended.");
  },
  onError(error) {
    console.error(error);
    endSession("Session failed.");
  },
});
try {
  await session.ready;
} catch (error) {
  endSession("Session failed to start.");
  await session.close();
  throw error;
}
session.send({
  type: "configure",
  instruction_version: 1,
  instruction: "Put the marker in the cup.",
  image_url: "https://raw.githubusercontent.com/fal-ai-community/fal-demos/main/fal_demos/video/wma/droid_observation.png",
});