wm-mcp 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- wm_mcp-0.3.0/.gitignore +7 -0
- wm_mcp-0.3.0/CHANGELOG.md +29 -0
- wm_mcp-0.3.0/DEMO.md +40 -0
- wm_mcp-0.3.0/LICENSE +21 -0
- wm_mcp-0.3.0/PKG-INFO +130 -0
- wm_mcp-0.3.0/README.md +103 -0
- wm_mcp-0.3.0/pyproject.toml +50 -0
- wm_mcp-0.3.0/src/wm/__init__.py +5 -0
- wm_mcp-0.3.0/src/wm/backend.py +132 -0
- wm_mcp-0.3.0/src/wm/cli.py +108 -0
- wm_mcp-0.3.0/src/wm/handles.py +49 -0
- wm_mcp-0.3.0/src/wm/http.py +105 -0
- wm_mcp-0.3.0/src/wm/server.py +87 -0
- wm_mcp-0.3.0/src/wm/service.py +169 -0
- wm_mcp-0.3.0/src/wm/video.py +148 -0
- wm_mcp-0.3.0/tests/conftest.py +58 -0
- wm_mcp-0.3.0/tests/test_cli.py +23 -0
- wm_mcp-0.3.0/tests/test_handles.py +39 -0
- wm_mcp-0.3.0/tests/test_http.py +83 -0
- wm_mcp-0.3.0/tests/test_service.py +91 -0
- wm_mcp-0.3.0/tests/test_video.py +92 -0
- wm_mcp-0.3.0/uv.lock +2344 -0
wm_mcp-0.3.0/.gitignore
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
# Changelog
|
|
2
|
+
|
|
3
|
+
## 0.3.0 — unreleased
|
|
4
|
+
|
|
5
|
+
- **`surprise` streams.** Windows are decoded on the fly — bounded memory,
|
|
6
|
+
unbounded video length. Hour-long recordings scan without silent frame caps
|
|
7
|
+
(the old code quietly degraded a 10-minute video to ~0.4fps). Progress is
|
|
8
|
+
logged to stderr every 20 windows. Output now reports `n_windows`.
|
|
9
|
+
- **Images as goals.** `embed_video` accepts still images (png/jpg/webp…),
|
|
10
|
+
treated as a static clip — so `predict(target_handle=...)` can answer "does
|
|
11
|
+
this video end up looking like this photo?"
|
|
12
|
+
|
|
13
|
+
## 0.2.0 — unreleased
|
|
14
|
+
|
|
15
|
+
- **Time-native everywhere.** `surprise` samples at a configurable `fps`
|
|
16
|
+
(default 8) so window spans are consistent regardless of source frame rate,
|
|
17
|
+
and every segment now carries `start_s`/`end_s` timestamps. `embed_video`
|
|
18
|
+
and `surprise` accept `start_s`/`end_s` to scan a slice of a long video.
|
|
19
|
+
`embed_video` reports the time span it embedded.
|
|
20
|
+
- **Eager model preload.** `wm-mcp serve` starts loading weights in a
|
|
21
|
+
background thread at startup (disable with `--no-eager`), so an agent's
|
|
22
|
+
first tool call doesn't eat the full model load. Loading is lock-guarded.
|
|
23
|
+
- Real encode/decode tests for the video loader (PyAV round-trip).
|
|
24
|
+
|
|
25
|
+
## 0.1.0 — 2026-07-07
|
|
26
|
+
|
|
27
|
+
- Initial probe. Three MCP tools (`embed_video`, `predict`, `surprise`) over
|
|
28
|
+
V-JEPA 2 ViT-L (fp16 on MPS/CUDA), server-side latent handles, localhost
|
|
29
|
+
HTTP API (`serve --http`), `warmup` command, live demo GIF.
|
wm_mcp-0.3.0/DEMO.md
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
# Recording the launch GIF
|
|
2
|
+
|
|
3
|
+
The GIF is one exchange: an LLM agent consulting a world model. Keep it under 30s.
|
|
4
|
+
|
|
5
|
+
## Setup (once)
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
claude mcp add wm -- uvx wm-mcp serve
|
|
9
|
+
# pre-warm so the GIF doesn't show a weight download:
|
|
10
|
+
uvx wm-mcp embed sample.mp4
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
Use a clip with a clear event mid-video (something falls, a cut, a jump scare).
|
|
14
|
+
Good sources: your own camera roll beats stock footage for authenticity.
|
|
15
|
+
|
|
16
|
+
## The take
|
|
17
|
+
|
|
18
|
+
Record the terminal (QuickTime screen recording, or `vhs` if installed; convert
|
|
19
|
+
with `ffmpeg -i in.mov -vf "fps=12,scale=900:-1" out.gif`).
|
|
20
|
+
|
|
21
|
+
1. Start `claude` in a clean directory.
|
|
22
|
+
2. Type exactly: **"Where does this clip stop being predictable? ./sample.mp4"**
|
|
23
|
+
3. Let it call `surprise`, wait for the answer, stop recording after the agent
|
|
24
|
+
summarizes the most surprising segment.
|
|
25
|
+
|
|
26
|
+
## Caption for the README/post
|
|
27
|
+
|
|
28
|
+
> An LLM agent consulting a video world model (V-JEPA 2) as a tool. The model
|
|
29
|
+
> predicts each window's second half from its first; the spike is where reality
|
|
30
|
+
> disagreed. Local, MIT, `uvx wm-mcp serve`.
|
|
31
|
+
|
|
32
|
+
## Launch checklist
|
|
33
|
+
|
|
34
|
+
- [ ] GIF at top of README
|
|
35
|
+
- [ ] `git tag v0.1.0 && git push --tags` (publishes to PyPI via CI — set up PyPI
|
|
36
|
+
trusted publishing for the repo first)
|
|
37
|
+
- [ ] Show HN: "Show HN: wm — run a world model locally, let any LLM agent call it"
|
|
38
|
+
- [ ] r/LocalLLaMA post (angle: first local world-model runner with MCP)
|
|
39
|
+
- [ ] MCP community/Discord (angle: first non-LLM model class behind MCP tools)
|
|
40
|
+
- [ ] Note launch date; **day-14 verdict is binding** (see README kill criteria)
|
wm_mcp-0.3.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Rian Herrmann
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
wm_mcp-0.3.0/PKG-INFO
ADDED
|
@@ -0,0 +1,130 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: wm-mcp
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: Run a world model locally, and let any LLM agent call it as a tool. V-JEPA 2 as an MCP server.
|
|
5
|
+
Project-URL: Homepage, https://github.com/rian-one/wm
|
|
6
|
+
Project-URL: Issues, https://github.com/rian-one/wm/issues
|
|
7
|
+
Author: Rian Herrmann
|
|
8
|
+
License-Expression: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: agents,jepa,mcp,model-context-protocol,v-jepa,video,world-model
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
15
|
+
Requires-Python: >=3.10
|
|
16
|
+
Requires-Dist: av>=12.0
|
|
17
|
+
Requires-Dist: fastapi>=0.110
|
|
18
|
+
Requires-Dist: huggingface-hub>=0.30
|
|
19
|
+
Requires-Dist: mcp>=1.2
|
|
20
|
+
Requires-Dist: numpy>=1.26
|
|
21
|
+
Requires-Dist: pillow>=10.0
|
|
22
|
+
Requires-Dist: torch>=2.4
|
|
23
|
+
Requires-Dist: torchvision>=0.19
|
|
24
|
+
Requires-Dist: transformers>=4.53
|
|
25
|
+
Requires-Dist: uvicorn>=0.30
|
|
26
|
+
Description-Content-Type: text/markdown
|
|
27
|
+
|
|
28
|
+
# wm
|
|
29
|
+
|
|
30
|
+
**Run a world model locally, and let any LLM agent call it as a tool.**
|
|
31
|
+
|
|
32
|
+

