Add BrainOmni weights converted from OpenTSLab/BrainOmni@9a4d3c70 (braindecode PR w41/brainomni-rehost @ d6d9804d)

README.md

@@ -0,0 +1,64 @@
+---
+license: mit
+library_name: braindecode
+tags:
+- braindecode
+- pytorch
+- safetensors
+- eeg
+- meg
+- brainomni
+---
+
+# brainomni-tiny-pretrained
+
+Weights of BrainOmni tiny (lm_dim 256, 8 heads, 12 blocks) with its frozen tokenizer, for
+[`braindecode.models.BrainOmni`](https://braindecode.org/stable/generated/braindecode.models.BrainOmni.html),
+converted from the authors' release. The classification head is not pretrained (seeded random init); fine-tune or linear-probe before use.
+
+```python
+from braindecode.models import BrainOmni
+model = BrainOmni.from_pretrained("braindecode/brainomni-tiny-pretrained", chs_info=raw.info["chs"], n_outputs=2)
+```
+
+`chs_info` must carry sensor positions (EEG) and coil orientations (MEG); the
+`chs_info` in `config.json` (19 EEG channels, 10-20) is only a default. Input is
+expected at 256 Hz, preprocessed as in the authors' code.
+
+## Source and conversion
+
+- Source: [OpenTSLab/BrainOmni](https://huggingface.co/OpenTSLab/BrainOmni) at
+  revision `9a4d3c70495370397ccfbfd6d2496f25647545a5`, file `tiny/BrainOmni.pt` (sha256 `62c67ba6a84ea0625e67a3b5e7463fe3930bfee88a612a225e9062a052542ffc`), MIT licence.
+- `convert_brainomni_checkpoints.py` (in this repository) renames the keys to
+  braindecode's, drops the pretraining-only mask predictor, stores the RoPE
+  cache as `(cos, sin)` pairs with zero sine (the released cache holds cosines
+  only and the released code uses it as loaded) and writes `config.json`,
+  `model.safetensors` and `pytorch_model.bin` with `save_pretrained`.
+- The converted model's outputs equal braindecode's loading of the original
+  file (max-abs difference 0.0, float32 and bfloat16).
+- Requires a braindecode version newer than 1.8.1.
+
+## Citation
+
+```bibtex
+@inproceedings{xiao2025brainomni,
+  title     = {BrainOmni: A Brain Foundation Model for Unified EEG and MEG Signals},
+  author    = {Xiao, Q. and Cui, Z. and Zhang, C. and Chen, S. and Wu, W. and
+               Thwaites, A. and Woolgar, A. and Zhou, B. and Zhang, C.},
+  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
+  year      = {2025},
+  note      = {arXiv:2505.18185},
+}
+
+@article{aristimunha2025braindecode,
+  title   = {Braindecode: a deep learning library for raw electrophysiological data},
+  author  = {Aristimunha, Bruno and others},
+  journal = {Zenodo},
+  year    = {2025},
+  doi     = {10.5281/zenodo.17699192},
+}
+```
+
+## License
+
+MIT, as the original BrainOmni release.

config.json

@@ -0,0 +1,451 @@
+{
+  "n_outputs": 2,
+  "n_chans": null,
+  "chs_info": [
+    {
+      "ch_name": "Fp1",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.030902595782959074,
+        0.1145851808674079,
+        0.027866574329093563,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "Fp2",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        0.028409487887091203,
+        0.11534630809630636,
+        0.02772125772253031,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "F7",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.07187663250970616,
+        0.07310353217275958,
+        0.02579045587891774,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "F3",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.05180904586237076,
+        0.0866879023191724,
+        0.07871409302883829,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "Fz",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.0012292786120998045,
+        0.09327445212894306,
+        0.10263929347682083,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "F4",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        0.05027427028109521,
+        0.087438385860065,
+        0.07727065434234637,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "F8",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        0.07143526180097565,
+        0.07450512184827826,
+        0.025101027393181548,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "T7",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.08598208540859081,
+        0.014871643830998479,
+        0.03117337373904913,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "C3",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.06714873038365092,
+        0.023358234962598565,
+        0.10451068106025274,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "Cz",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.0013741408533525957,
+        0.027617093717945417,
+        0.1401994941744,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "C4",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        0.06532887326785189,
+        0.02357310071942504,
+        0.10369243420048097,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "T8",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        0.08326136240706139,
+        0.015258183861365685,
+        0.030972968524064372,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "P7",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.07445797197401204,
+        -0.0421231600703902,
+        0.041273631387385656,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "P3",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.05503823925315043,
+        -0.04421030104524479,
+        0.09990897692201317,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "Pz",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.0017094524972804195,
+        -0.04521299319629392,
+        0.12667292418210063,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "P4",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        0.053636008012841616,
+        -0.044334525348021875,
+        0.10051602773215543,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "P8",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        0.07103246343848865,
+        -0.04225998064272267,
+        0.04119886278195534,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "O1",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        -0.0315735653251896,
+        -0.08056835005318111,
+        0.054789645624052044,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    },
+    {
+      "ch_name": "O2",
+      "kind": 2,
+      "coil_type": 1,
+      "unit": 107,
+      "cal": 1.0,
+      "range": 1.0,
+      "loc": [
+        0.027683093109090593,
+        -0.08048884161275736,
+        0.05473408147779854,
+        0.0,
+        0.0,
+        0.0,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN,
+        NaN
+      ]
+    }
+  ],
+  "n_times": 512,
+  "input_window_seconds": null,
+  "sfreq": 256.0,
+  "emb_dim": 256,
+  "n_neuro": 16,
+  "window_length": 512,
+  "overlap_ratio": 0.25,
+  "n_filters": 32,
+  "ratios": [
+    8,
+    4,
+    2
+  ],
+  "kernel_size": 5,
+  "last_kernel_size": 5,
+  "tokenizer_num_heads": 4,
+  "codebook_dim": 256,
+  "codebook_size": 512,
+  "num_quantizers": 4,
+  "rotation_trick": true,
+  "tokenizer_drop_prob": 0.0,
+  "lm_dim": 256,
+  "num_heads": 8,
+  "depth": 12,
+  "drop_prob": 0.1,
+  "activation": "torch.nn.modules.activation.SELU",
+  "braindecode_version": "1.8.1"
+}

convert_brainomni_checkpoints.py

@@ -0,0 +1,109 @@
+"""Convert the released BrainOmni checkpoints to braindecode-native files.
+
+Source: https://huggingface.co/OpenTSLab/BrainOmni at revision
+9a4d3c70495370397ccfbfd6d2496f25647545a5 (MIT), files
+``braintokenizer/BrainTokenizer.pt``, ``tiny/BrainOmni.pt``, ``base/BrainOmni.pt``
+and their ``model_cfg.json``.
+
+Usage::
+
+    python convert_brainomni_checkpoints.py OUT_DIR [SOURCE_DIR]
+
+writes ``OUT_DIR/{braintokenizer,brainomni-tiny,brainomni-base}-pretrained``
+with ``config.json``, ``model.safetensors`` and ``pytorch_model.bin``
+(``save_pretrained``). Without ``SOURCE_DIR`` the files are downloaded.
+
+Key changes: ``quantizer.rvq.`` -> ``quantizer.``, ``decoder.`` ->
+``final_layer.``, the doubled ``conv``/``convtr`` wrappers and SEANet's
+``.model`` are flattened, ``weight_g``/``weight_v`` -> weight-norm
+parametrizations, the feed-forward linears are renumbered as
+``FeedForwardBlock``. The pretraining-only ``mask_token`` and ``predict_head``
+are dropped. The RoPE cache ``rotate`` is stored by the release as cosines only
+(the released code copies it into a complex cache, so the sine is zero); it is
+kept as ``(cos, 0)`` pairs. The BrainOmni classification head is not
+pretrained: it is a seeded random ``nn.Linear`` default init.
+The stored ``chs_info`` (19 EEG channels, 10-20) is only a default; pass
+your own ``chs_info`` to ``from_pretrained``.
+"""
+
+import json
+import sys
+from pathlib import Path
+
+import mne
+import torch
+
+from braindecode.models import BrainOmni, BrainTokenizer
+
+REPO, REVISION = "OpenTSLab/BrainOmni", "9a4d3c70495370397ccfbfd6d2496f25647545a5"
+RENAMES = {
+    "n_dim": "emb_dim",
+    "n_head": "tokenizer_num_heads",
+    "lm_head": "num_heads",
+    "lm_depth": "depth",
+    "lm_dropout": "drop_prob",
+}
+UNUSED = {"mask_ratio", "num_quantizers_used", "quantize_optimize_method"}
+CH_NAMES = (
+    "Fp1 Fp2 F7 F3 Fz F4 F8 T7 C3 Cz C4 T8 P7 P3 Pz P4 P8 O1 O2".split()
+)
+
+
+def fetch(source, path):
+    if source is not None:
+        return Path(source) / path
+    from huggingface_hub import hf_hub_download
+
+    return Path(hf_hub_download(REPO, path, revision=REVISION))
+
+
+def rename(key):
+    key = key.replace("quantizer.rvq.", "quantizer.")
+    if key.startswith("decoder."):
+        key = "final_layer." + key.removeprefix("decoder.")
+    key = key.replace("tokenizer.decoder.", "tokenizer.final_layer.")
+    key = key.replace(".convtr.convtr.", ".convtr.").replace(".conv.conv.", ".conv.")
+    key = key.replace("seanet_encoder.model.", "seanet_encoder.")
+    key = key.replace("seanet_decoder.model.", "seanet_decoder.")
+    key = key.replace(".weight_g", ".parametrizations.weight.original0")
+    key = key.replace(".weight_v", ".parametrizations.weight.original1")
+    for name in ("ff", "aggregate_mlp"):
+        key = key.replace(f"{name}.layer.0.", f"{name}.0.")
+        key = key.replace(f"{name}.layer.2.", f"{name}.3.")
+    return key
+
+
+def convert(source, folder, cls, out, **kwargs):
+    config = json.loads(fetch(source, f"{folder}/model_cfg.json").read_text())
+    config = {RENAMES.get(k, k): v for k, v in config.items() if k not in UNUSED}
+    dropout = config.pop("dropout")
+    config["drop_prob" if cls is BrainTokenizer else "tokenizer_drop_prob"] = dropout
+    info = mne.create_info(CH_NAMES, 256.0, "eeg")
+    info.set_montage("standard_1020")
+    torch.manual_seed(0)  # the BrainOmni head is a seeded random init
+    model = cls(chs_info=info["chs"], n_times=512, sfreq=256.0, **config, **kwargs)
+
+    official = torch.load(
+        fetch(source, f"{folder}/{cls.__name__}.pt"), weights_only=True
+    )
+    state = {}
+    for key, value in official.items():
+        if key == "mask_token" or key.startswith("predict_head."):
+            continue
+        if key.endswith("rope_embedding_layer.rotate"):
+            value = torch.stack((value, torch.zeros_like(value)), dim=-1)
+        state[rename(key)] = value
+    for key, value in model.state_dict().items():
+        if key.startswith("final_layer.") and cls is BrainOmni:
+            state[key] = value
+    model.load_state_dict(state, strict=True)
+    model.save_pretrained(out)
+    return model
+
+
+if __name__ == "__main__":
+    out = Path(sys.argv[1])
+    source = sys.argv[2] if len(sys.argv) > 2 else None
+    convert(source, "braintokenizer", BrainTokenizer, out / "braintokenizer-pretrained")
+    for size in ("tiny", "base"):
+        convert(source, size, BrainOmni, out / f"brainomni-{size}-pretrained", n_outputs=2)

model.safetensors

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

pytorch_model.bin

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