Skip to content

Cosmos-H-Dreams#

License Weights HuggingFace Paper Python

Real-time action-conditioned surgical video simulation via WebRTC, built on FlashDreams.

Overview#

Cosmos-H-Dreams is a fine-tuned variant of Cosmos-H-Surgical-Simulator, with its own checkpoint and a serving layer in a streaming server, enabling live surgical simulation driven by keyboard or Meta Quest controller input. Given a conditional first frame from a surgical procedure and a live stream of instrument action vectors, the model rolls forward in blocks of generated frames and streams the output to a browser or VR headset in real time via WebRTC.

The system is built on top of FlashDreams, NVIDIA's high-performance inference and serving library for autoregressive video models. It uses a fine-tuned checkpoint from Cosmos-H-Surgical-Simulator and supports two modes:

  • Offline batch inference — feed a JSON manifest of {input_video, input_action, output_video} entries and produce MP4 + raw tensor outputs.
  • Interactive WebRTC — drive the rollout live from a browser (keyboard) or a Meta Quest headset (WebXR), with no action .npy needed.

News#

  • [July, 2026] — Initial release of Cosmos-H-Dreams

Runner Configurations#

Slugs follow the pattern cosmosHDreams-[chunk{N}-][2steps-]{encoder}-{decoder} across four independent axes:

Axis Choices Notes
Chunk size (chunk{N}) chunk2, chunk3 Latent frames per DiT forward pass. Higher = fewer DiT calls for the same output. chunk3 is recommended for throughput.
Schedule 4steps, 2steps 4-step is closer to the training distribution; 2-step gives ~2× DiT speedup at some fidelity cost. Omitting defaults to 4steps.
Encoder vae Full Wan2.1 VAE.
Decoder vae, lighttae Full Wan2.1 VAE vs. TAEHV lighttae (~10× faster decode, modest quality drop).

The recommended chunk3 variants:

Slug Schedule Encoder Decoder
cosmosHDreams-chunk3-4steps-vae-vae ⭐ 4-step vae vae
cosmosHDreams-chunk3-4steps-vae-lighttae 4-step vae lighttae
cosmosHDreams-chunk3-2steps-vae-vae 2-step vae vae
cosmosHDreams-chunk3-2steps-vae-lighttae 2-step vae lighttae

All 12 configs share the same checkpoint and DiT geometry. Run uv run flashdreams-run --help to list every available slug.

Quick Start#

1. Build the container#

docker build -t cosmos-h-dreams:latest docker/

2. Start the container#

Place the following assets under the repo root before launching:

  • checkpoints/ — CosmosH .pt checkpoint(s)
docker run --rm -it \
  --network host \
  --gpus all \
  -e NVIDIA_DRIVER_CAPABILITIES=compute,utility,video \
  -v .:/workspace/cosmos-h-dreams \
  -w /workspace/cosmos-h-dreams \
  cosmos-h-dreams:latest /bin/bash

3. Install dependencies#

uv sync --extra dev --extra runners --group lint

4. Mode A — Offline batch inference#

uv run flashdreams-run cosmosHDreams-chunk3-vae-vae \
  --input-json assets/example_data/offline/suturebot_inference_manifest.json \
  --cr1-embeddings-path checkpoints/cr1_empty_string_text_embeddings.pt \
  --root-dir . \
  --total-blocks 20 \
  --save-comparison True \
  --resolution 288,512 \
  --pipeline.diffusion-model.transformer.checkpoint-path checkpoints/model_ema_bf16_jhutabletop_288x512_h73.pt

Each entry produces <name>.mp4, <name>_annotated.mp4, <name>.npy (raw [3, T, H, W] tensor in [-1, 1]), and <name>_latents.npy.

Note — first-run latency. The first block absorbs torch.compile JIT and the initial CUDA-graph capture (reported as [WARMUP]). Compiled artifacts are cached and reused on subsequent runs of the same configuration and resolution.

5. Mode B — Interactive WebRTC#

Keyboard (any browser):

uv run --package flash-cosmosHDreams python -m cosmosHDreams.webrtc.server \
  --config cosmosHDreams/configs/keyboard_tabletop.yaml

