worldkernels 0.1.0__py3-none-any.whl

worldkernels/__init__.py

@@ -0,0 +1,35 @@
+"""
+WorldKernels - GPU-first world model simulation engine.
+
+Serve learned world models (DiT, VAE) as interactive sessions.
+
+Example:
+    >>> from worldkernels import WorldKernel, Action, WorldConfig
+    >>> wk = WorldKernel(device="cuda")
+    >>> wk.load_world("Etched/oasis-500m")
+    >>> session = wk.create_session(world="oasis-500m", config=WorldConfig())
+    >>> obs = session.step(Action("keyboard", {"keys": ["W"]}))
+"""
+
+from worldkernels.core.action import Action
+from worldkernels.core.config import ServerConfig, WorldConfig
+from worldkernels.core.engine import WorldKernel
+from worldkernels.core.observation import Observation
+from worldkernels.core.session import Session, SessionStatus
+
+try:
+    from worldkernels._version import __version__, __version_tuple__
+except ImportError:
+    __version__ = "0.1.0.dev0"
+    __version_tuple__ = (0, 1, 0, "dev0")
+
+__all__ = [
+    "WorldKernel",
+    "Session",
+    "SessionStatus",
+    "Action",
+    "Observation",
+    "WorldConfig",
+    "ServerConfig",
+    "__version__",
+]

worldkernels/_version.py

@@ -0,0 +1,3 @@
+# Version of the worldkernels package.
+__version__ = "0.1.0"
+__version_tuple__ = (0, 1, 0)

worldkernels/cli/main.py

@@ -0,0 +1,44 @@
+"""WorldKernels CLI."""
+
+from __future__ import annotations
+
+import sys
+
+
+def app() -> None:
+    """Main CLI entry point."""
+    args = sys.argv[1:]
+
+    if not args or args[0] in ("-h", "--help", "help"):
+        print(
+            """worldkernels - GPU-first world model simulation engine
+
+Usage:
+    worldkernels serve [--host HOST] [--port PORT]   Start HTTP/WebSocket server
+    worldkernels version                             Show version
+    worldkernels --help                              Show this help
+
+Examples:
+    worldkernels serve --host 0.0.0.0 --port 8000
+"""
+        )
+        return
+
+    if args[0] in ("version", "--version", "-V"):
+        from worldkernels import __version__
+
+        print(f"worldkernels {__version__}")
+        return
+
+    if args[0] == "serve":
+        print("Server not yet implemented. Coming soon!")
+        print("Install with: pip install worldkernels[serve]")
+        return
+
+    print(f"Unknown command: {args[0]}")
+    print("Run 'worldkernels --help' for usage.")
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    app()

worldkernels/core/action.py

@@ -0,0 +1,34 @@
+"""Action schema for world model inputs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class Action:
+    """
+    Generalized action container.
+
+    Not RL-specific — just "control input to world".
+
+    Examples:
+        >>> Action("keyboard", {"keys": ["W", "SPACE"]})
+        >>> Action("continuous", {"velocity": [0.1, 0.0, 0.5]})
+        >>> Action("camera", {"trajectory": "orbit_left_30"})
+        >>> Action("text", {"command": "open the door"})
+    """
+
+    action_type: str
+    """Type of action, e.g., 'keyboard', 'continuous', 'text', 'pose'."""
+
+    payload: dict[str, Any] = field(default_factory=dict)
+    """Schema depends on action_type."""
+
+    timestamp: float | None = None
+    """Optional timing information."""
+
+    def __post_init__(self) -> None:
+        if not self.action_type:
+            raise ValueError("action_type cannot be empty")

worldkernels/core/config.py

@@ -0,0 +1,42 @@
+"""Configuration dataclasses for WorldKernel."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+
+@dataclass
+class WorldConfig:
+    """Per-session generation parameters."""
+
+    # Resolution / format
+    height: int = 480
+    width: int = 848
+    fps: int = 24
+
+    # Generation quality
+    num_inference_steps: int = 4
+    guidance_scale: float = 1.0
+
+    # Chunking (how many frames per step call)
+    frames_per_step: int = 8
+
+    # Conditioning
+    initial_prompt: str | None = None
+    initial_image: str | None = None  # path or base64
+
+    # Compute budget
+    max_vram_gb: float | None = None
+    precision: Literal["bf16", "fp16", "fp32"] = "bf16"
+
+
+@dataclass
+class ServerConfig:
+    """Configuration for the HTTP/WebSocket server."""
+
+    host: str = "0.0.0.0"
+    port: int = 8000
+    max_sessions: int = 4
+    api_key: str | None = None
+    cors_origins: list[str] = field(default_factory=lambda: ["*"])