|
|
33
|
+
|
|
34
|
+
*Live and unedited: Claude calls the `surprise` tool; V-JEPA 2 watches the clip and
|
|
35
|
+
reports where its predictions broke (z=3.1 at frames 56–71 — the exact splice point
|
|
36
|
+
of the test video). Wait time hidden, output verbatim.*
|
|
37
|
+
|
|
38
|
+
LLMs can *describe* what happens when a cup gets knocked off a table. A world model
|
|
39
|
+
can *watch* it happen and tell you, frame by frame, where reality stopped matching
|
|
40
|
+
its prediction. `wm` puts Meta's [V-JEPA 2](https://huggingface.co/facebook/vjepa2-vitl-fpc64-256)
|
|
41
|
+
(MIT-licensed, 0.3B params) behind an [MCP](https://modelcontextprotocol.io) server,
|
|
42
|
+
so Claude, or any MCP client, can consult a video world model the way it consults a
|
|
43
|
+
database.
|
|
44
|
+
|
|
45
|
+
```
|
|
46
|
+
uvx wm-mcp serve
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
That's the whole install. First run downloads PyTorch and ~1.2GB of model weights —
|
|
50
|
+
**minutes of downloading, not a hang.**
|
|
51
|
+
|
|
52
|
+
## The demo
|
|
53
|
+
|
|
54
|
+
Ask your agent:
|
|
55
|
+
|
|
56
|
+
> "Where does this clip stop being predictable? /path/to/video.mp4"
|
|
57
|
+
|
|
58
|
+
The agent calls `surprise()`, the world model slides a window over the video,
|
|
59
|
+
predicts each window's second half from its first, and reports the segment where
|
|
60
|
+
its prediction failed hardest — z-scored per video, so scores mean something.
|
|
61
|
+
|
|
62
|
+
## The three tools
|
|
63
|
+
|
|
64
|
+
| Tool | What it does |
|
|
65
|
+
|---|---|
|
|
66
|
+
| `embed_video(source)` | Video → latent sequence, stored server-side under an opaque handle. |
|
|
67
|
+
| `predict(handle, horizon_frames, target_handle?)` | Predict the trailing frames of a clip from its leading context (masked in-window prediction — horizon capped at clip length, this is not open-ended rollout). With `target_handle`, also returns cosine similarity between the prediction and a target embedding — goal-similarity scoring, computed server-side. |
|
|
68
|
+
| `surprise(source, window, stride)` | Per-segment "how wrong was the model" scores + the most surprising segment. Anomaly spotting, event detection, "did anything unexpected happen." |
|
|
69
|
+
|
|
70
|
+
Latents never cross the wire: tools exchange handle IDs, tensors stay in the server.
|
|
71
|
+
|
|
72
|
+
## Setup with Claude Code
|
|
73
|
+
|
|
74
|
+
```bash
|
|
75
|
+
claude mcp add wm -- uvx wm-mcp serve
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
Then: *"embed this video and tell me how surprising the ending is"*.
|
|
79
|
+
|
|
80
|
+
Also usable directly, no agent:
|
|
81
|
+
|
|
82
|
+
```bash
|
|
83
|
+
uvx wm-mcp surprise clip.mp4
|
|
84
|
+
uvx wm-mcp embed clip.mp4
|
|
85
|
+
uvx wm-mcp warmup # pre-download weights so first real call is fast
|
|
86
|
+
uvx wm-mcp info
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
## HTTP API (non-MCP consumers)
|
|
90
|
+
|
|
91
|
+
The same three primitives over localhost HTTP — for scripts, notebooks, or
|
|
92
|
+
anything that doesn't speak MCP:
|
|
93
|
+
|
|
94
|
+
```bash
|
|
95
|
+
uvx wm-mcp serve --http # 127.0.0.1:8642
|
|
96
|
+
curl -s localhost:8642/health
|
|
97
|
+
curl -s -X POST localhost:8642/surprise -H 'content-type: application/json' \
|
|
98
|
+
-d '{"source": "clip.mp4"}'
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
Endpoints: `GET /health`, `GET /info`, `POST /embed`, `POST /predict`,
|
|
102
|
+
`POST /surprise` — one-to-one with the MCP tools. Binds to localhost, no auth;
|
|
103
|
+
don't expose it to the internet.
|
|
104
|
+
|
|
105
|
+
## Hardware & honest numbers
|
|
106
|
+
|
|
107
|
+
- Apple Silicon (MPS): ≥16GB unified memory. Measured on an M2 Pro / 16GB:
|
|
108
|
+
`surprise` on a 6s clip = **~18s end to end** (11 windows, fp16). A full
|
|
109
|
+
64-frame `embed_video` is the heavyweight case — expect minutes, not seconds.
|
|
110
|
+
- NVIDIA: ≥8GB VRAM, substantially faster.
|
|
111
|
+
- CPU works; it's just slow.
|
|
112
|
+
- The model loads in fp16 on GPU devices (fp32 attention over 8k tokens will
|
|
113
|
+
swap a 16GB machine — we measured it, you don't want it).
|
|
114
|
+
|
|
115
|
+
## What this is (and isn't)
|
|
116
|
+
|
|
117
|
+
This is a **probe** — the smallest useful bridge between LLM agents and video world
|
|
118
|
+
models, shipped to find out whether anyone wants that bridge. It is not a robotics
|
|
119
|
+
stack, not a leaderboard, not a hosted service, and `predict` is honest about what
|
|
120
|
+
V-JEPA 2's predictor actually does (masked latent prediction within a clip), not a
|
|
121
|
+
marketing claim about simulating the future.
|
|
122
|
+
|
|
123
|
+
Public kill criteria, decided before launch: if by day 14 this has under 200 stars
|
|
124
|
+
AND under 10 substantive issues/PRs from strangers AND under 3 unprompted
|
|
125
|
+
integrations — it gets archived with a public retro. If it resonates, the next step
|
|
126
|
+
is a protocol, not a feature list.
|
|
127
|
+
|
|
128
|
+
## License
|
|
129
|
+
|
|
130
|
+
MIT. V-JEPA 2 weights are MIT-licensed by Meta FAIR.
|
wm_mcp-0.3.0/README.md
ADDED
|
@@ -0,0 +1,103 @@
|
|
|
1
|
+
# wm
|
|
2
|
+
|
|
3
|
+
**Run a world model locally, and let any LLM agent call it as a tool.**
|
|
4
|
+
|
|
5
|
+

|
|
6
|
+
|
|
7
|
+
*Live and unedited: Claude calls the `surprise` tool; V-JEPA 2 watches the clip and
|
|
8
|
+
reports where its predictions broke (z=3.1 at frames 56–71 — the exact splice point
|
|
9
|
+
of the test video). Wait time hidden, output verbatim.*
|
|
10
|
+
|
|
11
|
+
LLMs can *describe* what happens when a cup gets knocked off a table. A world model
|
|
12
|
+
can *watch* it happen and tell you, frame by frame, where reality stopped matching
|
|
13
|
+
its prediction. `wm` puts Meta's [V-JEPA 2](https://huggingface.co/facebook/vjepa2-vitl-fpc64-256)
|
|
14
|
+
(MIT-licensed, 0.3B params) behind an [MCP](https://modelcontextprotocol.io) server,
|
|
15
|
+
so Claude, or any MCP client, can consult a video world model the way it consults a
|
|
16
|
+
database.
|
|
17
|
+
|
|
18
|
+
```
|
|
19
|
+
uvx wm-mcp serve
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
That's the whole install. First run downloads PyTorch and ~1.2GB of model weights —
|
|
23
|
+
**minutes of downloading, not a hang.**
|
|
24
|
+
|
|
25
|
+
## The demo
|
|
26
|
+
|
|
27
|
+
Ask your agent:
|
|
28
|
+
|
|
29
|
+
> "Where does this clip stop being predictable? /path/to/video.mp4"
|
|
30
|
+
|
|
31
|
+
The agent calls `surprise()`, the world model slides a window over the video,
|
|
32
|
+
predicts each window's second half from its first, and reports the segment where
|
|
33
|
+
its prediction failed hardest — z-scored per video, so scores mean something.
|
|
34
|
+
|
|
35
|
+
## The three tools
|
|
36
|
+
|
|
37
|
+
| Tool | What it does |
|
|
38
|
+
|---|---|
|
|
39
|
+
| `embed_video(source)` | Video → latent sequence, stored server-side under an opaque handle. |
|
|
40
|
+
| `predict(handle, horizon_frames, target_handle?)` | Predict the trailing frames of a clip from its leading context (masked in-window prediction — horizon capped at clip length, this is not open-ended rollout). With `target_handle`, also returns cosine similarity between the prediction and a target embedding — goal-similarity scoring, computed server-side. |
|
|
41
|
+
| `surprise(source, window, stride)` | Per-segment "how wrong was the model" scores + the most surprising segment. Anomaly spotting, event detection, "did anything unexpected happen." |
|
|
42
|
+
|
|
43
|
+
Latents never cross the wire: tools exchange handle IDs, tensors stay in the server.
|
|
44
|
+
|
|
45
|
+
## Setup with Claude Code
|
|
46
|
+
|
|
47
|
+
```bash
|
|
48
|
+
claude mcp add wm -- uvx wm-mcp serve
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
Then: *"embed this video and tell me how surprising the ending is"*.
|
|
52
|
+
|
|
53
|
+
Also usable directly, no agent:
|
|
54
|
+
|
|
55
|
+
```bash
|
|
56
|
+
uvx wm-mcp surprise clip.mp4
|
|
57
|
+
uvx wm-mcp embed clip.mp4
|
|
58
|
+
uvx wm-mcp warmup # pre-download weights so first real call is fast
|
|
59
|
+
uvx wm-mcp info
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
## HTTP API (non-MCP consumers)
|
|
63
|
+
|
|
64
|
+
The same three primitives over localhost HTTP — for scripts, notebooks, or
|
|
65
|
+
anything that doesn't speak MCP:
|
|
66
|
+
|
|
67
|
+
```bash
|
|
68
|
+
uvx wm-mcp serve --http # 127.0.0.1:8642
|
|
69
|
+
curl -s localhost:8642/health
|
|
70
|
+
curl -s -X POST localhost:8642/surprise -H 'content-type: application/json' \
|
|
71
|
+
-d '{"source": "clip.mp4"}'
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
Endpoints: `GET /health`, `GET /info`, `POST /embed`, `POST /predict`,
|
|
75
|
+
`POST /surprise` — one-to-one with the MCP tools. Binds to localhost, no auth;
|
|
76
|
+
don't expose it to the internet.
|
|
77
|
+
|
|
78
|
+
## Hardware & honest numbers
|
|
79
|
+
|
|
80
|
+
- Apple Silicon (MPS): ≥16GB unified memory. Measured on an M2 Pro / 16GB:
|
|
81
|
+
`surprise` on a 6s clip = **~18s end to end** (11 windows, fp16). A full
|
|
82
|
+
64-frame `embed_video` is the heavyweight case — expect minutes, not seconds.
|
|
83
|
+
- NVIDIA: ≥8GB VRAM, substantially faster.
|
|
84
|
+
- CPU works; it's just slow.
|
|
85
|
+
- The model loads in fp16 on GPU devices (fp32 attention over 8k tokens will
|
|
86
|
+
swap a 16GB machine — we measured it, you don't want it).
|
|
87
|
+
|
|
88
|
+
## What this is (and isn't)
|
|
89
|
+
|
|
90
|
+
This is a **probe** — the smallest useful bridge between LLM agents and video world
|
|
91
|
+
models, shipped to find out whether anyone wants that bridge. It is not a robotics
|
|
92
|
+
stack, not a leaderboard, not a hosted service, and `predict` is honest about what
|
|
93
|
+
V-JEPA 2's predictor actually does (masked latent prediction within a clip), not a
|
|
94
|
+
marketing claim about simulating the future.
|
|
95
|
+
|
|
96
|
+
Public kill criteria, decided before launch: if by day 14 this has under 200 stars
|
|
97
|
+
AND under 10 substantive issues/PRs from strangers AND under 3 unprompted
|
|
98
|
+
integrations — it gets archived with a public retro. If it resonates, the next step
|
|
99
|
+
is a protocol, not a feature list.
|
|
100
|
+
|
|
101
|
+
## License
|
|
102
|
+
|
|
103
|
+
MIT. V-JEPA 2 weights are MIT-licensed by Meta FAIR.
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "wm-mcp"
|
|
3
|
+
version = "0.3.0"
|
|
4
|
+
description = "Run a world model locally, and let any LLM agent call it as a tool. V-JEPA 2 as an MCP server."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = "MIT"
|
|
7
|
+
requires-python = ">=3.10"
|
|
8
|
+
authors = [{ name = "Rian Herrmann" }]
|
|
9
|
+
keywords = ["world-model", "v-jepa", "jepa", "mcp", "model-context-protocol", "video", "agents"]
|
|
10
|
+
classifiers = [
|
|
11
|
+
"Development Status :: 3 - Alpha",
|
|
12
|
+
"Intended Audience :: Developers",
|
|
13
|
+
"Programming Language :: Python :: 3",
|
|
14
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
15
|
+
]
|
|
16
|
+
dependencies = [
|
|
17
|
+
"mcp>=1.2",
|
|
18
|
+
"torch>=2.4",
|
|
19
|
+
"torchvision>=0.19",
|
|
20
|
+
"transformers>=4.53",
|
|
21
|
+
"numpy>=1.26",
|
|
22
|
+
"av>=12.0",
|
|
23
|
+
"pillow>=10.0",
|
|
24
|
+
"huggingface_hub>=0.30",
|
|
25
|
+
"fastapi>=0.110",
|
|
26
|
+
"uvicorn>=0.30",
|
|
27
|
+
]
|
|
28
|
+
|
|
29
|
+
[project.urls]
|
|
30
|
+
Homepage = "https://github.com/rian-one/wm"
|
|
31
|
+
Issues = "https://github.com/rian-one/wm/issues"
|
|
32
|
+
|
|
33
|
+
[project.scripts]
|
|
34
|
+
wm-mcp = "wm.cli:main"
|
|
35
|
+
|
|
36
|
+
[dependency-groups]
|
|
37
|
+
dev = ["pytest>=8.0", "httpx>=0.27"]
|
|
38
|
+
|
|
39
|
+
[build-system]
|
|
40
|
+
requires = ["hatchling"]
|
|
41
|
+
build-backend = "hatchling.build"
|
|
42
|
+
|
|
43
|
+
[tool.hatch.build.targets.wheel]
|
|
44
|
+
packages = ["src/wm"]
|
|
45
|
+
|
|
46
|
+
[tool.hatch.build.targets.sdist]
|
|
47
|
+
exclude = ["*.gif", "*.mp4", "*.png", "demo.tape", ".github"]
|
|
48
|
+
|
|
49
|
+
[tool.pytest.ini_options]
|
|
50
|
+
testpaths = ["tests"]
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
"""V-JEPA 2 backend: encoder + masked-latent predictor.
|
|
2
|
+
|
|
3
|
+
Semantics pinned to what the architecture actually does (transformers VJEPA2Model):
|
|
4
|
+
- the encoder embeds a clip into N latent tokens (temporal-major: N = T' * S where
|
|
5
|
+
T' = frames / tubelet_size and S = spatial patches per tubelet);
|
|
6
|
+
- the predictor fills in latents for *target* token positions given *context*
|
|
7
|
+
positions within the same clip. Future prediction = mask the trailing frames.
|
|
8
|
+
There is no arbitrary-horizon rollout; horizon is capped at the clip length.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
import numpy as np
|
|
14
|
+
|
|
15
|
+
from . import MODEL_ID
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def pick_device() -> str:
|
|
19
|
+
import torch
|
|
20
|
+
|
|
21
|
+
if torch.backends.mps.is_available():
|
|
22
|
+
return "mps"
|
|
23
|
+
if torch.cuda.is_available():
|
|
24
|
+
return "cuda"
|
|
25
|
+
return "cpu"
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class VJepa2Backend:
|
|
29
|
+
def __init__(self, model_id: str = MODEL_ID, device: str | None = None):
|
|
30
|
+
import threading
|
|
31
|
+
|
|
32
|
+
self.model_id = model_id
|
|
33
|
+
self.device = device or pick_device()
|
|
34
|
+
self._model = None
|
|
35
|
+
self._processor = None
|
|
36
|
+
self._load_lock = threading.Lock()
|
|
37
|
+
|
|
38
|
+
def _load(self):
|
|
39
|
+
# lock: eager background preload and a first tool call may race here
|
|
40
|
+
with self._load_lock:
|
|
41
|
+
if self._model is None:
|
|
42
|
+
import torch
|
|
43
|
+
from transformers import AutoModel, AutoVideoProcessor
|
|
44
|
+
|
|
45
|
+
self._processor = AutoVideoProcessor.from_pretrained(self.model_id)
|
|
46
|
+
# fp16 on GPU: fp32 attention over 8k tokens swaps a 16GB machine
|
|
47
|
+
# (measured 823s/clip); fp16 keeps activations in memory.
|
|
48
|
+
dtype = torch.float16 if self.device in ("mps", "cuda") else torch.float32
|
|
49
|
+
self._model = AutoModel.from_pretrained(
|
|
50
|
+
self.model_id, attn_implementation="sdpa", dtype=dtype
|
|
51
|
+
)
|
|
52
|
+
self._model.to(self.device).eval()
|
|
53
|
+
return self._model, self._processor
|
|
54
|
+
|
|
55
|
+
def preload_async(self) -> None:
|
|
56
|
+
"""Start loading weights in a daemon thread so the first call is fast."""
|
|
57
|
+
import threading
|
|
58
|
+
|
|
59
|
+
threading.Thread(target=self._load, daemon=True, name="wm-preload").start()
|
|
60
|
+
|
|
61
|
+
# -- layout ---------------------------------------------------------------
|
|
62
|
+
|
|
63
|
+
def layout(self, n_frames: int):
|
|
64
|
+
"""(tokens_per_tubelet, tubelet_size, n_tubelets) for a clip of n_frames."""
|
|
65
|
+
model, _ = self._load()
|
|
66
|
+
cfg = model.config
|
|
67
|
+
spatial = (cfg.crop_size // cfg.patch_size) ** 2
|
|
68
|
+
tubelet = cfg.tubelet_size
|
|
69
|
+
return spatial, tubelet, max(1, n_frames // tubelet)
|
|
70
|
+
|
|
71
|
+
# -- ops ------------------------------------------------------------------
|
|
72
|
+
|
|
73
|
+
def preprocess(self, frames: np.ndarray):
|
|
74
|
+
"""(T, H, W, C) uint8 -> pixel tensor on device, in the model's dtype."""
|
|
75
|
+
model, processor = self._load()
|
|
76
|
+
inputs = processor(frames, return_tensors="pt")
|
|
77
|
+
pixels = inputs["pixel_values_videos"].to(self.device)
|
|
78
|
+
return pixels.to(next(model.parameters()).dtype)
|
|
79
|
+
|
|
80
|
+
def encode(self, pixels) -> "object":
|
|
81
|
+
"""Pixel tensor -> encoder latents (N, D) on CPU."""
|
|
82
|
+
import torch
|
|
83
|
+
|
|
84
|
+
model, _ = self._load()
|
|
85
|
+
with torch.inference_mode():
|
|
86
|
+
out = model(pixel_values_videos=pixels, skip_predictor=True)
|
|
87
|
+
return out.last_hidden_state[0].float().cpu()
|
|
88
|
+
|
|
89
|
+
def predict_masked(self, pixels, n_frames: int, horizon_frames: int):
|
|
90
|
+
"""Predict latents for the trailing `horizon_frames` from the leading context.
|
|
91
|
+
|
|
92
|
+
Returns (predicted, actual): both (N_target, D) on CPU, where actual is the
|
|
93
|
+
encoder's real latents at the same target positions.
|
|
94
|
+
"""
|
|
95
|
+
import torch
|
|
96
|
+
|
|
97
|
+
model, _ = self._load()
|
|
98
|
+
spatial, tubelet, n_tubelets = self.layout(n_frames)
|
|
99
|
+
|
|
100
|
+
horizon_tubelets = max(1, min(horizon_frames // tubelet, n_tubelets - 1))
|
|
101
|
+
ctx_tokens = (n_tubelets - horizon_tubelets) * spatial
|
|
102
|
+
total_tokens = n_tubelets * spatial
|
|
103
|
+
|
|
104
|
+
dev = pixels.device
|
|
105
|
+
context_mask = [torch.arange(0, ctx_tokens, device=dev).unsqueeze(0)]
|
|
106
|
+
target_mask = [torch.arange(ctx_tokens, total_tokens, device=dev).unsqueeze(0)]
|
|
107
|
+
|
|
108
|
+
with torch.inference_mode():
|
|
109
|
+
out = model(
|
|
110
|
+
pixel_values_videos=pixels,
|
|
111
|
+
context_mask=context_mask,
|
|
112
|
+
target_mask=target_mask,
|
|
113
|
+
)
|
|
114
|
+
predicted = out.predictor_output.last_hidden_state[0].float().cpu()
|
|
115
|
+
actual = out.predictor_output.target_hidden_state[0].float().cpu()
|
|
116
|
+
return predicted, actual
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
_backend: VJepa2Backend | None = None
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def get_backend() -> VJepa2Backend:
|
|
123
|
+
global _backend
|
|
124
|
+
if _backend is None:
|
|
125
|
+
_backend = VJepa2Backend()
|
|
126
|
+
return _backend
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def set_backend(backend) -> None:
|
|
130
|
+
"""Test seam: swap in a fake backend."""
|
|
131
|
+
global _backend
|
|
132
|
+
_backend = backend
|
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
"""CLI: `wm-mcp serve` (MCP over stdio) plus direct commands for demos and debugging."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import json
|
|
7
|
+
import sys
|
|
8
|
+
|
|
9
|
+
from . import MODEL_ID, __version__
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def main(argv: list[str] | None = None) -> int:
|
|
13
|
+
parser = argparse.ArgumentParser(
|
|
14
|
+
prog="wm-mcp",
|
|
15
|
+
description="Run a world model locally, and let any LLM agent call it as a tool.",
|
|
16
|
+
)
|
|
17
|
+
parser.add_argument("--version", action="version", version=f"wm-mcp {__version__} ({MODEL_ID})")
|
|
18
|
+
sub = parser.add_subparsers(dest="command")
|
|
19
|
+
|
|
20
|
+
p_serve = sub.add_parser("serve", help="start the MCP server (stdio), or HTTP with --http")
|
|
21
|
+
p_serve.add_argument("--http", action="store_true", help="serve a localhost HTTP API instead of MCP stdio")
|
|
22
|
+
p_serve.add_argument("--host", default="127.0.0.1")
|
|
23
|
+
p_serve.add_argument("--port", type=int, default=8642)
|
|
24
|
+
p_serve.add_argument(
|
|
25
|
+
"--no-eager",
|
|
26
|
+
action="store_true",
|
|
27
|
+
help="don't preload model weights in the background on startup",
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
sub.add_parser("warmup", help="download weights and load the model so first use is fast")
|
|
31
|
+
|
|
32
|
+
p_embed = sub.add_parser("embed", help="embed a video, print handle info as JSON")
|
|
33
|
+
p_embed.add_argument("source", help="video path or URL")
|
|
34
|
+
|
|
35
|
+
p_surprise = sub.add_parser("surprise", help="find where a video stops being predictable")
|
|
36
|
+
p_surprise.add_argument("source", help="video path or URL")
|
|
37
|
+
p_surprise.add_argument("--window", type=int, default=16)
|
|
38
|
+
p_surprise.add_argument("--stride", type=int, default=8)
|
|
39
|
+
p_surprise.add_argument("--fps", type=float, default=8.0, help="sampling rate (frames/sec)")
|
|
40
|
+
p_surprise.add_argument("--start", type=float, default=None, help="clip start (seconds)")
|
|
41
|
+
p_surprise.add_argument("--end", type=float, default=None, help="clip end (seconds)")
|
|
42
|
+
|
|
43
|
+
sub.add_parser("info", help="print device/model info as JSON")
|
|
44
|
+
|
|
45
|
+
args = parser.parse_args(argv)
|
|
46
|
+
|
|
47
|
+
if args.command == "serve":
|
|
48
|
+
if not args.no_eager:
|
|
49
|
+
from .backend import get_backend
|
|
50
|
+
|
|
51
|
+
get_backend().preload_async()
|
|
52
|
+
if args.http:
|
|
53
|
+
from .http import serve_http
|
|
54
|
+
|
|
55
|
+
serve_http(host=args.host, port=args.port)
|
|
56
|
+
else:
|
|
57
|
+
from .server import serve
|
|
58
|
+
|
|
59
|
+
serve()
|
|
60
|
+
return 0
|
|
61
|
+
|
|
62
|
+
if args.command == "warmup":
|
|
63
|
+
import time
|
|
64
|
+
|
|
65
|
+
from .backend import get_backend
|
|
66
|
+
|
|
67
|
+
t0 = time.time()
|
|
68
|
+
backend = get_backend()
|
|
69
|
+
backend._load()
|
|
70
|
+
print(json.dumps({"device": backend.device, "model": MODEL_ID, "load_s": round(time.time() - t0, 1)}))
|
|
71
|
+
return 0
|
|
72
|
+
|
|
73
|
+
if args.command == "embed":
|
|
74
|
+
from . import service
|
|
75
|
+
|
|
76
|
+
print(json.dumps(service.embed_video(args.source), indent=2))
|
|
77
|
+
return 0
|
|
78
|
+
|
|
79
|
+
if args.command == "surprise":
|
|
80
|
+
from . import service
|
|
81
|
+
|
|
82
|
+
print(
|
|
83
|
+
json.dumps(
|
|
84
|
+
service.surprise(
|
|
85
|
+
args.source,
|
|
86
|
+
window=args.window,
|
|
87
|
+
stride=args.stride,
|
|
88
|
+
fps=args.fps,
|
|
89
|
+
start_s=args.start,
|
|
90
|
+
end_s=args.end,
|
|
91
|
+
),
|
|
92
|
+
indent=2,
|
|
93
|
+
)
|
|
94
|
+
)
|
|
95
|
+
return 0
|
|
96
|
+
|
|
97
|
+
if args.command == "info":
|
|
98
|
+
from .backend import pick_device
|
|
99
|
+
|
|
100
|
+
print(json.dumps({"version": __version__, "model": MODEL_ID, "device": pick_device()}, indent=2))
|
|
101
|
+
return 0
|
|
102
|
+
|
|
103
|
+
parser.print_help()
|
|
104
|
+
return 1
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
if __name__ == "__main__":
|
|
108
|
+
sys.exit(main())
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
"""Server-side latent store. Tools exchange opaque handle IDs, never tensors.
|
|
2
|
+
|
|
3
|
+
Latents are multi-MB tensors; MCP transports JSON. Handles keep the tensors in
|
|
4
|
+
process memory and give agents something cheap to pass around.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import secrets
|
|
10
|
+
from collections import OrderedDict
|
|
11
|
+
from dataclasses import dataclass, field
|
|
12
|
+
from typing import Any
|
|
13
|
+
|
|
14
|
+
MAX_HANDLES = 32
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
@dataclass
|
|
18
|
+
class Entry:
|
|
19
|
+
latents: Any # torch.Tensor (N, D) — mean-poolable latent sequence
|
|
20
|
+
pixels: Any = None # processed pixel tensor, kept only when predict() may need it
|
|
21
|
+
meta: dict = field(default_factory=dict)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class HandleStore:
|
|
25
|
+
def __init__(self, max_handles: int = MAX_HANDLES):
|
|
26
|
+
self._store: OrderedDict[str, Entry] = OrderedDict()
|
|
27
|
+
self._max = max_handles
|
|
28
|
+
|
|
29
|
+
def put(self, entry: Entry, prefix: str = "lat") -> str:
|
|
30
|
+
handle = f"{prefix}_{secrets.token_hex(4)}"
|
|
31
|
+
self._store[handle] = entry
|
|
32
|
+
while len(self._store) > self._max:
|
|
33
|
+
self._store.popitem(last=False) # evict oldest
|
|
34
|
+
return handle
|
|
35
|
+
|
|
36
|
+
def get(self, handle: str) -> Entry:
|
|
37
|
+
if handle not in self._store:
|
|
38
|
+
raise KeyError(
|
|
39
|
+
f"unknown handle: {handle!r} (it may have been evicted; "
|
|
40
|
+
f"the store keeps the {self._max} most recent)"
|
|
41
|
+
)
|
|
42
|
+
self._store.move_to_end(handle)
|
|
43
|
+
return self._store[handle]
|
|
44
|
+
|
|
45
|
+
def __len__(self) -> int:
|
|
46
|
+
return len(self._store)
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
STORE = HandleStore()
|