The Net

.gitattributes

@@ -0,0 +1,3 @@
+* text=auto
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.st filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 JacobLinCool
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,63 @@
+---
+library_name: transformers
+license: mit
+tags:
+- custom_code
+- multimodal
+- mixture-of-experts
+- weight-sharing
+---
+
+# The Net
+
+20T-scale multimodal architecture. Text, audio, images, video and EEG. 686 attention/FFN blocks, a three-level sparse MoE core with 128 experts per layer and top-2 routing, and recurrent memory. PyTorch implementation and checkpoint files released under the MIT license.
+
+## Architecture
+
+The Net has **11 major components**: five modality encoders, multiscale fusion, a shared world model, three media decoders and a memory controller.
+
+```text
+Text · Audio · Images · Video · EEG
+                ↓
+     Modality-specific encoders
+                ↓
+       Multiscale fusion
+                ↓
+   Hierarchical sparse MoE core ↔ Memory
+                ↓
+ Text head · Audio / Image / Video decoders
+```
+
+Encoders combine modality-specific stems with learned chunking, pooling and dechunking. Fusion aligns their representations before the shared core. The core uses attention, SwiGLU experts and RMSNorm across widths of 8,192, 12,288 and 16,384. Residual detail paths connect hierarchy levels; recurrent memory carries state between segments.
+
+| Modality | Interface |
+| --- | --- |
+| Text | Byte-level tokens and next-token logits |
+| Audio | Waveform frames, spectral features and codebook logits |
+| Images | RGB patches and image latents |
+| Video | Spatiotemporal tubelets and video latents |
+| EEG | Multichannel signals with electrode adjacency |
+
+## Implementation
+
+The repository includes the PyTorch architecture, Transformers configuration,
+tokenizer and checkpoint files under the MIT license.
+
+| File | Contents |
+| --- | --- |
+| `config.json` | Model configuration |
+| `configuration_omni.py` | `OmniConfig` |
+| `modeling_omni.py` | `OmniModel` module implementation |
+| `tokenizer.json` | Byte-level tokenizer |
+| `tensor_shapes.json` | Architecture tensor names and dimensions |
+
+## FAQ
+
+**Q: Which modalities does the architecture cover?**  
+Text, audio, images, video and EEG, connected through a shared hierarchical core.
+
+**Q: How are the experts organized?**  
+The shared core spans three hierarchy levels, with 128 experts per layer and top-2 routing. Multiscale fusion connects the modality encoders to the core; media decoders and a text head provide the output interfaces.
+
+**Q: What is this release?**  
+A **Safetensors & Xet storage experiment** with constant-initialized, shared weights.

__init__.py

@@ -0,0 +1,4 @@
+from .configuration_omni import OmniConfig
+from .modeling_omni import OmniModel
+
+__all__ = ["OmniConfig", "OmniModel"]

architecture.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aba8d5f96cf32f5945cb1edc7cfb94fbfc95b5759daa3208e35a8a0844ea859d
+size 115368746

config.json

