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.
@@ -0,0 +1,7 @@
1
+ __pycache__/
2
+ *.pyc
3
+ .venv/
4
+ dist/
5
+ *.egg-info/
6
+ .pytest_cache/
7
+ .DS_Store
@@ -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
+ ![An LLM agent consulting a video world model](demo.gif)
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
+ ![An LLM agent consulting a video world model](demo.gif)
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,5 @@
1
+ """wm — run a world model locally, and let any LLM agent call it as a tool."""
2
+
3
+ __version__ = "0.3.0"
4
+
5
+ MODEL_ID = "facebook/vjepa2-vitl-fpc64-256"
@@ -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()