cortexflowx 0.1.0__py3-none-any.whl

cortexflow/__init__.py

@@ -0,0 +1,78 @@
+"""CortexFlow — Brain decoding with Diffusion Transformers & Flow Matching.
+
+Reconstruct what someone saw, heard, or thought from fMRI brain activity.
+
+Architecture: fMRI → BrainEncoder → DiT (flow matching) → stimulus
+
+Modules:
+    dit             Diffusion Transformer backbone (AdaLN-Zero, QK-Norm, SwiGLU)
+    flow_matching   Rectified flow training & ODE sampling
+    vae             Latent space compression (2D images, 1D audio)
+    brain_encoder   fMRI → conditioning embeddings (global + tokens)
+    brain2img       Full brain → image pipeline
+    brain2audio     Full brain → audio pipeline
+    brain2text      Full brain → text pipeline (autoregressive)
+    training        Training loops, schedulers, synthetic data
+"""
+
+from cortexflow._types import (
+    BrainData,
+    CortexFlowError,
+    DiTConfig,
+    FlowConfig,
+    Modality,
+    ReconstructionResult,
+    TrainingConfig,
+    VAEConfig,
+)
+from cortexflow.brain2audio import AudioDiT, Brain2Audio, build_brain2audio
+from cortexflow.brain2img import Brain2Image, build_brain2img
+from cortexflow.brain2text import Brain2Text, BrainTextDecoder, build_brain2text
+from cortexflow.brain_encoder import (
+    BrainEncoder,
+    ROIBrainEncoder,
+    SubjectAdapter,
+    make_synthetic_brain_data,
+)
+from cortexflow.dit import DiffusionTransformer
+from cortexflow.flow_matching import EMAModel, RectifiedFlowMatcher
+from cortexflow.training import SyntheticBrainDataset, Trainer, WarmupCosineScheduler
+from cortexflow.vae import AudioVAE, LatentVAE
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Types
+    "BrainData",
+    "ReconstructionResult",
+    "DiTConfig",
+    "FlowConfig",
+    "VAEConfig",
+    "TrainingConfig",
+    "Modality",
+    "CortexFlowError",
+    # Core
+    "DiffusionTransformer",
+    "RectifiedFlowMatcher",
+    "EMAModel",
+    "LatentVAE",
+    "AudioVAE",
+    "BrainEncoder",
+    "ROIBrainEncoder",
+    "SubjectAdapter",
+    # Pipelines
+    "Brain2Image",
+    "Brain2Audio",
+    "Brain2Text",
+    "AudioDiT",
+    "BrainTextDecoder",
+    # Factories
+    "build_brain2img",
+    "build_brain2audio",
+    "build_brain2text",
+    "make_synthetic_brain_data",
+    # Training
+    "Trainer",
+    "WarmupCosineScheduler",
+    "SyntheticBrainDataset",
+]

cortexflow/_types.py

