bruAristimunha's picture
Add BrainOmni weights converted from OpenTSLab/BrainOmni@9a4d3c70 (braindecode PR w41/brainomni-rehost @ d6d9804d)
945aa48 verified
|
Raw History Blame Contribute Delete
2.37 kB
metadata
license: mit
library_name: braindecode
tags:
  - braindecode
  - pytorch
  - safetensors
  - eeg
  - meg
  - brainomni

braintokenizer-pretrained

Weights of the BrainOmni EEG/MEG VQ-VAE tokenizer (encoder, residual VQ and decoder), for braindecode.models.BrainTokenizer, converted from the authors' release.

from braindecode.models import BrainTokenizer
model = BrainTokenizer.from_pretrained("braindecode/braintokenizer-pretrained", chs_info=raw.info["chs"])

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 at revision 9a4d3c70495370397ccfbfd6d2496f25647545a5, file braintokenizer/BrainTokenizer.pt (sha256 d41c44c14c3f3b11fd0fb660752e356dff4cb4bc5f32a05f470f503ffddc7b1a), 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

@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.