Crate zuna_rs 

zuna-rs — ZUNA EEG Foundation Model inference in Rust

Pure-Rust inference for the ZUNA EEG foundation model, built on Burn 0.20 and exg for FIF preprocessing.

Three entry points

Type	Loads	Use case
ZunaInference	encoder + decoder	full encode → diffuse → decode pipeline
ZunaEncoder	encoder only	produce latent embeddings, save memory
ZunaDecoder	decoder only	reconstruct from stored embeddings

Quick start — full pipeline

use zuna_rs::{ZunaInference, InferenceResult};

let (model, _ms) = ZunaInference::<B>::load(
    Path::new("config.json"),
    Path::new("model.safetensors"),
    device,
)?;
let result: InferenceResult = model.run_fif(Path::new("recording.fif"), 50, 1.0, 10.0)?;
result.save_safetensors("output.safetensors")?;

Quick start — encode only

use zuna_rs::{ZunaEncoder, EncodingResult};

let (enc, _ms) = ZunaEncoder::<B>::load(
    Path::new("config.json"),
    Path::new("model.safetensors"),
    device,
)?;
let result: EncodingResult = enc.encode_fif(Path::new("recording.fif"), 10.0)?;
result.save_safetensors("data/embeddings.safetensors")?;

Quick start — decode from stored embeddings

use zuna_rs::{ZunaDecoder, encoder::EncodingResult};

let embeddings = EncodingResult::load_safetensors("data/embeddings.safetensors")?;
let (dec, _ms) = ZunaDecoder::<B>::load(
    Path::new("config.json"),
    Path::new("model.safetensors"),
    device,
)?;
let result = dec.decode_embeddings(&embeddings, 50, 1.0, 10.0)?;
result.save_safetensors("output.safetensors")?;

Embedding regularisation

The encoder uses an MMD (Maximum Mean Discrepancy) bottleneck: during training an MMD loss constrains the embedding distribution toward N(0, I). At inference the bottleneck is a pure passthrough — no reparameterisation is applied. Embeddings from ZunaEncoder or ZunaInference::encode_fif are therefore already in the regularised latent space and can be used directly for downstream tasks.

Re-exports

pub use inference::ZunaInference;

pub use inference::EpochOutput;

pub use inference::InferenceResult;

pub use encoder::ZunaEncoder;

pub use encoder::EpochEmbedding;

pub use encoder::EncodingResult;

pub use decoder::ZunaDecoder;

pub use config::ModelConfig;

pub use config::DataConfig;

pub use config::InferConfig;

pub use data::InputBatch;

pub use data::FifInfo;

pub use channel_positions::channel_xyz;

pub use channel_positions::MontageLayout;

pub use channel_positions::montage_channels;

pub use channel_positions::nearest_channel;

pub use channel_positions::normalise;

pub use csv_loader::load_from_csv;

pub use csv_loader::load_from_raw_tensor;

pub use csv_loader::load_from_named_tensor;

pub use csv_loader::PaddingStrategy;

pub use csv_loader::CsvLoadOptions;

pub use csv_loader::CsvInfo;

pub use csv_export::fif_to_csv;

Modules

channel_positions

EEG channel position lookup from embedded standard montage files.

config

csv_export

Export FIF recordings to CSV files.

csv_loader

CSV and raw-tensor loading for ZUNA inference.

data

decoder

Standalone ZUNA decoder — reconstruct EEG signals from latent embeddings.

encoder

Standalone ZUNA encoder — produce latent EEG embeddings.

inference

ZunaInference<B> — single entry point for the ZUNA EEG model.

model

weights