ZUNA — EEG Foundation Model (NeuroTechX HF-compatible build)

ZUNA is the Zyphra EEG foundation model — an encoder–decoder architecture (encoder → latent embeddings; flow-matching diffusion decoder), 382 M parameters. This is a self-contained, HuggingFace trust_remote_code build prepared by NeuroTechX so the model loads anywhere with only torch + transformers + a few public pip deps — no private packages required.

Part of the NeuroTechX EEG-FM Collection.

Install

pip install torch transformers vector_quantize_pytorch einx einops

Usage

from transformers import AutoModel
model = AutoModel.from_pretrained("NeuroTechX/zuna", trust_remote_code=True).cuda().eval()

# encoder-only latents (the typical EEG-FM use):
latents = model.encode(encoder_input, seq_lens, ...)   # [B, L', encoder_output_dim]

Or, in the NeuroTechX stack: make_hf_encoder("NeuroTechX/zuna") (emeg-fm).

What this build changes

A faithful mirror of the weights in mhough/zuna-base, made standalone. The original modeling_zuna.py imported a private zuna package, so it could not load via trust_remote_code for anyone without that environment. Here the EncoderDecoder architecture is vendored from the Zyphra source + Meta's lingua (MIT), imports rewritten to be self-contained, and the train-only activation-probe stubbed to a no-op at inference. Weights are byte-identical to the original.

Attribution & license

• Architecture & weights: Zyphra (ZUNA).
• Vendored framework: Meta lingua (MIT).
• Licensing of the ZUNA weights should be confirmed with Zyphra; this build is offered in good faith for open EEG-FM research and reproducibility.