Open http://0.0.0.0:8080/keyboard in any browser to start controlling the surgical robot.

Meta Quest (WebXR):

WebXR requires HTTPS. Generate a self-signed certificate once, replacing <host-ip> with the LAN IP of the server:

openssl req -x509 -newkey rsa:2048 -nodes \
  -keyout key.pem -out cert.pem -days 365 \
  -subj "/CN=quest" \
  -addext "subjectAltName=IP:<host-ip>"

uv run --package flash-cosmosHDreams python -m cosmosHDreams.webrtc.server_quest \
  --config cosmosHDreams/configs/quest_tabletop.yaml

Open https://<host-ip>:8443/quest in the Quest browser and click Enter VR. Hold B to reset the simulation; hold Y for 1 s to exit immersive mode.

Unified (keyboard + Quest on one port):

uv run --package flash-cosmosHDreams python -m cosmosHDreams.webrtc.server_unified \
  --config cosmosHDreams/configs/unified_tabletop.yaml

Serves /keyboard, /quest, /viewer, and / on a single HTTPS port (default 8443). Whichever client connects most recently drives; the other is paused until it reconnects (takeover semantics).

Note — first-input latency. In interactive mode the DiT is compiled with torch.compile on the first forward pass triggered by user input. Expect the first generation to take significantly longer than steady-state; subsequent generations run at normal speed.

How to create your own Surgical Dreams model#

Two paths are supported depending on your starting point.

Start from the Cosmos-H lineage, which already understands instrument kinematics and tissue dynamics from multi-embodiment surgical training.

  1. Collect embodiment data. Record paired videos and instrument action vectors from your robot or procedure. The Open-H Dataset shows the expected format.
  2. Fine-tune. Fine-tune Cosmos-H-Surgical-Simulator on your data following its training guide.
  3. Distill. Distill the fine-tuned checkpoint to a 2- or 4-step diffusion schedule for real-time throughput. See the distillation guide.
  4. Deploy. Pass your distilled checkpoint via --pipeline.diffusion-model.transformer.checkpoint-path to any Cosmos-H-Dreams slug or create a new config for it under configs.

Path B — Bring your own model#

Use this path if you are starting from a different video diffusion architecture or a model trained on a proprietary surgical corpus.

  1. Fine-tune on your embodiment. Train or fine-tune your video generative model on paired video and action data from your target robot and procedure.
  2. Distill to few-step diffusion. Reduce the sampling steps to 2–4 using consistency distillation or flow-matching distillation to meet real-time latency targets. See the distillation guide.
  3. Integrate with FlashDreams. Register your model as a FlashDreams runner. FlashDreams provides AR-cache management, torch.compile/CUDA-graph wrapping, and fast VAE/TAEHV decode. We recommend reading the FlashDreams documentation before starting this process.
  4. Serve. Once registered, your model runs under the same WebRTC serving layer and is controllable from the same keyboard and Meta Quest clients.

Documentation#

Guide Description
cosmosHDreams/GUIDE.md Container setup, dependency sync, config selection, full flag reference for offline inference and all three WebRTC servers
cosmosHDreams/README.md Package reference: run commands, DataChannel protocol, key bindings
cosmosHDreams/configs/ Annotated YAML schemas (runtime, scenes, server, video settings)

System Requirements#

  • NVIDIA GPU with at least 12GB of VRAM.
  • NVIDIA driver R580 series or newer (CUDA 13.x)
  • Python >= 3.12
  • Docker with the NVIDIA Container Toolkit

License#

Component License
Source code Apache 2.0
Cosmos-H-Dreams model weights NVIDIA Open Model

Resources#

Known Issues#

WebRTC stream flickers with wrong colors — the video stream in the browser or Quest headset shows color artifacts or flickers between frames. This happens when the GPU is saturated by the diffusion process and cannot simultaneously run the NVENC hardware encoder reliably. Fix: force the CPU encoder in the server YAML:

video:
  encoder: cpu_libav

Contributing#

See CONTRIBUTING.md for guidelines on reporting bugs and submitting changes.