@@ -0,0 +1,245 @@
+{
+  "architectures": [
+    "OmniModel"
+  ],
+  "audio_codebook_size": 1024,
+  "audio_codebooks": 8,
+  "audio_frame_size": 320,
+  "auto_map": {
+    "AutoConfig": "configuration_omni.OmniConfig",
+    "AutoModel": "modeling_omni.OmniModel"
+  },
+  "checkpoint_format": "the-net-shared-aliases-v2",
+  "context_dim": 256,
+  "eeg_channels": 64,
+  "head_dim": 128,
+  "initialization": "constant, untrained",
+  "model_type": "the_net_omni",
+  "patch_size": 4,
+  "stages": {
+    "audio_decoder": [
+      [
+        4096,
+        12,
+        1,
+        4,
+        1
+      ],
+      [
+        8192,
+        84,
+        8,
+        2,
+        2
+      ]
+    ],
+    "audio_encoder": [
+      [
+        4096,
+        12,
+        1,
+        4,
+        1
+      ],
+      [
+        8192,
+        24,
+        1,
+        4,
+        1
+      ],
+      [
+        12288,
+        12,
+        1,
+        4,
+        1
+      ]
+    ],
+    "eeg_encoder": [
+      [
+        2048,
+        6,
+        1,
+        4,
+        1
+      ],
+      [
+        4096,
+        12,
+        1,
+        4,
+        1
+      ],
+      [
+        8192,
+        6,
+        1,
+        4,
+        1
+      ]
+    ],
+    "fusion": [
+      [
+        16384,
+        30,
+        8,
+        3,
+        2
+      ]
+    ],
+    "image_decoder": [
+      [
+        4096,
+        8,
+        1,
+        4,
+        1
+      ],
+      [
+        12288,
+        52,
+        16,
+        2,
+        2
+      ]
+    ],
+    "image_encoder": [
+      [
+        4096,
+        8,
+        1,
+        4,
+        1
+      ],
+      [
+        8192,
+        24,
+        1,
+        4,
+        1
+      ],
+      [
+        12288,
+        24,
+        1,
+        4,
+        1
+      ]
+    ],
+    "memory_control": [
+      [
+        8192,
+        38,
+        32,
+        2,
+        2
+      ]
+    ],
+    "text_encoder": [
+      [
+        2048,
+        6,
+        1,
+        4,
+        1
+      ],
+      [
+        4096,
+        12,
+        1,
+        4,
+        1
+      ],
+      [
+        8192,
+        6,
+        1,
+        4,
+        1
+      ]
+    ],
+    "video_decoder": [
+      [
+        8192,
+        8,
+        1,
+        4,
+        1
+      ],
+      [
+        16384,
+        88,
+        16,
+        2,
+        2
+      ]
+    ],
+    "video_encoder": [
+      [
+        4096,
+        8,
+        1,
+        4,
+        1
+      ],
+      [
+        8192,
+        32,
+        1,
+        4,
+        1
+      ],
+      [
+        16384,
+        44,
+        1,
+        4,
+        1
+      ]
+    ],
+    "world_model": [
+      [
+        8192,
+        48,
+        128,
+        2,
+        2
+      ],
+      [
+        12288,
+        48,
+        128,
+        2,
+        2
+      ],
+      [
+        16384,
+        34,
+        128,
+        2,
+        2
+      ]
+    ]
+  },
+  "transformers_version": "4.57.6",
+  "vocab_size": 259,
+  "quantization_config": {
+    "quant_method": "compressed-tensors",
+    "format": "pack-quantized",
+    "config_groups": {
+      "shared_bank": {
+        "targets": [
+          "shared_bank"
+        ],
+        "weights": {
+          "num_bits": 1,
+          "type": "int",
+          "symmetric": false,
+          "strategy": "tensor",
+          "dynamic": false
+        }
+      }
+    }
+  }
+}

configuration_omni.py