@@ -0,0 +1,115 @@
+"""Core data types for cortexflow."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+import torch
+
+
+class Modality(str, Enum):
+    """Output modality for brain decoding."""
+
+    IMAGE = "image"
+    AUDIO = "audio"
+    TEXT = "text"
+
+
+@dataclass
+class BrainData:
+    """Container for fMRI brain activity data.
+
+    Attributes:
+        voxels: Tensor of shape ``(batch, n_voxels)`` or ``(n_voxels,)``.
+        subject_id: Optional subject identifier for subject-specific adapters.
+        roi_mask: Optional boolean mask indicating which voxels belong to ROIs.
+        tr: Repetition time in seconds (fMRI temporal resolution).
+    """
+
+    voxels: torch.Tensor
+    subject_id: str | None = None
+    roi_mask: torch.Tensor | None = None
+    tr: float = 2.0
+
+    @property
+    def n_voxels(self) -> int:
+        return self.voxels.shape[-1]
+
+    @property
+    def batch_size(self) -> int:
+        if self.voxels.ndim == 1:
+            return 1
+        return self.voxels.shape[0]
+
+
+@dataclass
+class ReconstructionResult:
+    """Result of a brain decoding reconstruction."""
+
+    modality: Modality
+    output: torch.Tensor  # decoded output (image / waveform / token ids)
+    brain_condition: torch.Tensor  # the conditioning embeddings used
+    n_steps: int = 50  # diffusion steps used
+    cfg_scale: float = 1.0  # classifier-free guidance scale
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class TrainingConfig:
+    """Configuration for training a brain decoder."""
+
+    learning_rate: float = 1e-4
+    batch_size: int = 16
+    n_epochs: int = 100
+    warmup_steps: int = 500
+    weight_decay: float = 0.01
+    ema_decay: float = 0.9999
+    grad_clip: float = 1.0
+    mixed_precision: bool = False
+    log_every: int = 100
+    save_every: int = 1000
+    eval_every: int = 500
+
+
+@dataclass
+class DiTConfig:
+    """Configuration for the Diffusion Transformer."""
+
+    in_channels: int = 4  # latent channels
+    hidden_dim: int = 768
+    depth: int = 12
+    num_heads: int = 12
+    patch_size: int = 2
+    cond_dim: int = 768  # brain conditioning dimension
+    mlp_ratio: float = 4.0
+    qk_norm: bool = True
+    use_cross_attn: bool = True
+
+
+@dataclass
+class VAEConfig:
+    """Configuration for the latent VAE."""
+
+    in_channels: int = 3  # RGB
+    latent_channels: int = 4
+    hidden_dims: list[int] = field(default_factory=lambda: [64, 128, 256, 512])
+    kl_weight: float = 1e-6
+
+
+@dataclass
+class FlowConfig:
+    """Configuration for flow matching."""
+
+    num_steps: int = 50
+    sigma_min: float = 1e-5
+    logit_normal: bool = True  # SD3-style timestep sampling
+    logit_normal_mean: float = 0.0
+    logit_normal_std: float = 1.0
+    solver: str = "euler"  # "euler" or "midpoint"
+    cfg_scale: float = 4.0
+
+
+class CortexFlowError(Exception):
+    """Base exception for cortexflow."""

cortexflow/brain2audio.py

@@ -0,0 +1,298 @@
+"""Brain → Audio reconstruction pipeline.
+
+Reconstruct what someone heard from their fMRI activity.
+
+Architecture: fMRI → BrainEncoder → DiT (flow matching on mel spectrograms)
+→ Griffin-Lim or learned vocoder → waveform.
+
+The DiT operates on 1D latent sequences (compressed mel spectrograms)
+rather than 2D spatial maps.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cortexflow._types import (
+    BrainData,
+    DiTConfig,
+    FlowConfig,
+    Modality,
+    ReconstructionResult,
+)
+from cortexflow.brain_encoder import BrainEncoder
+from cortexflow.flow_matching import RectifiedFlowMatcher
+
+
+# ── 1D DiT for audio ────────────────────────────────────────────────────
+
+
+class DiTBlock1D(nn.Module):
+    """DiT block adapted for 1D sequences (audio spectrograms)."""
+
+    def __init__(self, hidden_dim: int, num_heads: int, cond_dim: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.adaLN = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, 6 * hidden_dim))
+        self.norm1 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True)
+        self.norm_cross = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.cross_attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True)
+        self.cond_proj = nn.Linear(cond_dim, hidden_dim) if cond_dim != hidden_dim else nn.Identity()
+        self.norm2 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        mlp_h = int(hidden_dim * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_dim, mlp_h), nn.SiLU(),
+            nn.Linear(mlp_h, hidden_dim),
+        )
+        nn.init.zeros_(self.adaLN[-1].weight)
+        nn.init.zeros_(self.adaLN[-1].bias)
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor, brain_tokens: torch.Tensor | None = None):
+        s1, sc1, g1, s2, sc2, g2 = self.adaLN(c).chunk(6, dim=-1)
+
+        h = self.norm1(x) * (1 + sc1.unsqueeze(1)) + s1.unsqueeze(1)
+        h, _ = self.attn(h, h, h, need_weights=False)
+        x = x + g1.unsqueeze(1) * h
+
+        if brain_tokens is not None:
+            h = self.norm_cross(x)
+            kv = self.cond_proj(brain_tokens)
+            h, _ = self.cross_attn(h, kv, kv, need_weights=False)
+            x = x + h
+
+        h = self.norm2(x) * (1 + sc2.unsqueeze(1)) + s2.unsqueeze(1)
+        h = self.mlp(h)
+        x = x + g2.unsqueeze(1) * h
+        return x
+
+
+class AudioDiT(nn.Module):
+    """1D Diffusion Transformer for mel spectrogram generation.
+
+    Operates on a sequence of mel frames (or compressed latents).
+    """
+
+    def __init__(
+        self,
+        n_mels: int = 80,
+        seq_len: int = 128,
+        hidden_dim: int = 256,
+        depth: int = 6,
+        num_heads: int = 8,
+        cond_dim: int = 256,
+    ):
+        super().__init__()
+        self.n_mels = n_mels
+        self.seq_len = seq_len
+        self.input_proj = nn.Linear(n_mels, hidden_dim)
+        self.pos_embed = nn.Parameter(torch.zeros(1, seq_len, hidden_dim))
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+
+        # Timestep embedding
+        self.time_embed = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim * 4), nn.SiLU(),
+            nn.Linear(hidden_dim * 4, hidden_dim),
+        )
+        self.time_dim = hidden_dim
+        self.cond_proj = nn.Sequential(nn.Linear(cond_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim))
+
+        self.blocks = nn.ModuleList([
+            DiTBlock1D(hidden_dim, num_heads, cond_dim) for _ in range(depth)
+        ])
+        self.final_norm = nn.LayerNorm(hidden_dim, eps=1e-6)
+        self.output_proj = nn.Linear(hidden_dim, n_mels)
+        nn.init.zeros_(self.output_proj.weight)
+        nn.init.zeros_(self.output_proj.bias)
+
+    def _sinusoidal_embed(self, t: torch.Tensor) -> torch.Tensor:
+        half = self.time_dim // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device).float() / half)
+        args = t.view(-1, 1) * freqs.unsqueeze(0)
+        return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        brain_global: torch.Tensor,
+        brain_tokens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Predict velocity for mel spectrogram flow matching.
+
+        Args:
+            x: ``(B, n_mels, T)`` noisy mel spectrogram.
+            t: ``(B,)`` timestep.
+            brain_global: ``(B, cond_dim)``
+            brain_tokens: ``(B, K, cond_dim)``
+
+        Returns:
+            ``(B, n_mels, T)`` predicted velocity.
+        """
+        B = x.shape[0]
+        # (B, n_mels, T) → (B, T, n_mels) → (B, T, hidden)
+        x = x.transpose(1, 2)
+        x = self.input_proj(x) + self.pos_embed[:, :x.shape[1]]
+
+        t_emb = self.time_embed(self._sinusoidal_embed(t))
+        c = t_emb + self.cond_proj(brain_global)
+
+        for block in self.blocks:
+            x = block(x, c, brain_tokens)
+
+        x = self.final_norm(x)
+        x = self.output_proj(x)  # (B, T, n_mels)
+        return x.transpose(1, 2)  # (B, n_mels, T)
+
+
+# ── Brain2Audio Pipeline ────────────────────────────────────────────────
+
+
+class Brain2Audio(nn.Module):
+    """Reconstruct audio from brain activity using 1D DiT + flow matching.
+
+    Pipeline::
+
+        fMRI → BrainEncoder → AudioDiT (flow matching) → mel spectrogram
+        → Griffin-Lim → waveform
+
+    Args:
+        n_voxels: Number of fMRI voxels.
+        n_mels: Number of mel frequency bins.
+        audio_len: Number of mel spectrogram frames.
+        sample_rate: Audio sample rate (for mel spectrogram computation).
+    """
+
+    def __init__(
+        self,
+        n_voxels: int = 1024,
+        n_mels: int = 80,
+        audio_len: int = 128,
+        hidden_dim: int = 256,
+        depth: int = 6,
+        num_heads: int = 8,
+        sample_rate: int = 16000,
+    ):
+        super().__init__()
+        self.n_mels = n_mels
+        self.audio_len = audio_len
+        self.sample_rate = sample_rate
+
+        cond_dim = hidden_dim
+        self.brain_encoder = BrainEncoder(
+            n_voxels=n_voxels, cond_dim=cond_dim, n_tokens=16,
+        )
+        self.dit = AudioDiT(
+            n_mels=n_mels, seq_len=audio_len, hidden_dim=hidden_dim,
+            depth=depth, num_heads=num_heads, cond_dim=cond_dim,
+        )
+        self.flow_matcher = RectifiedFlowMatcher(FlowConfig())
+
+        # Unconditional embeddings for CFG
+        self.uncond_global = nn.Parameter(torch.zeros(1, cond_dim))
+        self.uncond_tokens = nn.Parameter(torch.zeros(1, 16, cond_dim))
+
+    def training_loss(self, mel: torch.Tensor, brain_data: BrainData) -> torch.Tensor:
+        """Flow matching loss on mel spectrograms.
+
+        Args:
+            mel: Target mel spectrogram ``(B, n_mels, T)``.
+            brain_data: Corresponding fMRI.
+        """
+        brain_global, brain_tokens = self.brain_encoder(brain_data.voxels)
+        return self.flow_matcher.compute_loss(self.dit, mel, brain_global, brain_tokens)
+
+    @torch.no_grad()
+    def reconstruct(
+        self,
+        brain_data: BrainData,
+        num_steps: int = 50,
+        cfg_scale: float = 3.0,
+    ) -> ReconstructionResult:
+        """Reconstruct audio mel spectrogram from brain activity."""
+        B = brain_data.batch_size
+        device = brain_data.voxels.device
+        brain_global, brain_tokens = self.brain_encoder(brain_data.voxels)
+
+        mel_shape = (B, self.n_mels, self.audio_len)
+        mel = self.flow_matcher.sample(
+            self.dit, mel_shape, brain_global, brain_tokens,
+            num_steps=num_steps, cfg_scale=cfg_scale,
+            brain_global_uncond=self.uncond_global.expand(B, -1),
+            brain_tokens_uncond=self.uncond_tokens.expand(B, -1, -1),
+        )
+        return ReconstructionResult(
+            modality=Modality.AUDIO,
+            output=mel,
+            brain_condition=brain_global,
+            n_steps=num_steps,
+            cfg_scale=cfg_scale,
+        )
+
+    @staticmethod
+    def mel_to_waveform(mel: torch.Tensor, n_fft: int = 1024, hop_length: int = 256) -> torch.Tensor:
+        """Convert mel spectrogram to waveform via Griffin-Lim.
+
+        Args:
+            mel: ``(B, n_mels, T)`` mel spectrogram (linear scale).
+            n_fft: FFT size.
+            hop_length: Hop length.
+
+        Returns:
+            ``(B, samples)`` waveform.
+        """
+        B, n_mels, T = mel.shape
+        # Create simple mel filterbank inverse (pseudo-inverse approach)
+        n_freqs = n_fft // 2 + 1
+        # Approximate: project mel back to linear spectrogram
+        mel_basis = torch.linspace(0, 1, n_mels, device=mel.device).unsqueeze(1)
+        freq_basis = torch.linspace(0, 1, n_freqs, device=mel.device).unsqueeze(0)
+        filterbank = torch.exp(-0.5 * ((mel_basis - freq_basis) / 0.05) ** 2)
+        filterbank = filterbank / (filterbank.sum(dim=0, keepdim=True) + 1e-8)
+
+        # Mel → linear spectrogram (pseudo-inverse)
+        fb_pinv = filterbank.T  # (n_freqs, n_mels)
+        magnitude = torch.matmul(fb_pinv.unsqueeze(0), mel.clamp(min=0))  # (B, n_freqs, T)
+
+        # Griffin-Lim: iterative phase estimation
+        window = torch.hann_window(n_fft, device=mel.device)
+        out_length = T * hop_length
+        phase = torch.randn(B, n_freqs, T, device=mel.device) * 2 * math.pi
+        for _ in range(32):
+            stft = magnitude * torch.exp(1j * phase)
+            waveform = torch.istft(
+                stft, n_fft=n_fft, hop_length=hop_length,
+                window=window, length=out_length,
+            )
+            new_stft = torch.stft(
+                waveform, n_fft=n_fft, hop_length=hop_length,
+                window=window, return_complex=True,
+            )
+            # Trim or pad to match magnitude shape
+            nt = min(new_stft.shape[-1], T)
+            phase = torch.zeros_like(magnitude)
+            phase[:, :, :nt] = new_stft[:, :, :nt].angle()
+
+        stft = magnitude * torch.exp(1j * phase)
+        waveform = torch.istft(
+            stft, n_fft=n_fft, hop_length=hop_length,
+            window=window, length=out_length,
+        )
+        return waveform
+
+
+def build_brain2audio(
+    n_voxels: int = 1024,
+    n_mels: int = 80,
+    audio_len: int = 128,
+    hidden_dim: int = 256,
+    depth: int = 6,
+) -> Brain2Audio:
+    """Build a Brain2Audio model with sensible defaults."""
+    return Brain2Audio(
+        n_voxels=n_voxels, n_mels=n_mels, audio_len=audio_len,
+        hidden_dim=hidden_dim, depth=depth,
+    )