worldkernels/core/engine.py

@@ -0,0 +1,113 @@
+"""WorldKernel main engine class."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from worldkernels.core.config import WorldConfig
+    from worldkernels.core.session import Session
+
+
+class WorldKernel:
+    """
+    GPU-first world model simulation engine.
+
+    This is the main entry point for the WorldKernels library.
+    It manages world model loading, session creation, and resource lifecycle.
+
+    Example:
+        >>> from worldkernels import WorldKernel, WorldConfig
+        >>> wk = WorldKernel(device="cuda")
+        >>> wk.load_world("Etched/oasis-500m")
+        >>> session = wk.create_session(world="oasis-500m", config=WorldConfig())
+        >>> obs = session.step(Action("keyboard", {"keys": ["W"]}))
+    """
+
+    def __init__(
+        self,
+        device: str = "cuda",
+        max_sessions: int = 4,
+        offload_idle: bool = True,
+    ) -> None:
+        """
+        Initialize the WorldKernel engine.
+
+        Args:
+            device: Device to run on ('cuda', 'cuda:0', 'cpu').
+            max_sessions: Maximum concurrent sessions.
+            offload_idle: Whether to offload paused sessions to CPU.
+        """
+        self.device = device
+        self.max_sessions = max_sessions
+        self.offload_idle = offload_idle
+        self._worlds: dict[str, object] = {}
+        self._sessions: dict[str, Session] = {}
+
+    def load_world(
+        self,
+        model_id: str,
+        alias: str | None = None,
+        trust_remote_code: bool = False,
+    ) -> None:
+        """
+        Load a world model from HuggingFace Hub or local path.
+
+        Args:
+            model_id: HuggingFace model ID (e.g., 'Etched/oasis-500m') or local path.
+            alias: Optional alias for referencing this world.
+            trust_remote_code: Whether to trust remote code from the model repo.
+        """
+        # Stub implementation
+        key = alias or model_id.split("/")[-1]
+        self._worlds[key] = {"model_id": model_id, "loaded": False}
+
+    def create_session(
+        self,
+        world: str,
+        config: WorldConfig | None = None,
+        seed: int | None = None,
+    ) -> Session:
+        """
+        Create a new simulation session.
+
+        Args:
+            world: Name/alias of the loaded world model.
+            config: Session configuration.
+            seed: Random seed for reproducibility.
+
+        Returns:
+            A new Session instance.
+        """
+        from worldkernels.core.config import WorldConfig as WC
+        from worldkernels.core.session import Session
+
+        if world not in self._worlds:
+            raise ValueError(f"World '{world}' not loaded. Call load_world() first.")
+
+        if len(self._sessions) >= self.max_sessions:
+            raise RuntimeError(f"Max sessions ({self.max_sessions}) reached.")
+
+        cfg = config or WC()
+        session = Session(world_id=world, config=cfg, seed=seed or 0)
+        self._sessions[session.session_id] = session
+        return session
+
+    def get_session(self, session_id: str) -> Session | None:
+        """Get a session by ID."""
+        return self._sessions.get(session_id)
+
+    def list_sessions(self) -> list[str]:
+        """List all active session IDs."""
+        return list(self._sessions.keys())
+
+    def list_worlds(self) -> list[str]:
+        """List all loaded world model aliases."""
+        return list(self._worlds.keys())
+
+    def shutdown(self) -> None:
+        """Cleanup all sessions and resources."""
+        for session in self._sessions.values():
+            session.close()
+        self._sessions.clear()
+        self._worlds.clear()

worldkernels/core/observation.py

@@ -0,0 +1,39 @@
+"""Observation types returned by world models."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class Observation:
+    """
+    Multi-modal observation output from a world model step.
+
+    This is NOT RL-specific (no reward/done by default).
+    """
+
+    step_index: int
+    """Which step this observation corresponds to."""
+
+    generation_time_ms: float
+    """Time taken to generate this observation."""
+
+    frames: list[bytes] | None = None
+    """RGB video frames (encoded or raw numpy bytes)."""
+
+    latent: bytes | None = None
+    """Optional: return latent state for chaining."""
+
+    audio: bytes | None = None
+    """Optional: audio track."""
+
+    depth: bytes | None = None
+    """Optional: depth map."""
+
+    segmentation: bytes | None = None
+    """Optional: semantic segmentation."""
+
+    structured: dict[str, Any] | None = field(default_factory=dict)
+    """Optional: structured state (positions, velocities, etc.)."""

