yuntian-deng/NeuralOS-DiT-Mamba

NeuralOS DiT/Mamba

DiT-XL + Mamba-2.7B weights for NeuralOS, which generates desktop screen frames in response to mouse and keyboard inputs.

NeuralOS paper — Luke Rivard, Sun Sun, Hongyu Guo, Wenhu Chen, and Yuntian Deng, ICLR 2026.

Checkpoints

• Latest saved checkpoint: model.safetensors — v6, step 360, saved May 27, 2026. This run continued from the v4 checkpoint below with the Mamba backbone frozen. The step counter is local to the v6 run.
• Demo checkpoint: v4-demo/model.safetensors — v4, step 9,360. This is the checkpoint used by the NeuralOS DiT/Mamba web demo.

Each file is approximately 13.59 GB in FP32 and includes the DiT renderer, Mamba temporal encoder, and autoencoder. No separate pretrained-model download is needed at inference. The original checkpoint tensor values and dtypes are preserved exactly. These are inference weight exports; they do not contain optimizer or training-resume state.

The saved evaluation logs record PSNR at rollout step 32 of 37.35 for v6 and 40.33 for v4. These are historical evaluations, not a new benchmark. Use v4 to reproduce the demo checkpoint; use v6 to try the newest checkpoint on the source server.

checkpoint.json and v4-demo/checkpoint.json record checkpoint steps, source and export hashes, and export checks. SHA256SUMS covers the release files.

Run locally

The bundled runtime removes the need for a private source checkout. It generates 1024 × 768 PNG frames and a JSON run record. Mouse and keyboard inputs condition generated images; the model does not execute commands in an operating system.

Use Linux, Python 3.11, an NVIDIA GPU, and a CUDA toolkit compatible with the installed PyTorch build. The release was tested with PyTorch 2.11.0 / torchvision 0.26.0, CUDA 13.0, Mamba 2.3.1, and causal-conv1d 1.6.1. The commands below reproduce those package versions; building the Mamba kernels requires a compiler and CUDA development toolkit.

python3.11 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install huggingface-hub==1.22.0

# Download the latest checkpoint and runtime.
hf download yuntian-deng/NeuralOS-DiT-Mamba \
  --local-dir NeuralOS-DiT-Mamba --exclude "v4-demo/*"
cd NeuralOS-DiT-Mamba

pip install torch==2.11.0 torchvision==0.26.0 \
  --index-url https://download.pytorch.org/whl/cu130
pip install -r requirements.txt
pip install ninja packaging psutil
pip install --no-build-isolation causal-conv1d==1.6.1 mamba-ssm==2.3.1

python infer.py --model-dir . --frames 8 --output frames-v6

To run the demo checkpoint with the same runtime:

hf download yuntian-deng/NeuralOS-DiT-Mamba \
  v4-demo/model.safetensors v4-demo/checkpoint.json --local-dir .
python infer.py --model-dir v4-demo --frames 8 --output frames-v4

The inference config uses 255 padding steps to initialize the recurrent state and two DDIM sampling steps per frame. --steps changes the sampling count; --seed sets the random seed. --device cuda:0 selects a GPU.

Mouse and keyboard inputs

Pass --actions actions.json with a JSON list containing at least one action per requested frame:

[
  {"x": 512, "y": 384, "left": false, "right": false, "keys": []},
  {"x": 512, "y": 384, "left": true, "right": false, "keys": []},
  {"x": 700, "y": 384, "left": false, "right": false, "keys": []}
]

python infer.py --model-dir . --actions actions.json --frames 3 --output frames

Coordinates are integer display pixels: x from 0 through 1023 and y from 0 through 767. Each action is held for one generated frame. With no actions file, the script generates a mouse sweep. Keyboard names follow the vocabulary in runtime/latent_diffusion/ldm/modules/encoders/mamba_temporal_encoder.py.

Validation

Both exports preserve all 1,245 source tensors exactly and pass finite-value checks. Each checkpoint passed offline loading, 255 recurrent-state initialization steps, and generation of three 1024 × 768 frames with the bundled runtime. All state keys matched, with no missing or unexpected keys. See validation.json for the recorded checks.

Source and licenses

The runtime is a minimal inference subset of the NeuralOS DiT/Mamba implementation. It includes the required source, model config, and latent normalization statistics.

Component licenses are documented in THIRD_PARTY_NOTICES.md: CompVis latent diffusion uses MIT, adapted DiT code uses CC-BY-NC-4.0, and Mamba uses Apache-2.0. Full license texts and source attributions are included. These component notices do not assign a blanket license to the released weights.

Citation

@inproceedings{rivard2026neuralos,
  title={NeuralOS: Towards Simulating Operating Systems via Neural Generative Models},
  author={Luke Rivard and Sun Sun and Hongyu Guo and Wenhu Chen and Yuntian Deng},
  booktitle={International Conference on Learning Representations},
  year={2026}
}