Upload README.md with huggingface_hub

README.md

@@ -1,3 +1,187 @@
 ---
 license: apache-2.0
 ---
+
+# Astronex-World-5B
+
+Real-time interactive video world model for autonomous driving and embodied intelligence: a 5B controllable video DiT
+on the Wan2.2-TI2V-5B backbone, with camera, action, event and text control. One model, two inference forms — a
+*bidirectional* form with full temporal attention and a *causal* form that streams a rollout block by block with
+cross-block KV caching. 832x480 and 1280x704 at 24 fps.
+
+[Project page](https://world.astronex.com.cn/)  | 
+[Code](https://github.com/Astronex-Robotics/Astronex-World)  | 
+[Weights](https://huggingface.co/Astronex-Lab/Astronex-World)
+
+## Gallery
+
+Every clip is generated by these weights. Click a frame to play it on the [project page](https://world.astronex.com.cn/).
+
+| | | |
+|:---:|:---:|:---:|
+| [![Rain-soaked megacity pursuit](https://world.astronex.com.cn/assets/poster/Astronex_01_720p24.jpg)](https://world.astronex.com.cn/assets/videos/Astronex_01_720p24.mp4) | [![Storm lighthouse at dusk](https://world.astronex.com.cn/assets/poster/Astronex_04_720p24.jpg)](https://world.astronex.com.cn/assets/videos/Astronex_04_720p24.mp4) | [![Portrait in a gentle breeze](https://world.astronex.com.cn/assets/poster/Astronex_08_480p.jpg)](https://world.astronex.com.cn/assets/videos/Astronex_08_480p.mp4) |
+| [![Alpine lake at sunset](https://world.astronex.com.cn/assets/poster/Astronex_09_480p.jpg)](https://world.astronex.com.cn/assets/videos/Astronex_09_480p.mp4) | [![Ocean cliff at golden hour](https://world.astronex.com.cn/assets/poster/Astronex_10_480p.jpg)](https://world.astronex.com.cn/assets/videos/Astronex_10_480p.mp4) | [![Open-world coastal pursuit](https://world.astronex.com.cn/assets/poster/Astronex_11_480p.jpg)](https://world.astronex.com.cn/assets/videos/Astronex_11_480p.mp4) |
+| [![Wormhole approach](https://world.astronex.com.cn/assets/poster/Astronex_12_480p.jpg)](https://world.astronex.com.cn/assets/videos/Astronex_12_480p.mp4) | [![Supermassive black hole](https://world.astronex.com.cn/assets/poster/Astronex_13_480p.jpg)](https://world.astronex.com.cn/assets/videos/Astronex_13_480p.mp4) | [![Coastal highway under action control](https://world.astronex.com.cn/assets/poster/cosmos-coast.jpg)](https://world.astronex.com.cn/assets/videos/cosmos-coast.mp4) |
+
+Each frame in the montage below is a separate project output; the full gallery, including 33-second rollouts and
+action-controlled driving, is on the [project page](https://world.astronex.com.cn/).
+
+![Astronex-World sample montage](https://world.astronex.com.cn/assets/img/teaser.jpg)
+
+## The model
+
+Given text or a first frame, the model consumes frame-aligned camera trajectories, a 64-dimensional continuous action
+stream and an embodiment identifier to predict future visual states under control. Cameras are injected with PRoPE in
+all 30 Transformer layers; actions modulate every layer. Beyond steering a rollout, the framework exposes an open
+action-sequence output head — see [Action-sequence output](#action-sequence-output-open-post-training-head).
+
+| | Value |
+| --- | --- |
+| Parameters | 5.35B |
+| Backbone | Wan2.2-TI2V-5B, 30 layers, hidden 3072, 24 heads, FFN 14,336 |
+| Latent | 48 channels, VAE 4x temporal, 16x16 spatial |
+| Conditioning | text, first frame, camera intrinsics/extrinsics (PRoPE), 64-D continuous action + 32 embodiment IDs, event text |
+| Causal form | block-causal attention with cross-block KV cache, 8 latent frames per block, 20-frame local window + 4 persistent sink frames, frustum-overlap retrieval for revisit trajectories |
+| Bidirectional form | full spatiotemporal attention over the clip, 20 latent frames |
+| Output | 832x480 at 24 fps; a 1280x704 preset uses the same weights |
+| Hardware | trained on 2x NVIDIA L20 48 GB, runs on 1x L20 48 GB |
+
+## What is in this repository
+
+| Path | Contents |
+| --- | --- |
+| `model-0000{1,2,3}-of-00003.safetensors`, `config.json`, `model.safetensors.index.json` | the model, bf16, 5.35B parameters, ~10.7 GB — one set of weights for both sampling modes |
+
+The export is merged: the LoRA adapters are folded into the weights, so there is no adapter file to load alongside and
+no `lora` entry in the configuration. Sampling in full-attention mode reads the same files — the checkpoint carries 898
+tensors: the 885 that mode uses, the 7 action-*input* tensors (`action_embedder`), and the 6 action-*output* head tensors
+described below. `config.json` records the head under `action_output_head`.
+
+## Quick start
+
+Load the causal weights directly with `safetensors`:
+
+```python
+import glob, json
+from safetensors.torch import load_file
+from wan_utils.wan_wrapper import WanDiffusionWrapper
+
+cfg = json.load(open("config.json"))
+model = WanDiffusionWrapper(**cfg["model_kwargs"], base_dtype="bfloat16")
+
+state = {}
+for shard in sorted(glob.glob("model-*-of-*.safetensors")):
+    state.update(load_file(shard))
+model.load_state_dict(state)          # strict: 892 tensors, exact match
+```
+
+To sample and to post-train, use the code release, which ships the streaming pipeline, the sampling presets, the
+post-training branches and a browser demo:
+
+```bash
+git clone https://github.com/Astronex-Robotics/Astronex-World
+cd Astronex-World
+
+python wan/check_components.py                 # verify the VAE and text encoder are in place
+bash scripts/infer_causal.sh                   # streaming rollout, 8 steps
+bash scripts/infer_bidirectional.sh            # full-attention form, same weights
+bash scripts/infer_event.sh                    # caption switch partway through a rollout
+```
+
+The Wan2.2 VAE and the UMT5 text encoder are not duplicated here; they come from the minWM `ckpts/` directory, or from
+the official Wan2.2 release.
+
+## Sampling settings that matter
+
+`config.json` carries the settings these weights were validated under:
+
+- **8 UniPC steps** is the validated setting, with CFG 3.0 and an 8-frame inference block. The sampler also accepts 4
+  steps for latency experiments; quality is lower and 8 steps is what the released material uses.
+- **Window 20 + sink 4 + retrieval.** The KV window is `sink + recent history + current block`, and the retrieval path
+  scores evicted history frames by camera-frustum overlap so a revisited place can be redrawn. Sink and window are
+  properties of the weights: sampling them with a different geometry changes the model's behaviour, so change them
+  together rather than one at a time.
+- **Frames arrive in whole blocks.** Generate in multiples of 8 latent frames (23 pixel-frame steps, not 20, once the
+  i2v reference frame is counted).
+
+Generation runs in the compressed video latent space; the VAE is used only for input encoding and final decoding.
+Cameras enter PRoPE as intrinsics plus world-to-camera transforms, with pixel-space intrinsics normalised by the source
+width and height.
+
+## Action-sequence output (open post-training head)
+
+The action pathway is bidirectional by design. One direction is action *input*: a commanded 64-D action steers the
+rollout. The other — `action_output` — is public in this framework and lets the model report the action it depicted,
+which is what makes the representation usable as an inverse-dynamics readout for driving and robot control.
+
+| | Interface |
+| --- | --- |
+| Head | `LayerNorm -> Linear(dim, hidden) -> SiLU -> Linear(hidden, 64)`, output layer zero-initialised so inserting it leaves the model unchanged |
+| Input | spatial hidden tokens mean-pooled per latent frame: `(B, L, dim) -> (B, F, 64)` |
+| Output | 64-D action vector per latent frame; trained with MSE on masked samples so the head cannot copy its input |
+| Install / train | `action_output: true` installs the head, `control_only: [action_out]` trains that head and nothing else |
+| Reference recipe | `configs/posttrain_action_output.yaml` — DROID robot footage, learning rate 2e-5 |
+
+```bash
+bash Astronex-World/scripts/posttrain.sh action_output     # trains the head against a frozen video model
+```
+
+Two settings decide whether the run means anything: `control_only: [action_out]` keeps the video model frozen so the
+head can be retrained at any time, and `action_output_input_dropout: 1.0` withholds the commanded action from the head's
+input, so the head has to read the frames rather than its own conditioning. Embodiment-specific adapters then connect
+the generic 64-D interface to joints, end effectors, grippers or a mobile base.
+
+The head ships in this repository: six tensors, zero-initialised (LayerNorm at 1, both linears at 0), written into the
+export so the parameters exist and are trainable as they stand — enable them with `model_kwargs.action_output: true` and
+train the branch above. Reading actions from the untrained head returns zeros, so train it first. Once trained it doubles
+as a probe: a head that recovers the action from generated frames is direct evidence that the action conditioning is
+really steering the rollout.
+
+## Benchmarks
+
+**WBench** (official evaluation code and default VLM judge, 8-step causal form):
+
+| Split | Average | Quality | Setting | Interaction | Consistency | Physical |
+| --- | --- | --- | --- | --- | --- | --- |
+| Navi 158 | **73.5** | 78.2 | 73.5 | 63.4 | 83.6 | 68.6 |
+| Full 289 | **70.0** | 78.3 | 73.8 | 47.6 | 82.4 | 68.1 |
+
+Astronex-World holds the highest Average among the published WBench entries on both splits. The strongest axis is
+long-horizon stability — background consistency 90.89, geometric consistency 89.64, temporal flicker 96.55 and motion
+smoothness 98.45 on Full 289 — which is what block-causal KV caching, the local window and the persistent sinks are
+built for.
+
+**VBench 1.0** (official evaluation code, custom-video mode):
+
+| Setting | Dimension | Score | Videos |
+| --- | --- | --- | --- |
+| Text-to-video | Motion smoothness | 0.990 | 240 |
+| Text-to-video | Temporal flickering | 0.987 | 240 |
+| Text-to-video | Imaging quality | 0.715 | 240 |
+| Image-to-video | I2V background | 0.997 | 15 |
+| Image-to-video | Temporal flickering | 0.995 | 15 |
+| Image-to-video | Background consistency | 0.985 | 15 |
+| Image-to-video | Camera motion | 0.485 | 103 |
+
+**Motion retention.** Few-step causal distillation trades some subject motion for streaming latency; the
+motion-preservation term in the DMD objective is what keeps rollouts from freezing (without it the same pipeline
+reached roughly half the amplitude it reaches now). The bidirectional form keeps the full motion of the teacher, which
+is why the release ships both forms: causal for interaction and length, bidirectional when the motion has to be exact.
+
+## Speed
+
+Real-time on a single NVIDIA L20 48 GB: the causal form streams block by block with cross-block KV caching and few-step
+UniPC sampling, so frames leave the sampler as the rollout is produced instead of after the whole clip is denoised. The
+4-step path and the 720p preset (same weights, larger frame) are the levers for trading quality against latency.
+
+<!-- Citation: add the technical-report BibTeX entry here once the preprint is posted. -->
+
+## License
+
+Released under the Apache License, Version 2.0 — see [LICENSE](LICENSE). Copyright 2026 Astronex Robotics.
+
+## Provenance of the causal export
+
+Base `online_unipc25_to12_step20_full.pt` plus the 1545 `generator_ema` adapter, merged as
+`W = W_base + (B @ A) * 128/64` in fp32 and cast to bf16. The merge was checked against the LoRA forward pass on real
+tensors before writing: worst relative error 3.7e-7, which is fp32 rounding.