Update README

README.md

@@ -11,80 +11,111 @@ license: other
 
 # Cosmos3 Action Viewer
 
-Docker/Viser Space for the Cosmos3 Action Viewer. The app runs the Viser
-frontend and local Cosmos3 generation in one process. It does not use OAuth, a
-FastAPI proxy, a Gradio backend, or ZeroGPU quota forwarding.
+An interactive, browser-based viewer for robot action datasets that runs Cosmos3
+world-model generation locally. Three modes are supported:
 
-## Runtime
+- **Forward dynamics** — predict future video frames from a sequence of actions.
+- **Inverse dynamics** — predict the actions that connect observed frames.
+- **Policy** — predict the next actions to take.
 
-The default runtime starts the Viser Action Viz app on `PORT` and immediately
-warms the persistent in-process Cosmos runner in the background. If a user
-clicks a generation button before warmup completes, that generation waits on the
-same runner lock and reuses the initialized pipeline. Packaged demo data lives
-under `assets/examples/`. Generated artifacts are written below
-`ACTION_VIZ_OUTPUT_ROOT`, loaded by the Viser session, and removed after each
-generation/session cleanup path.
+The app serves the Viser frontend and runs Cosmos3 generation in a single
+process. Packaged demo data ships under `assets/examples/`, so you can try every
+mode without downloading any datasets — you only need the model checkpoint.
 
-The hosted Space should be public for the Viser websocket to connect through the
-Hugging Face frontend. Private mode has been observed to serve the page while
-leaving the Viser websocket disconnected.
+## Run Locally
 
