zuna 1.1.1

Foundation model for EEG reconstruction and interpolation

ZUNA: EEG Foundation Model

ZUNA is a 380M-parameter masked diffusion autoencoder trained to reconstruct, denoise, and upsample scalp-EEG signals. Given a subset of EEG channels, ZUNA can:

• Denoise existing EEG channels
• Reconstruct missing EEG channels
• Predict novel channel signals given physical coordinates on the scalp

ZUNA was trained on approximately 2 million channel-hours of EEG data from a wide range of publicly available sources. At 380M parameters, it is lightweight enough to run on a consumer GPU and can be used on CPU for many workloads.

Performance

ZUNA significantly outperforms existing standard methods for channel denoising, reconstruction, and upsampling. We compared ZUNA to MNE's default spherical spline interpolation method. ZUNA outperforms MNE in reconstruction accuracy across a range of unseen datasets, even those with a different preprocessing pipeline. ZUNA's advantage is particularly striking for higher upsampling ratios, demonstrating that it is effectively using general priors learned through large-scale pretraining.

Installation

# (1). Download tutorial and sample data from GitHub	
git clone --depth 1 --filter=blob:none --sparse https://github.com/Zyphra/zuna.git && cd zuna && git sparse-checkout set tutorials

# (2). Pip Install zuna
pip install zuna

Or install in development mode:

# (1). Download Zuna codebase from GitHub
git clone https://github.com/Zyphra/zuna.git && cd zuna

# (2). Pip Install zuna in developer mode
pip install -e .

GPU support (PyTorch + CUDA)

zuna runs on the GPU via PyTorch, and PyPI cannot pick a PyTorch build that matches your GPU driver for you. If the automatically installed torch is built for a newer CUDA version than your NVIDIA driver supports, PyTorch silently falls back to CPU (very slow) with a warning like No CUDA runtime is found / CUDA initialization: The NVIDIA driver on your system is too old.

To avoid this, install a torch build that matches your driver before installing zuna. Check the CUDA version your driver supports (top-right of nvidia-smi), then install the matching wheel — for example, for CUDA 12.8:

# 1. Install a torch build matching your driver's CUDA version (see `nvidia-smi`).
#    Example for CUDA 12.8 — use cu121 / cu124 / cu126 / cu128 to match yours:
pip install torch --index-url https://download.pytorch.org/whl/cu128

# 2. Then install zuna (it will use the torch you already installed)
pip install zuna

If you already installed zuna and it's running on CPU, fix it by reinstalling the matching torch:

pip install --force-reinstall torch --index-url https://download.pytorch.org/whl/cu128

Verify GPU access:

python -c "import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available())"
# Expect the CUDA build (e.g. ...+cu128) and `True`.

Quick Start

See tutorials/run_zuna_pipeline_new.py for a complete working example.

Note that you can also find a version of this script here on Google Colaboratory for free GPU access.

Edit the paths and options, then run:

python tutorials/run_zuna_pipeline_new.py

Citation

For more information please see our technical whitepaper and blog. If you find ZUNA useful in your work, please cite accordingly.

Organizations or researchers interested in collaborating with Zyphra to improve future versions for specific needs or use cases should contact [e-mail].

Disclaimer

This software and related services ("Services") are provided for research use only and are not intended for use in the diagnosis, cure, mitigation, treatment, or prevention of any disease or health condition. The Services have not been validated for any medical or clinical use. The information provided through the Services is for general informational purposes only and is not a substitute for any professional medical or healthcare advice. We do not warrant that any information provided through the Services is accurate, complete, or useful to you. Any reliance you place on such information is strictly at your own risk.