""" Cosmos-Predict2.5-2B world model wrapper. Installation ------------ Option A — cosmos-framework (NVIDIA's own library): git clone https://github.com/NVIDIA/cosmos-framework pip install -e cosmos-framework huggingface-cli login # accept gated license on the HF model page Option B — diffusers (lighter, covers Cosmos 1.0 + Predict2 variants): pip install "diffusers>=0.33" pip install av imageio imageio-ffmpeg NOTE: When running alongside Cosmos-Policy, the policy already does its own video encoding internally. This world model is mainly used to: 1. Provide a latent context vector for other components. 2. Run "imagination" (predict future frames) for planning. Stub mode runs on CPU with no weights. """ from __future__ import annotations import logging from typing import Optional import numpy as np from models.base import BaseWorldModel log = logging.getLogger(__name__) # ─────────────────────────────────────────────────────────────── stub class _StubWorldModel(BaseWorldModel): def __init__(self, latent_dim: int, rng: np.random.Generator): self._dim = latent_dim self._rng = rng def encode(self, frames: np.ndarray) -> np.ndarray: """frames: (T, H, W, 3) float32 [0,1]""" mean_rgb = frames.mean(axis=(0, 1, 2)) # (3,) emb = np.zeros(self._dim, dtype=np.float32) for i, v in enumerate(mean_rgb): emb[i::3] = v emb += self._rng.normal(0, 0.01, self._dim).astype(np.float32) return emb / (np.linalg.norm(emb) + 1e-8) def predict_future( self, frames: np.ndarray, actions: Optional[np.ndarray] = None, ) -> np.ndarray: last = frames[-1:] if len(frames) else np.zeros((1, 224, 224, 3), np.float32) return np.repeat(last, 4, axis=0) # ─────────────────────────────────────────────────────────────── cosmos-framework class _CosmosFrameworkWorldModel(BaseWorldModel): """ Uses NVIDIA's cosmos-framework package. Install: git clone https://github.com/NVIDIA/cosmos-framework && pip install -e cosmos-framework """ def __init__(self, model_id: str, device: str, dtype_str: str, cfg: dict): import torch from cosmos_framework.models.video.predict2 import CosmosVideoPredict2 self._device = device self._dtype = getattr(torch, dtype_str, torch.bfloat16) self._latent_dim = cfg.get("latent_dim", 512) log.info("Loading world model via cosmos-framework: %s", model_id) self._model = CosmosVideoPredict2.from_pretrained( model_id, torch_dtype=self._dtype, device_map=device, ).eval() log.info("World model loaded.") def encode(self, frames: np.ndarray) -> np.ndarray: import torch # (T, H, W, 3) → (1, T, 3, H, W) t = torch.from_numpy(frames).permute(0, 3, 1, 2).unsqueeze(0).to( self._device, self._dtype ) with torch.no_grad(): out = self._model.encode_video(t) emb = out.last_hidden_state.float().mean(dim=1).squeeze(0).cpu().numpy() return self._fit_dim(emb) def predict_future( self, frames: np.ndarray, actions: Optional[np.ndarray] = None, ) -> np.ndarray: import torch t = torch.from_numpy(frames).permute(0, 3, 1, 2).unsqueeze(0).to( self._device, self._dtype ) action_t = None if actions is not None: action_t = torch.tensor(actions, dtype=self._dtype, device=self._device) with torch.no_grad(): pred = self._model.predict_future(t, action_conditioning=action_t, num_frames=4) return pred.frames[0].permute(0, 2, 3, 1).float().clamp(0, 1).cpu().numpy() def _fit_dim(self, emb: np.ndarray) -> np.ndarray: d = self._latent_dim if emb.shape[0] > d: emb = emb[:d] elif emb.shape[0] < d: emb = np.pad(emb, (0, d - emb.shape[0])) return (emb / (np.linalg.norm(emb) + 1e-8)).astype(np.float32) # ─────────────────────────────────────────────────────────────── diffusers class _DiffusersWorldModel(BaseWorldModel): """ Uses HuggingFace diffusers. Install: pip install "diffusers>=0.33" Provides video-to-latent encoding via the VAE; prediction via the full pipeline. Limited: no direct action conditioning (use cosmos-framework for that). """ def __init__(self, model_id: str, device: str, dtype_str: str, cfg: dict): import torch from diffusers import AutoPipelineForVideo2Video self._device = device self._dtype = getattr(torch, dtype_str, torch.bfloat16) self._latent_dim = cfg.get("latent_dim", 512) log.info("Loading world model via diffusers: %s", model_id) self._pipe = AutoPipelineForVideo2Video.from_pretrained( model_id, torch_dtype=self._dtype, ).to(device) log.info("World model loaded (diffusers).") def encode(self, frames: np.ndarray) -> np.ndarray: import torch # Use VAE encoder on the last frame as a compact world state proxy last = frames[-1] # (H, W, 3) t = torch.from_numpy(last).permute(2, 0, 1).unsqueeze(0).to( self._device, self._dtype ) * 2 - 1 # [0,1] → [-1,1] with torch.no_grad(): z = self._pipe.vae.encode(t).latent_dist.mean emb = z.float().flatten().cpu().numpy() d = self._latent_dim if emb.shape[0] > d: emb = emb[:d] elif emb.shape[0] < d: emb = np.pad(emb, (0, d - emb.shape[0])) return (emb / (np.linalg.norm(emb) + 1e-8)).astype(np.float32) def predict_future( self, frames: np.ndarray, actions: Optional[np.ndarray] = None, ) -> np.ndarray: from PIL import Image pil = Image.fromarray((frames[-1] * 255).astype(np.uint8)) result = self._pipe(image=pil, num_frames=4, output_type="np") return result.frames[0] # (4, H, W, 3) float32 # ─────────────────────────────────────────────────────────────── factory class CosmosWorldModel: """ Factory — tries loaders in order: 1. cosmos-framework (full feature set, action conditioning) 2. diffusers (lighter, video encoding + generation) 3. raises RuntimeError with install instructions """ def __new__(cls, config: dict) -> BaseWorldModel: # type: ignore[misc] if config.get("stub", True): log.info("WorldModel: STUB mode") return _StubWorldModel(config.get("latent_dim", 512), np.random.default_rng(42)) model_id = config["model_id"] device = config.get("device", "cuda") dtype_str = config.get("dtype", "bfloat16") # Try cosmos-framework try: import cosmos_framework # noqa: F401 return _CosmosFrameworkWorldModel(model_id, device, dtype_str, config) except ImportError: log.warning( "cosmos-framework not found — trying diffusers.\n" "For full Cosmos-Predict2.5 support (incl. action conditioning):\n" " git clone https://github.com/NVIDIA/cosmos-framework\n" " pip install -e cosmos-framework" ) # Try diffusers try: import diffusers # noqa: F401 return _DiffusersWorldModel(model_id, device, dtype_str, config) except ImportError: pass raise RuntimeError( f"Cannot load world model '{model_id}'.\n" "Install one of:\n" " A) git clone https://github.com/NVIDIA/cosmos-framework && pip install -e cosmos-framework\n" " B) pip install 'diffusers>=0.33'\n" "Or run with --stub for development." )