@@ -0,0 +1,94 @@
+"""Concrete heterogeneous omni architecture; all parameter counts derive from its modules."""
+
+from copy import deepcopy
+
+from transformers import PretrainedConfig
+
+# (width, block count, expert count, FFN expansion, selected experts)
+STAGES = {
+    "text_encoder": [(2048, 6, 1, 4, 1), (4096, 12, 1, 4, 1), (8192, 6, 1, 4, 1)],
+    "audio_encoder": [(4096, 12, 1, 4, 1), (8192, 24, 1, 4, 1), (12288, 12, 1, 4, 1)],
+    "image_encoder": [(4096, 8, 1, 4, 1), (8192, 24, 1, 4, 1), (12288, 24, 1, 4, 1)],
+    "video_encoder": [(4096, 8, 1, 4, 1), (8192, 32, 1, 4, 1), (16384, 44, 1, 4, 1)],
+    "eeg_encoder": [(2048, 6, 1, 4, 1), (4096, 12, 1, 4, 1), (8192, 6, 1, 4, 1)],
+    "fusion": [(16384, 30, 8, 3, 2)],
+    "world_model": [(8192, 48, 128, 2, 2), (12288, 48, 128, 2, 2), (16384, 34, 128, 2, 2)],
+    "audio_decoder": [(4096, 12, 1, 4, 1), (8192, 84, 8, 2, 2)],
+    "image_decoder": [(4096, 8, 1, 4, 1), (12288, 52, 16, 2, 2)],
+    "video_decoder": [(8192, 8, 1, 4, 1), (16384, 88, 16, 2, 2)],
+    "memory_control": [(8192, 38, 32, 2, 2)],
+}
+
+
+class OmniConfig(PretrainedConfig):
+    model_type = "the_net_omni"
+
+    def __init__(
+        self,
+        stages=None,
+        head_dim=128,
+        vocab_size=259,
+        eeg_channels=64,
+        context_dim=256,
+        audio_frame_size=320,
+        patch_size=4,
+        audio_codebooks=8,
+        audio_codebook_size=1024,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.stages = deepcopy(STAGES if stages is None else stages)
+        self.head_dim = head_dim
+        self.vocab_size = vocab_size
+        self.eeg_channels = eeg_channels
+        self.context_dim = context_dim
+        self.audio_frame_size = audio_frame_size
+        self.patch_size = patch_size
+        self.audio_codebooks = audio_codebooks
+        self.audio_codebook_size = audio_codebook_size
+        for name in (
+            "head_dim",
+            "vocab_size",
+            "eeg_channels",
+            "context_dim",
+            "audio_frame_size",
+            "patch_size",
+            "audio_codebooks",
+            "audio_codebook_size",
+        ):
+            value = getattr(self, name)
+            if type(value) is not int or value <= 0:
+                raise ValueError(f"{name} must be a positive integer")
+        if set(self.stages) != set(STAGES):
+            raise ValueError("All eleven named components are required")
+        for component, rows in self.stages.items():
+            if not rows:
+                raise ValueError(f"Empty component: {component}")
+            for row in rows:
+                if len(row) != 5 or any(type(n) is not int or n < 1 for n in row):
+                    raise ValueError(f"Invalid stage: {row}")
+                width, _, experts, _, selected = row
+                if width % head_dim or selected > experts:
+                    raise ValueError(f"Invalid heads or routing: {row}")
+
+    @classmethod
+    def tiny(cls):
+        stages = {
+            name: [
+                (32 * (i + 1), 1, min(experts, 2), expansion, min(selected, 2))
+                for i, (_, _, experts, expansion, selected) in enumerate(rows)
+            ]
+            for name, rows in STAGES.items()
+        }
+        return cls(
+            stages=stages,
+            head_dim=16,
+            context_dim=32,
+            audio_frame_size=32,
+            eeg_channels=8,
+            audio_codebooks=2,
+            audio_codebook_size=16,
+        )
+
+
+OmniConfig.register_for_auto_class()

constant_checkpoint.py

@@ -0,0 +1,58 @@
+"""Constant-tensor checkpoint reader."""
+
+import json
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+from .configuration_omni import OmniConfig
+from .modeling_omni import OmniModel
+
+
+class ConstantCheckpoint:
+    def __init__(self, directory):
+        self.directory = Path(directory)
+        self.shapes = json.loads((self.directory / "tensor_shapes.json").read_text())
+        self.values = load_file(self.directory / "initialization.st")
+        if self.shapes.keys() != self.values.keys():
+            raise ValueError("Shape and value keys differ")
+        self.config = OmniConfig.from_pretrained(self.directory, local_files_only=True)
+
+    def slice(self, name, start, length):
+        size = 1
+        for dimension in self.shapes[name]:
+            size *= dimension
+        if (
+            type(start) is not int
+            or type(length) is not int
+            or start < 0
+            or length < 0
+            or start + length > size
+        ):
+            raise ValueError("Invalid tensor slice")
+        if length > 1_000_000:
+            raise ValueError("Read at most one million elements per slice")
+        return self.values[name].expand(length).clone()
+
+    def meta_model(self):
+        with torch.device("meta"):
+            model = OmniModel(self.config)
+        parameters = dict(model.named_parameters())
+        if parameters.keys() != self.shapes.keys():
+            raise ValueError("Checkpoint and model parameter keys differ")
+        for name, parameter in parameters.items():
+            if list(parameter.shape) != self.shapes[name]:
+                raise ValueError(f"Shape mismatch: {name}")
+        return model
+
+    def materialize(self, max_parameters=10_000_000):
+        model = self.meta_model()
+        total = sum(p.numel() for p in model.parameters())
+        if total > max_parameters:
+            raise ValueError(f"{total:,} parameters exceed the materialization limit")
+        model.to_empty(device="cpu")
+        with torch.no_grad():
+            for name, parameter in model.named_parameters():
+                parameter.fill_(self.values[name].item())
+        return model

initialization.st

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4d29dea3bfe42eb4de7a144b1afbe4b451a79d6009c80f235843c435db8b0d0
+size 8768280

modeling_omni.py

@@ -0,0 +1,391 @@
+"""PyTorch implementation of the heterogeneous hierarchical omni architecture."""
+
+import itertools
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration_omni import OmniConfig
+
+
+class Attention(nn.Module):
+    def __init__(self, width, head_dim):
+        super().__init__()
+        self.head_dim = head_dim
+        self.q_proj = nn.Linear(width, width, bias=False)
+        self.k_proj = nn.Linear(width, width, bias=False)
+        self.v_proj = nn.Linear(width, width, bias=False)
+        self.o_proj = nn.Linear(width, width, bias=False)
+
+    def forward(self, x, causal=False):
+        b, s, d = x.shape
+
+        def heads(value):
+            return value.reshape(b, s, d // self.head_dim, self.head_dim).transpose(1, 2)
+
+        q, k, v = (heads(layer(x)) for layer in (self.q_proj, self.k_proj, self.v_proj))
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
+        return self.o_proj(y.transpose(1, 2).reshape_as(x))
+
+
+class Expert(nn.Module):
+    def __init__(self, width, intermediate):
+        super().__init__()
+        self.gate_proj = nn.Linear(width, intermediate, bias=False)
+        self.up_proj = nn.Linear(width, intermediate, bias=False)
+        self.down_proj = nn.Linear(intermediate, width, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class SparseMLP(nn.Module):
+    def __init__(self, width, experts, expansion, selected):
+        super().__init__()
+        self.selected = selected
+        self.router = nn.Linear(width, experts, bias=False) if experts > 1 else None
+        self.experts = nn.ModuleList([Expert(width, width * expansion) for _ in range(experts)])
+
+    def forward(self, x):
+        if self.router is None:
+            return self.experts[0](x)
+        flat = x.reshape(-1, x.shape[-1])
+        scores, ids = self.router(flat).topk(self.selected, dim=-1)
+        scores = scores.softmax(-1)
+        result = torch.zeros_like(flat)
+        for expert_id, expert in enumerate(self.experts):
+            tokens, slots = torch.where(ids == expert_id)
+            if tokens.numel():
+                result = result.index_add(
+                    0, tokens, expert(flat[tokens]) * scores[tokens, slots, None]
+                )
+        return result.reshape_as(x)
+
+
+class Block(nn.Module):
+    def __init__(self, width, experts, expansion, selected, head_dim):
+        super().__init__()
+        self.attention_norm = nn.RMSNorm(width)
+        self.attention = Attention(width, head_dim)
+        self.ffn_norm = nn.RMSNorm(width)
+        self.mlp = SparseMLP(width, experts, expansion, selected)
+
+    def forward(self, x, causal=False):
+        x = x + self.attention(self.attention_norm(x), causal)
+        return x + self.mlp(self.ffn_norm(x))
+
+
+class Chunker(nn.Module):
+    """Learned boundaries, weighted pooling, and explicit source-to-chunk mapping.
+
+    Single-example sequences avoid treating padding as real observations. Processing
+    a complete observed prefix is bidirectional; next-token targets are not inputs.
+    """
+
+    def __init__(self, width, context_dim):
+        super().__init__()
+        self.boundary = nn.Linear(width, 1)
+        self.condition = nn.Linear(context_dim, width, bias=False)
+
+    def forward(self, x, context):
+        gate = torch.sigmoid(self.boundary(x + self.condition(context)[:, None]))
+        starts = gate[0, :, 0] > 0.5
+        starts = starts.clone()
+        starts[0] = True
+        assignment = starts.long().cumsum(0) - 1
+        groups = int(assignment[-1]) + 1
+        weights = gate[0]
+        sums = x.new_zeros(groups, x.shape[-1]).index_add(0, assignment, x[0] * weights)
+        counts = x.new_zeros(groups, 1).index_add(0, assignment, weights)
+        pooled = (sums / counts.clamp_min(torch.finfo(x.dtype).eps))[None]
+        return pooled, assignment
+
+
+class Stage(nn.Module):
+    def __init__(self, row, config, adaptive):
+        super().__init__()
+        width, depth, experts, expansion, selected = row
+        self.blocks = nn.ModuleList(
+            [Block(width, experts, expansion, selected, config.head_dim) for _ in range(depth)]
+        )
+        self.chunker = Chunker(width, config.context_dim) if adaptive else None
+
+    def forward(self, x, context, causal=False):
+        for block in self.blocks:
+            x = block(x, causal)
+        detail = x
+        if self.chunker is None:
+            return x, None, detail
+        pooled, assignment = self.chunker(x, context)
+        return pooled, assignment, detail
+
+
+class Hierarchy(nn.Module):
+    def __init__(self, name, config, adaptive=False):
+        super().__init__()
+        rows = config.stages[name]
+        self.stages = nn.ModuleList([Stage(row, config, adaptive) for row in rows])
+        self.transitions = nn.ModuleList(
+            [nn.Linear(a[0], b[0], bias=False) for a, b in itertools.pairwise(rows)]
+        )
+        self.dechunk_projections = (
+            nn.ModuleList([nn.Linear(b[0], a[0], bias=False) for a, b in itertools.pairwise(rows)])
+            if adaptive
+            else nn.ModuleList()
+        )
+
+    def forward(self, x, context, causal=False):
+        details, assignments = [], []
+        for i, stage in enumerate(self.stages):
+            x, assignment, detail = stage(x, context, causal)
+            details.append(detail)
+            assignments.append(assignment)
+            if i < len(self.transitions):
+                x = self.transitions[i](x)
+        return x, (details, assignments)
+
+    def dechunk(self, x, trace):
+        details, assignments = trace
+        for i in reversed(range(len(self.stages))):
+            if assignments[i] is not None:
+                x = x[:, assignments[i]] + details[i]
+            if i:
+                x = self.dechunk_projections[i - 1](x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, name, config):
+        super().__init__()
+        width = config.stages[name][0][0]
+        self.name = name
+        self.config = config
+        if name == "text_encoder":
+            self.embedding = nn.Embedding(config.vocab_size, width)
+        elif name == "image_encoder":
+            self.patch_embedding = nn.Conv2d(
+                3, width, config.patch_size, stride=config.patch_size, bias=False
+            )
+        elif name == "video_encoder":
+            self.tubelet_embedding = nn.Conv3d(
+                3,
+                width,
+                (2, config.patch_size, config.patch_size),
+                stride=(2, config.patch_size, config.patch_size),
+                bias=False,
+            )
+        elif name == "audio_encoder":
+            self.waveform_projection = nn.Linear(config.audio_frame_size, width, bias=False)
+            self.spectrum_projection = nn.Linear(
+                config.audio_frame_size // 2 + 1, width, bias=False
+            )
+        elif name == "eeg_encoder":
+            self.electrode_projection = nn.Linear(config.eeg_channels, width, bias=False)
+            self.graph_mix = nn.Linear(config.eeg_channels, config.eeg_channels, bias=False)
+        self.hierarchy = Hierarchy(name, config, adaptive=True)
+
+    def forward(self, value, context, adjacency=None):
+        rank = {
+            "text_encoder": 2,
+            "audio_encoder": 2,
+            "image_encoder": 4,
+            "video_encoder": 5,
+            "eeg_encoder": 3,
+        }[self.name]
+        if value.ndim != rank or value.shape[0] != 1 or 0 in value.shape:
+            raise ValueError(
+                f"{self.name} requires a nonempty rank-{rank} input with batch size one"
+            )
+        if self.name in ("image_encoder", "video_encoder") and (
+            value.shape[1] != 3 or any(n % self.config.patch_size for n in value.shape[-2:])
+        ):
+            raise ValueError("RGB spatial dimensions must be divisible by patch_size")
+        if self.name == "video_encoder" and value.shape[2] % 2:
+            raise ValueError("Video frame count must be divisible by two")
+        if self.name == "eeg_encoder" and value.shape[1] != self.config.eeg_channels:
+            raise ValueError("EEG channel count must match config")
+        if self.name == "text_encoder":
+            x = self.embedding(value)
+        elif self.name == "image_encoder":
+            x = self.patch_embedding(value).flatten(2).transpose(1, 2)
+        elif self.name == "video_encoder":
+            x = self.tubelet_embedding(value).flatten(2).transpose(1, 2)
+        elif self.name == "audio_encoder":
+            frame = self.config.audio_frame_size
+            if value.shape[-1] % frame:
+                raise ValueError("Audio must contain complete frames")
+            frames = value.unfold(-1, frame, frame)
+            spectrum = torch.fft.rfft(frames.float()).abs().to(value.dtype)
+            x = self.waveform_projection(frames) + self.spectrum_projection(spectrum)
+        else:
+            channels = self.config.eeg_channels
+            if adjacency is None or adjacency.shape != (channels, channels):
+                raise ValueError("EEG requires an explicit electrode adjacency matrix")
+            if (
+                (adjacency < 0).any()
+                or not torch.isfinite(adjacency).all()
+                or (adjacency.sum(-1) <= 0).any()
+            ):
+                raise ValueError("Invalid EEG adjacency")
+            graph = adjacency / adjacency.sum(-1, keepdim=True)
+            mixed = (graph @ value).transpose(1, 2)
+            x = self.electrode_projection(mixed + self.graph_mix(mixed))
+        position = torch.arange(x.shape[1], device=x.device)[:, None]
+        rate = torch.exp(-math.log(10000) * torch.arange(x.shape[2], device=x.device) / x.shape[2])
+        x = x + torch.sin(position * rate).to(x.dtype)[None]
+        high, trace = self.hierarchy(x, context)
+        # Fuse reconstructed detail features with the high-level representation.
+        details = self.hierarchy.dechunk(high, trace)
+        return high, details, trace
+
+
+class Fusion(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        width = config.stages["fusion"][0][0]
+        encoders = sorted(k for k in config.stages if k.endswith("_encoder"))
+        self.high_projections = nn.ModuleDict(
+            {k: nn.Linear(config.stages[k][-1][0], width, bias=False) for k in encoders}
+        )
+        self.detail_projections = nn.ModuleDict(
+            {k: nn.Linear(config.stages[k][0][0], width, bias=False) for k in encoders}
+        )
+        self.modality_embedding = nn.Embedding(5, width)
+        self.hierarchy = Hierarchy("fusion", config)
+
+    def forward(self, features, context):
+        tokens = []
+        for name, (high, detail, _) in features.items():
+            i = list(self.high_projections).index(name)
+            token = self.high_projections[name](high) + self.detail_projections[name](
+                detail.mean(1, keepdim=True)
+            )
+            tokens.append(token + self.modality_embedding.weight[i])
+        return self.hierarchy(torch.cat(tokens, dim=1), context)[0]
+
+
+class Decoder(nn.Module):
+    def __init__(self, name, config, output_width):
+        super().__init__()
+        self.input_projection = nn.Linear(
+            config.stages["world_model"][-1][0], config.stages[name][0][0], bias=False
+        )
+        self.hierarchy = Hierarchy(name, config)
+        self.output_projection = nn.Linear(config.stages[name][-1][0], output_width, bias=False)
+
+    def forward(self, x, context):
+        x = self.input_projection(x)
+        return self.output_projection(self.hierarchy(x, context, causal=True)[0])
+
+
+@dataclass
+class OmniOutput(ModelOutput):
+    next_token_logits: torch.Tensor | None = None
+    audio_codes: torch.Tensor | None = None
+    image_latents: torch.Tensor | None = None
+    video_latents: torch.Tensor | None = None
+    memory_state: torch.Tensor | None = None
+
+
+class OmniModel(PreTrainedModel):
+    config_class = OmniConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        for name in (k for k in config.stages if k.endswith("_encoder")):
+            setattr(self, name, Encoder(name, config))
+        self.fusion = Fusion(config)
+        self.world_input = nn.Linear(
+            config.stages["fusion"][-1][0], config.stages["world_model"][0][0], bias=False
+        )
+        self.world_model = Hierarchy("world_model", config, adaptive=True)
+        self.world_detail_projection = nn.Linear(
+            config.stages["world_model"][0][0], config.stages["world_model"][-1][0], bias=False
+        )
+        self.text_head = nn.Linear(
+            config.stages["world_model"][-1][0], config.vocab_size, bias=False
+        )
+        self.audio_decoder = Decoder(
+            "audio_decoder", config, config.audio_codebooks * config.audio_codebook_size
+        )
+        self.image_decoder = Decoder("image_decoder", config, 64)
+        self.video_decoder = Decoder("video_decoder", config, 128)
+        self.memory_input = nn.Linear(
+            config.stages["world_model"][-1][0], config.stages["memory_control"][0][0], bias=False
+        )
+        self.memory_control = Hierarchy("memory_control", config)
+        self.memory_output = nn.Linear(
+            config.stages["memory_control"][-1][0], config.stages["fusion"][0][0], bias=False
+        )
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d, nn.Conv3d)):
+            nn.init.normal_(module.weight, std=0.02)
+            if getattr(module, "bias", None) is not None:
+                nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        input_ids,
+        pixel_values=None,
+        audio_values=None,
+        video_values=None,
+        eeg_values=None,
+        eeg_adjacency=None,
+        task_vector=None,
+        memory_state=None,
+        output_modalities=(),
+    ):
+        if input_ids.ndim != 2 or input_ids.shape[0] != 1 or input_ids.shape[1] == 0:
+            raise ValueError("Reference execution requires one nonempty observed prefix per call")
+        context = (
+            self.text_head.weight.new_zeros(1, self.config.context_dim)
+            if task_vector is None
+            else task_vector
+        )
+        if context.shape != (1, self.config.context_dim):
+            raise ValueError("Invalid task_vector shape")
+        values = {
+            "text_encoder": input_ids,
+            "image_encoder": pixel_values,
+            "audio_encoder": audio_values,
+            "video_encoder": video_values,
+            "eeg_encoder": eeg_values,
+        }
+        features = {
+            name: getattr(self, name)(value, context, eeg_adjacency)
+            for name, value in values.items()
+            if value is not None
+        }
+        fused = self.fusion(features, context)
+        if memory_state is not None:
+            fused = fused + self.memory_output(memory_state)[:, None]
+        world, trace = self.world_model(self.world_input(fused), context)
+        detail = self.world_model.dechunk(world, trace)
+        world = world + self.world_detail_projection(detail.mean(1, keepdim=True))
+        summary = world.mean(1)
+        memory = self.memory_control(self.memory_input(summary[:, None]), context)[0][:, 0]
+        outputs = {}
+        if set(output_modalities) - {"audio", "image", "video"}:
+            raise ValueError("Unknown output modality")
+        for name in output_modalities:
+            output = getattr(self, name + "_decoder")(world, context)
+            key = "audio_codes" if name == "audio" else name + "_latents"
+            outputs[key] = output
+        return OmniOutput(next_token_logits=self.text_head(summary), memory_state=memory, **outputs)
+
+    @torch.no_grad()
+    def generate(self, input_ids, max_new_tokens=8, **kwargs):
+        for _ in range(max_new_tokens):
+            token = self(input_ids, **kwargs).next_token_logits.argmax(-1, keepdim=True)
+            input_ids = torch.cat((input_ids, token), dim=-1)
+        return input_ids
+
+
+OmniModel.register_for_auto_class("AutoModel")

requirements.txt

@@ -0,0 +1,3 @@
+torch==2.15.0
+transformers==4.57.6
+safetensors==0.8.0

shards/000/model-00001-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00002-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00003-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00004-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00005-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00006-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00007-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00008-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00009-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00010-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00011-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00012-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00013-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00014-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00015-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00016-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00017-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00018-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00019-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00020-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00021-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00022-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00023-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00024-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00025-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00026-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00027-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00028-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00029-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00030-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00031-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00032-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00033-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00034-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00035-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00036-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00037-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640

shards/000/model-00038-of-09400.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098438fd8e0a7d1416dcf33e4873f7ba42fa87184b40606141e6d35e8d426181
+size 268435640