-## Model Paths
+You need an NVIDIA GPU with >= 48GiB memory and the
+[NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
+installed.
 
-Local Docker runs should mount pre-downloaded app data:
+One time, download the Cosmos3-Nano checkpoint:
+```bash
+hf download nvidia/Cosmos3-Nano --local-dir app_data/models/Cosmos3-Nano
+```
 
+To run:
 ```bash
+# Build the image.
+docker build -t cosmos3-action-viewer .
+
+# Launch docker, mounting the directory that holds the model.
 docker run --rm --gpus all -p 8020:7860 \
-  -v /ssd2/cosmos/action_viz_app_data:/app_data \
+  -v "$(pwd)/app_data:/app_data" \
   cosmos3-action-viewer
 ```
 
-Preferred local model locations:
+Then open <http://localhost:8020>.
 
-- `/ssd2/cosmos/action_viz_app_data/models/nvidia--Cosmos3-Nano`
-- `/ssd2/cosmos/action_viz_app_data/models/Cosmos3-Nano`
+The container looks for the checkpoint under `/app_data/models/` (the mount
+above), so any of these on-disk layouts work:
 
-The app fails fast if the checkpoint is missing. It does not download
-`nvidia/Cosmos3-Nano` during startup or generation.
+- `app_data/models/Cosmos3-Nano`
+- `app_data/models/nvidia--Cosmos3-Nano`
 
-Useful environment variables:
+On first launch the model is warmed up in the background. A generation requested
+before warmup finishes simply waits for it to complete, then reuses the loaded
+pipeline — so the first generation is the slowest and later ones are faster.
 
-- `ACTION_VIZ_CHECKPOINT_LOCAL_PATH`: explicit checkpoint path or colon-separated candidates.
-- `ACTION_VIZ_MODEL_MOUNT_PATH`: explicit Hugging Face model mount path.
-- `ACTION_VIZ_APP_DATA_ROOT`: default `/app_data`.
-- `ACTION_VIZ_OUTPUT_ROOT`: default `/tmp/action_viz_sessions`.
-- `ACTION_VIZ_BACKGROUND_WARMUP`: default `1`; set to `0` for lightweight UI-only local starts.
-- `ACTION_VIZ_GENERATE_IMAGE_SIZE`: default `480`.
-- `ACTION_VIZ_GENERATE_NUM_FRAMES`: default `24`.
-- `ACTION_VIZ_GENERATE_NUM_STEPS`: default `30`.
-- `ACTION_VIZ_DEFAULT_DATASET`: default `bridge`.
+## Customizing a Launch
 
-## Diagnostics
+All behavior is controlled through environment variables, which you can pass to
+`docker run` with `-e NAME=value`. The most useful ones:
 
-Checkpoint resolution is read-only:
+Checkpoint / data locations:
 
-```bash
-python /app/scripts/checkpoint_diagnostic.py
-```
+- `ACTION_VIZ_CHECKPOINT_LOCAL_PATH` — explicit checkpoint path, or a
+  colon-separated list of candidate paths to try in order.
+- `ACTION_VIZ_MODEL_MOUNT_PATH` — explicit path to a mounted Hugging Face model.
+- `ACTION_VIZ_APP_DATA_ROOT` — root for app data inside the container (default
+  `/app_data`); models are resolved under `<root>/models/`.
+- `ACTION_VIZ_OUTPUT_ROOT` — where generated artifacts are written (default
+  `/tmp/action_viz_sessions`). Artifacts are cleaned up after each
+  generation/session.
+
+Generation settings:
+
+- `ACTION_VIZ_GENERATE_IMAGE_SIZE` — output resolution (default `480`).
+- `ACTION_VIZ_GENERATE_NUM_FRAMES` — frames per generation (default `24`).
+- `ACTION_VIZ_GENERATE_NUM_STEPS` — sampling steps (default `30`).
+- `ACTION_VIZ_DEFAULT_DATASET` — dataset selected on load (default `bridge`).
+- `ACTION_VIZ_USE_TORCH_COMPILE` — `torch.compile` the pipeline (default `1`).
+  The first generation pays a one-time compile cost in exchange for faster
+  subsequent runs.
+- `ACTION_VIZ_BACKGROUND_WARMUP` — warm the model on startup (default `1`); set
+  to `0` for a lightweight UI-only start that skips loading the model.
 
-One-step generation smoke tests run without the browser UI and clean their
-temporary artifacts by default:
+Example — lower resolution, fewer steps, and a different default dataset:
 
 ```bash
-python /app/scripts/local_generation_smoke.py --model-mode forward_dynamics --num-steps 1
-python /app/scripts/local_generation_smoke.py --model-mode policy --num-steps 1
+docker run --rm --gpus all -p 8020:7860 \
+  -v "$(pwd)/app_data:/app_data" \
+  -e ACTION_VIZ_GENERATE_IMAGE_SIZE=320 \
+  -e ACTION_VIZ_GENERATE_NUM_STEPS=15 \
+  -e ACTION_VIZ_DEFAULT_DATASET=droid \
+  cosmos3-action-viewer
 ```
 
-These smoke tests use `/app_data/huggingface` for auxiliary HF cache files and
-the same local checkpoint resolution as the Viser app.
+## Diagnostics
 
-For local development from the repository root:
+Check how the checkpoint is being resolved (read-only, no model load):
 
 ```bash
-docker build -t cosmos3-action-viewer .
-docker run --rm -p 8020:7860 \
-  -v /ssd2/cosmos/action_viz_app_data:/app_data \
-  cosmos3-action-viewer
+docker run --rm --gpus all \
+  -v "$(pwd)/app_data:/app_data" \
+  cosmos3-action-viewer python /app/scripts/checkpoint_diagnostic.py
 ```
 
-Open <http://localhost:8020>.
+Run a one-step generation smoke test without the browser UI (temporary artifacts
+are cleaned up by default):
+
+```bash
+docker run --rm --gpus all \
+  -v "$(pwd)/app_data:/app_data" \
+  cosmos3-action-viewer \
+  python /app/scripts/local_generation_smoke.py --model-mode forward_dynamics --num-steps 1
+
+docker run --rm --gpus all \
+  -v "$(pwd)/app_data:/app_data" \
+  cosmos3-action-viewer \
+  python /app/scripts/local_generation_smoke.py --model-mode policy --num-steps 1
+```