Will Lin
1fe50c0092
The v2 top level had grown to ~27 dirs + 13 loose files — half of them
single-concept abstraction shells (memory/ 86 LOC, transport/ 222, parity/ 178,
extend/ 226) and half vendored fastvideo code sitting as peers to the actual v2
design. Regroup to mirror README section 3 "Planes & dependency order":
core/ enums+types, card, loop, program, parity, request, parallel
(the model-native contracts; no kernels)
runtime/ + folded-in substrate: cache, memory, transport, extend
(the import graph shows only runtime consumes them; compile/cudagraph
already lived here)
serving/ + deploy/ (products / fleet)
_vendor/ all copied fastvideo: models, layers, attention, configs, distributed,
platforms, api, hooks, logging_utils, third_party + fastvideo_args/
utils/logger/envs/forward_context — internal layout unchanged, still
mirrors upstream for diffing
28 dirs + 13 root files -> 8 dirs + 6 root files. Pure mechanical move: 947
absolute-import paths rewritten (v2.X -> v2.{core,runtime,serving,_vendor}.X),
boundary-anchored so platform/ (native dispatch) and platforms/ (vendored CUDA
detect) no longer collide and neither does hooks/ vs extend/. Deleted the empty
v2/loader/. README section 3 + 16 updated; stale test count corrected
(34 files/216 tests -> 22 files/143 tests).
Validated: `import v2` stays torch-free; `v2/run_tests.py` and `pytest v2/tests/`
-> 143 passed, 2 skipped on the numpy toy backend. (On a GPU box force the toy
backend with CUDA_VISIBLE_DEVICES="" or detect() picks cuda.) No external importer
changed — examples use the re-exported `from v2 import VideoGenerator`.
100 lines
4.6 KiB
Python
100 lines
4.6 KiB
Python
"""WorldModelSession — a long-lived interactive world-model rollout.
|
|
|
|
The realtime/interactive plane, made concrete. A one-shot ``engine.run`` generates a clip and forgets;
|
|
a *session* continues the world: each ``.act(action)`` resumes the causal ``chunk_rollout`` loop from
|
|
the **persistent world state** carried on the ``Session`` (its ``kv_handle``), streams frames as chunks
|
|
complete, and honors **step-boundary cancellation**. This exercises the seams a one-shot path never
|
|
touches — cross-request persistent context, the ``CHUNK_ROLLOUT`` loop driven incrementally, the
|
|
``Stream`` frame channel, and ``CancelScope`` — using the exact runtime primitives the engine uses
|
|
(``RuntimeLoopContext`` + ``LoopRunner``), so interleaving two sessions is still parity-safe.
|
|
|
|
Cancellation is transactional: the persistent world state advances only on a *completed* act, so a
|
|
cancelled rollout leaves the session resumable from the last good frame.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
from collections.abc import Callable
|
|
|
|
from v2.core.enums import ExecutionProfile
|
|
from v2.core.loop.driver import LoopRunner
|
|
from v2.recipes.common import text_encode_node_fn
|
|
from v2.core.request import DiffusionParams, TaskType, make_request
|
|
from v2.core.request.requests import Session
|
|
from v2.core.request.streams import Stream
|
|
from v2.runtime.context import RuntimeLoopContext
|
|
|
|
|
|
class WorldModelSession:
|
|
|
|
def __init__(self,
|
|
engine: Any,
|
|
model_id: str,
|
|
*,
|
|
session_id: str = "world",
|
|
window: int = 2,
|
|
loop_id: str = "chunk_rollout"):
|
|
self.engine = engine
|
|
self.model_id = model_id
|
|
self.window = window # how many prior chunks of world state to carry forward
|
|
self.loop_id = loop_id
|
|
self.session = Session(session_id=session_id)
|
|
self.stream = Stream(session_id)
|
|
self.session.streams["video"] = self.stream
|
|
self.world_context: list = [] # persistent world chunks (the cross-request KV)
|
|
self.session.kv_handle = self.world_context
|
|
self.step = 0
|
|
|
|
@property
|
|
def instance(self) -> Any:
|
|
if self.model_id not in self.engine._registry:
|
|
raise KeyError(f"no model {self.model_id!r} registered on engine")
|
|
return self.engine._registry[self.model_id][0]
|
|
|
|
def act(self,
|
|
action: str,
|
|
*,
|
|
seed: int | None = None,
|
|
step_hook: Callable[[WorldModelSession, int], None] | None = None) -> Any:
|
|
"""Advance the world by one action, continuing from persistent context; stream frames.
|
|
|
|
Raises ``request.Cancelled`` if cancelled at a step boundary — and because the world state is
|
|
only committed after the loop completes, a cancelled act leaves the session resumable."""
|
|
inst = self.instance
|
|
req = make_request(TaskType.V2W,
|
|
self.model_id,
|
|
action,
|
|
diffusion=DiffusionParams(seed=self.step if seed is None else seed))
|
|
slots: dict[str, Any] = {}
|
|
text_encode_node_fn(inst, slots, req, None)
|
|
slots["world_context"] = list(self.world_context) # thread the persistent state in
|
|
ctx = RuntimeLoopContext(inst,
|
|
observers=self.engine.observers,
|
|
interceptors=self.engine.interceptors,
|
|
slots=slots,
|
|
stream=self.stream,
|
|
cancel_scope=self.session.cancel_scope,
|
|
profile=ExecutionProfile.SERVE,
|
|
metrics={},
|
|
request_id=f"{self.session.session_id}.{self.step}")
|
|
runner = LoopRunner(inst.loop(self.loop_id), ctx, req, inst)
|
|
n = 0
|
|
while not runner.done:
|
|
if step_hook is not None:
|
|
step_hook(self, n) # test seam: may call self.cancel() at a boundary
|
|
runner.step() # peek() raises request.Cancelled if cancelled
|
|
n += 1
|
|
result = runner.result
|
|
new_chunks = list(result.outputs.get("chunks") or [])
|
|
self.world_context = (self.world_context + new_chunks)[-self.window:] # slide window; commit
|
|
self.session.kv_handle = self.world_context
|
|
self.step += 1
|
|
return result
|
|
|
|
def frames(self) -> list:
|
|
"""All streamed media chunks so far across the session (the realtime frame channel)."""
|
|
return [e for e in self.stream.events if e.type == "media.chunk"]
|
|
|
|
def cancel(self) -> None:
|
|
self.session.cancel_scope.cancel()
|