worldkernels/core/session.py

@@ -0,0 +1,95 @@
+"""Session lifecycle management."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from worldkernels.core.action import Action
+    from worldkernels.core.config import WorldConfig
+    from worldkernels.core.observation import Observation
+
+
+class SessionStatus(str, Enum):
+    """Lifecycle status of a session."""
+
+    ACTIVE = "active"
+    PAUSED = "paused"
+    TERMINATED = "terminated"
+
+
+@dataclass
+class LatentState:
+    """Opaque tensor bundle representing the current world state."""
+
+    data: Any = None  # Will hold GPU tensors when implemented
+    device: str = "cpu"
+
+
+@dataclass
+class Session:
+    """
+    A stateful simulation session.
+
+    Represents an interactive connection to a running world model.
+    """
+
+    session_id: str = field(default_factory=lambda: f"sess_{uuid.uuid4().hex[:12]}")
+    """Unique identifier for this session."""
+
+    world_id: str = ""
+    """Which world model this session uses (registry key)."""
+
+    config: WorldConfig | None = None
+    """Session configuration (resolution, fps, etc.)."""
+
+    # Latent state (kept on GPU or offloaded)
+    state: LatentState = field(default_factory=LatentState)
+    """Current latent state of the simulation."""
+
+    step_index: int = 0
+    """How many steps have been taken."""
+
+    # Provenance
+    seed: int = 0
+    action_history: list[Action] = field(default_factory=list)
+    """History of actions for replay/debug."""
+
+    parent_session_id: str | None = None
+    """If this session was branched from another."""
+
+    # Lifecycle
+    created_at: datetime = field(default_factory=datetime.now)
+    last_active_at: datetime = field(default_factory=datetime.now)
+    status: SessionStatus = SessionStatus.ACTIVE
+
+    # --- Methods (stubs for now) ---
+
+    def step(self, action: Action) -> Observation:
+        """Execute one step with the given action."""
+        raise NotImplementedError("Session.step() not yet implemented")
+
+    def checkpoint(self) -> str:
+        """Save current state, return checkpoint ID."""
+        raise NotImplementedError("Session.checkpoint() not yet implemented")
+
+    def restore(self, checkpoint_id: str) -> None:
+        """Restore session to a saved checkpoint."""
+        raise NotImplementedError("Session.restore() not yet implemented")
+
+    def branch(self) -> Session:
+        """Clone this session's state into a new session."""
+        raise NotImplementedError("Session.branch() not yet implemented")
+
+    def close(self) -> None:
+        """Terminate this session and release resources."""
+        self.status = SessionStatus.TERMINATED
+
+    @property
+    def id(self) -> str:
+        """Alias for session_id."""
+        return self.session_id