oraculumai/ZUNA1.1-CoreML-Apple

ZUNA1.1 Core ML: Enumerated Apple Profiles

This repository contains Core ML conversions of ZUNA1.1 for Apple-native inference on iPhone, visionOS, and macOS, organized under enumerated profile folders.

ZUNA1.1 is a 380M-parameter masked diffusion autoencoder for scalp EEG reconstruction and superresolution. Given a subset of channels and their 3D electrode coordinates, the model can:

• Denoise observed EEG channels
• Reconstruct dropped or missing channels
• Predict signals at novel scalp positions from physical coordinates

The base model extends ZUNA1 with variable-length masking (windows up to 30 seconds), temporal-segment reconstruction, and a more flexible corruption/dropout training scheme. This Core ML release preserves the pretrained behavior for Apple deployment using the fixed 5-second (seq_len=1280) profile contract described below.

Model Overview

The base model follows the same high-level inference pattern as upstream ZUNA:

• Inputs are EEG windows of 5 seconds @ 256 Hz (seq_len=1280).
• Signals are tokenized with num_fine_time_pts=32, so each channel produces 1280 / 32 = 40 coarse time tokens.
• tok_idx encodes {x, y, z, tc} (electrode position + coarse time index).
• Inference performs:
  1. Encoder forward pass (once)
  2. Decoder denoising loop (N diffusion steps)
  3. Token-to-signal reconstruction

As reported in the original ZUNA paper, the base architecture is a ~380M-parameter position-aware diffusion autoencoder trained on a large harmonized public EEG corpus.

This release preserves the base model tensor contract and publishes profile-specific shapes for deterministic Apple deployment.

Preprocessing Contract

For best parity with upstream behavior, keep the same preprocessing assumptions used by ZUNA:

• EEG montage must include 3D channel positions
• Sampling rate: 256 Hz
• Epoch length: 5 seconds (1280 samples)
• Token chunk size: 32 (40 coarse tokens per channel)
• Normalization aligned with upstream inference (data_norm=10.0)

These assumptions are what the released pretrained weights were optimized for.

Getting Started

Profile artifacts are organized as:

• profiles/14ch/fp32/...
• profiles/16ch/fp32/...
• profiles/32ch/fp32/...
• profiles/64ch/fp32/...

Each profile contains:

• ZunaEncoder.mlpackage
• ZunaDecoderStep.mlpackage
• ZunaDecoderStepUpdate.mlpackage
• coreml_export_metadata.json

Model split

• ZunaEncoder: Encodes tokenized EEG context
• ZunaDecoderStep: One denoising step in the diffusion loop
• ZunaDecoderStepUpdate: Decoder step + Euler update (z_next = z - dt * v_c)

Use DecoderStepUpdate when you want a minimal host-side loop and fewer host tensor ops.

Available profiles

Profile	Channels	Precision	Token Count	Encoder/Decoder Tensor Shape	final-z rel_l2 vs PyTorch
14ch-fp32	14	fp32	560	[1, 560, 32]	0.000003
16ch-fp32	16	fp32	640	[1, 640, 32]	0.000003
32ch-fp32	32	fp32	1280	[1, 1280, 32]	0.000003
64ch-fp32	64	fp32	2560	[1, 2560, 32]	0.000002

See profiles/index.json for machine-readable profile discovery.

Validation

All published profiles are checked against the original PyTorch weights using a 20-step diffusion parity run.

Profile	MAE	RMSE	max_abs	rel_l2	Threshold	Gate
14ch-fp32	0.000002	0.000003	0.000022	0.000003	0.005000	PASS
16ch-fp32	0.000002	0.000002	0.000015	0.000003	0.005000	PASS
32ch-fp32	0.000001	0.000002	0.000068	0.000003	0.005000	PASS
64ch-fp32	0.000001	0.000001	0.000058	0.000002	0.005000	PASS

Parity Visualization

Waveform Overlay + Residual

Representative sample from 64ch-fp32 (channel 0), final-z step-loop output.

Runtime Notes

• 14ch/fp32 is included for native Emotiv EPOC X-style 14-channel workflows.
• Profiles are FP32-only: the base model is bf16-trained and its decoder residual stream exceeds the fp16 range (max 65504), so fp16 conversions fail parity and are not published.
• Throughput/latency depends strongly on diffusion steps.
• Lower diffusion steps are useful for rapid iteration; higher steps improve reconstruction quality.

Upstream Resources

• Original model card: https://huggingface.co/Zyphra/ZUNA1.1
• Original repository and tutorials: https://github.com/Zyphra/zuna
• Technical paper page: https://www.zyphra.com/zuna-technical-paper

Citation

Please cite and credit the original ZUNA1.1 model and Zyphra resources:

• Base model: https://huggingface.co/Zyphra/ZUNA1.1
• Repository: https://github.com/Zyphra/zuna
• Paper: ZUNA: Flexible EEG Superresolution with Position-Aware Diffusion Autoencoders
• Technical page: https://www.zyphra.com/zuna-technical-paper

Disclaimer

This conversion is for research and engineering use only. It is not validated for medical diagnosis, treatment, or clinical decision-making.

Use at your own risk and follow the base model's usage and licensing terms.