NVIDIA’s Cosmos-Framework Tutorial: Designing Colab-Friendly Miniature Cosmos 3 World Models with Omnimodal Mixture-of-Transformers

import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
torch.manual_seed(0)
@dataclass
class Cfg:
   d_model:   int = 192
   n_head:    int = 6
   n_layer:   int = 4
   ffn_mult:  int = 2
   n_mod:     int = 3
   text_vocab:int = 16
   vis_dim:   int = 8
   act_dim:   int = 4
   Lt:        int = 8
   Lv:        int = 8
   La:        int = 6
cfg = Cfg()
class RMSNorm(nn.Module):
   def __init__(self, d, eps=1e-6):
       super().__init__(); self.w = nn.Parameter(torch.ones(d)); self.eps = eps
   def forward(self, x):
       return self.w * x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def build_rope(T, hd, device, base=10000.0):
   pos  = torch.arange(T, device=device, dtype=torch.float32)[:, None]
   idx  = torch.arange(0, hd, 2, device=device, dtype=torch.float32)[None, :]
   freq = 1.0 / (base ** (idx / hd))
   ang  = pos * freq
   cos  = torch.cos(ang).repeat(1, 2)[None, None]
   sin  = torch.sin(ang).repeat(1, 2)[None, None]
   return cos, sin
def rotate_half(x):
   hd = x.shape[-1]; x1, x2 = x[..., :hd // 2], x[..., hd // 2:]
   return torch.cat([-x2, x1], -1)
def apply_rope(q, k, cos, sin):
   return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin
class Attention(nn.Module):
   """Shared cross-modal causal self-attention with rotary embeddings."""
   def __init__(self, c: Cfg):
       super().__init__()
       self.H, self.hd = c.n_head, c.d_model // c.n_head
       self.qkv  = nn.Linear(c.d_model, 3 * c.d_model, bias=False)
       self.proj = nn.Linear(c.d_model, c.d_model, bias=False)
   def forward(self, x, cos, sin, mask):
       B, T, D = x.shape
       q, k, v = self.qkv(x).chunk(3, -1)
       q = q.view(B, T, self.H, self.hd).transpose(1, 2)
       k = k.view(B, T, self.H, self.hd).transpose(1, 2)
       v = v.view(B, T, self.H, self.hd).transpose(1, 2)
       q, k = apply_rope(q, k, cos, sin)
       att = (q @ k.transpose(-2, -1)) / math.sqrt(self.hd)
       att = att.masked_fill(mask, float("-inf")).softmax(-1)
       o = (att @ v).transpose(1, 2).reshape(B, T, D)
       return self.proj(o)
class Expert(nn.Module):
   """A per-modality SwiGLU feed-forward 'transformer expert'."""
   def __init__(self, d, mult):
       super().__init__(); h = d * mult
       self.w1 = nn.Linear(d, h, bias=False)
       self.w3 = nn.Linear(d, h, bias=False)
       self.w2 = nn.Linear(h, d, bias=False)
   def forward(self, x):
       return self.w2(F.silu(self.w1(x)) * self.w3(x))
class MoTBlock(nn.Module):
   """Shared attention + Mixture-of-Transformers (per-modality expert) routing."""
   def __init__(self, c: Cfg):
       super().__init__()
       self.attn_norm = RMSNorm(c.d_model)
       self.attn      = Attention(c)
       self.ffn_norm  = nn.ModuleList([RMSNorm(c.d_model) for _ in range(c.n_mod)])
       self.experts   = nn.ModuleList([Expert(c.d_model, c.ffn_mult) for _ in range(c.n_mod)])
   def forward(self, x, cos, sin, mask, mod_id):
       x = x + self.attn(self.attn_norm(x), cos, sin, mask)
       out = torch.zeros_like(x)
       for i, exp in enumerate(self.experts):
           sel = (mod_id == i).view(1, -1, 1).to(x.dtype)
           out = out + sel * exp(self.ffn_norm[i](x))
       return x + out
class OmniMoT(nn.Module):
   def __init__(self, c: Cfg):
       super().__init__(); self.c = c
       self.text_emb = nn.Embedding(c.text_vocab, c.d_model)
       self.vis_in   = nn.Linear(c.vis_dim, c.d_model)
       self.act_in   = nn.Linear(c.act_dim, c.d_model)
       self.mod_emb  = nn.Embedding(c.n_mod, c.d_model)
       self.blocks   = nn.ModuleList([MoTBlock(c) for _ in range(c.n_layer)])
       self.norm     = RMSNorm(c.d_model)
       self.text_head = nn.Linear(c.d_model, c.text_vocab, bias=False)
       self.vis_head  = nn.Linear(c.d_model, c.vis_dim,  bias=False)
       self.act_head  = nn.Linear(c.d_model, c.act_dim,  bias=False)
       ids = torch.cat([torch.zeros(c.Lt), torch.ones(c.Lv), torch.full((c.La,), 2)]).long()
       self.register_buffer("mod_id", ids, persistent=False)
   def forward(self, text, vis, act):
       c = self.c
       x = torch.cat([self.text_emb(text), self.vis_in(vis), self.act_in(act)], 1)
       x = x + self.mod_emb(self.mod_id)[None]
       B, T, D = x.shape
       cos, sin = build_rope(T, D // c.n_head, x.device)
       mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)[None, None]
       for blk in self.blocks:
           x = blk(x, cos, sin, mask, self.mod_id)
       x = self.norm(x)
       ht = self.text_head(x[:, :c.Lt])
       hv = self.vis_head(x[:, c.Lt:c.Lt + c.Lv])
       ha = self.act_head(x[:, c.Lt + c.Lv:])
       return ht, hv, ha
model = OmniMoT(cfg).to(DEVICE)
n_params = sum(p.numel() for p in model.parameters())
print(f"Model built: OmniMoT  |  {n_params/1e6:.2f}M params  |  {cfg.n_layer} MoT blocks "
     f"x {cfg.n_mod} experts  |  device={DEVICE}")