cortexflowx 0.1.0__py3-none-any.whl

cortexflow/brain2img.py

@@ -0,0 +1,234 @@
+"""Brain → Image reconstruction pipeline.
+
+End-to-end pipeline: fMRI voxels → BrainEncoder → DiT (flow matching)
+→ VAE decoder → RGB image.
+
+This is the core brain decoding pipeline. Given measured fMRI activity
+while a subject views an image, reconstruct what they saw.
+
+Reference architecture inspired by:
+- MindEye (Scotti et al. 2023) for brain-to-image conditioning
+- SD3 / FLUX for the DiT + flow matching backbone
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from cortexflow._types import (
+    BrainData,
+    DiTConfig,
+    FlowConfig,
+    Modality,
+    ReconstructionResult,
+    VAEConfig,
+)
+from cortexflow.brain_encoder import BrainEncoder
+from cortexflow.dit import DiffusionTransformer
+from cortexflow.flow_matching import RectifiedFlowMatcher
+from cortexflow.vae import LatentVAE
+
+
+class Brain2Image(nn.Module):
+    """Reconstruct images from brain activity using DiT + Flow Matching.
+
+    Full pipeline::
+
+        fMRI voxels
+            → BrainEncoder (global + tokens)
+            → DiffusionTransformer (flow matching ODE)
+            → LatentVAE.decode()
+            → RGB image
+
+    Args:
+        n_voxels: Number of fMRI voxels.
+        img_size: Output image spatial resolution.
+        dit_config: DiT architecture configuration.
+        vae_config: VAE configuration.
+        flow_config: Flow matching configuration.
+        n_brain_tokens: Number of brain conditioning tokens.
+    """
+
+    def __init__(
+        self,
+        n_voxels: int = 1024,
+        img_size: int = 64,
+        dit_config: DiTConfig | None = None,
+        vae_config: VAEConfig | None = None,
+        flow_config: FlowConfig | None = None,
+        n_brain_tokens: int = 16,
+    ) -> None:
+        super().__init__()
+        self.img_size = img_size
+
+        # Configs
+        dit_cfg = dit_config or DiTConfig()
+        vae_cfg = vae_config or VAEConfig()
+        flow_cfg = flow_config or FlowConfig()
+
+        # VAE: determines latent spatial size
+        self.vae = LatentVAE(vae_cfg)
+        n_downsample = len(vae_cfg.hidden_dims)
+        latent_size = img_size // (2 ** n_downsample)
+
+        # Brain encoder
+        self.brain_encoder = BrainEncoder(
+            n_voxels=n_voxels,
+            cond_dim=dit_cfg.cond_dim,
+            n_tokens=n_brain_tokens,
+        )
+
+        # Unconditional embeddings for classifier-free guidance
+        self.uncond_global = nn.Parameter(torch.zeros(1, dit_cfg.cond_dim))
+        self.uncond_tokens = nn.Parameter(torch.zeros(1, n_brain_tokens, dit_cfg.cond_dim))
+
+        # DiT operates on VAE latent space
+        dit_cfg_updated = DiTConfig(
+            in_channels=vae_cfg.latent_channels,
+            hidden_dim=dit_cfg.hidden_dim,
+            depth=dit_cfg.depth,
+            num_heads=dit_cfg.num_heads,
+            patch_size=dit_cfg.patch_size,
+            cond_dim=dit_cfg.cond_dim,
+            mlp_ratio=dit_cfg.mlp_ratio,
+            qk_norm=dit_cfg.qk_norm,
+            use_cross_attn=dit_cfg.use_cross_attn,
+        )
+        self.dit = DiffusionTransformer(dit_cfg_updated, img_size=latent_size)
+
+        # Flow matcher
+        self.flow_matcher = RectifiedFlowMatcher(flow_cfg)
+
+        self._latent_size = latent_size
+        self._latent_channels = vae_cfg.latent_channels
+
+    def encode_brain(
+        self, brain_data: BrainData
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Encode fMRI to conditioning signals."""
+        return self.brain_encoder(brain_data.voxels)
+
+    def training_loss(
+        self,
+        images: torch.Tensor,
+        brain_data: BrainData,
+        cfg_dropout: float = 0.1,
+    ) -> torch.Tensor:
+        """Compute training loss for brain-conditioned image generation.
+
+        Args:
+            images: Target images ``(B, 3, H, W)``.
+            brain_data: Corresponding fMRI data.
+            cfg_dropout: Probability of dropping brain condition for CFG training.
+
+        Returns:
+            Scalar loss.
+        """
+        B = images.shape[0]
+
+        # Encode images to latent space (detach VAE — train separately or frozen)
+        with torch.no_grad():
+            z, _, _ = self.vae.encode(images)
+
+        # Encode brain
+        brain_global, brain_tokens = self.encode_brain(brain_data)
+
+        # CFG training: randomly drop conditioning
+        if cfg_dropout > 0 and self.training:
+            mask = torch.rand(B, device=images.device) < cfg_dropout
+            if mask.any():
+                brain_global = brain_global.clone()
+                brain_tokens = brain_tokens.clone()
+                brain_global[mask] = self.uncond_global.expand(mask.sum(), -1)
+                brain_tokens[mask] = self.uncond_tokens.expand(mask.sum(), -1, -1)
+
+        # Flow matching loss on latent space
+        return self.flow_matcher.compute_loss(
+            self.dit, z, brain_global, brain_tokens
+        )
+
+    @torch.no_grad()
+    def reconstruct(
+        self,
+        brain_data: BrainData,
+        num_steps: int = 50,
+        cfg_scale: float = 4.0,
+    ) -> ReconstructionResult:
+        """Reconstruct an image from brain activity.
+
+        Args:
+            brain_data: fMRI data to decode.
+            num_steps: Number of ODE solver steps.
+            cfg_scale: Classifier-free guidance scale.
+
+        Returns:
+            ReconstructionResult with the decoded image.
+        """
+        B = brain_data.batch_size
+        device = brain_data.voxels.device
+
+        # Encode brain
+        brain_global, brain_tokens = self.encode_brain(brain_data)
+
+        # Unconditional embeddings for CFG
+        uncond_global = self.uncond_global.expand(B, -1)
+        uncond_tokens = self.uncond_tokens.expand(B, -1, -1)
+
+        # Sample latents via flow matching
+        latent_shape = (B, self._latent_channels, self._latent_size, self._latent_size)
+        z = self.flow_matcher.sample(
+            self.dit,
+            shape=latent_shape,
+            brain_global=brain_global,
+            brain_tokens=brain_tokens,
+            num_steps=num_steps,
+            cfg_scale=cfg_scale,
+            brain_global_uncond=uncond_global,
+            brain_tokens_uncond=uncond_tokens,
+        )
+
+        # Decode latents to images
+        images = self.vae.decode(z)
+        images = images.clamp(0, 1)
+
+        return ReconstructionResult(
+            modality=Modality.IMAGE,
+            output=images,
+            brain_condition=brain_global,
+            n_steps=num_steps,
+            cfg_scale=cfg_scale,
+        )
+
+
+def build_brain2img(
+    n_voxels: int = 1024,
+    img_size: int = 64,
+    hidden_dim: int = 256,
+    depth: int = 6,
+    num_heads: int = 8,
+) -> Brain2Image:
+    """Build a Brain2Image model with sensible defaults.
+
+    Args:
+        n_voxels: Number of fMRI voxels.
+        img_size: Target image size (square).
+        hidden_dim: DiT hidden dimension.
+        depth: Number of DiT blocks.
+        num_heads: Attention heads.
+    """
+    dit_cfg = DiTConfig(
+        hidden_dim=hidden_dim,
+        depth=depth,
+        num_heads=num_heads,
+        cond_dim=hidden_dim,
+    )
+    vae_cfg = VAEConfig(hidden_dims=[32, 64])  # lightweight for small images
+    flow_cfg = FlowConfig()
+    return Brain2Image(
+        n_voxels=n_voxels,
+        img_size=img_size,
+        dit_config=dit_cfg,
+        vae_config=vae_cfg,
+        flow_config=flow_cfg,
+    )

cortexflow/brain2text.py

@@ -0,0 +1,278 @@
+"""Brain → Text reconstruction pipeline.
+
+Reconstruct what someone read or thought in linguistic form from fMRI.
+
+Architecture: fMRI → BrainEncoder → Transformer Decoder with
+cross-attention to brain tokens → autoregressive text generation.
+
+Unlike image/audio which use flow matching on continuous latents,
+text uses autoregressive decoding since language is inherently discrete.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cortexflow._types import BrainData, Modality, ReconstructionResult
+from cortexflow.brain_encoder import BrainEncoder
+
+
+class TextDecoderBlock(nn.Module):
+    """Transformer decoder block with causal self-attention + brain cross-attention."""
+
+    def __init__(self, hidden_dim: int, num_heads: int, cond_dim: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_dim, eps=1e-6)
+        self.self_attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(hidden_dim, 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.norm3 = nn.LayerNorm(hidden_dim, 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),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        brain_tokens: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Causal self-attention
+        h = self.norm1(x)
+        h, _ = self.self_attn(h, h, h, attn_mask=attn_mask, need_weights=False)
+        x = x + h
+
+        # Cross-attention to brain tokens
+        h = self.norm2(x)
+        kv = self.cond_proj(brain_tokens)
+        h, _ = self.cross_attn(h, kv, kv, need_weights=False)
+        x = x + h
+
+        # Feedforward
+        h = self.norm3(x)
+        x = x + self.mlp(h)
+        return x
+
+
+class BrainTextDecoder(nn.Module):
+    """Transformer decoder for brain → text.
+
+    Uses a simple character/word-piece vocabulary for self-contained
+    operation without external tokenizers.
+    """
+
+    def __init__(
+        self,
+        vocab_size: int = 256,
+        max_len: int = 128,
+        hidden_dim: int = 256,
+        depth: int = 6,
+        num_heads: int = 8,
+        cond_dim: int = 256,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.max_len = max_len
+        self.hidden_dim = hidden_dim
+
+        self.token_embed = nn.Embedding(vocab_size, hidden_dim)
+        self.pos_embed = nn.Embedding(max_len, hidden_dim)
+
+        self.blocks = nn.ModuleList([
+            TextDecoderBlock(hidden_dim, num_heads, cond_dim) for _ in range(depth)
+        ])
+        self.final_norm = nn.LayerNorm(hidden_dim, eps=1e-6)
+        self.lm_head = nn.Linear(hidden_dim, vocab_size, bias=False)
+
+        # Tie weights
+        self.lm_head.weight = self.token_embed.weight
+
+    def _causal_mask(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        """Generate causal attention mask."""
+        mask = torch.triu(torch.full((seq_len, seq_len), float("-inf"), device=device), diagonal=1)
+        return mask
+
+    def forward(
+        self,
+        token_ids: torch.Tensor,
+        brain_tokens: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward pass for training (teacher forcing).
+
+        Args:
+            token_ids: ``(B, T)`` input token IDs.
+            brain_tokens: ``(B, K, cond_dim)`` brain conditioning tokens.
+
+        Returns:
+            ``(B, T, vocab_size)`` logits.
+        """
+        B, T = token_ids.shape
+        positions = torch.arange(T, device=token_ids.device).unsqueeze(0)
+        x = self.token_embed(token_ids) + self.pos_embed(positions)
+
+        mask = self._causal_mask(T, token_ids.device)
+        for block in self.blocks:
+            x = block(x, brain_tokens, attn_mask=mask)
+
+        x = self.final_norm(x)
+        return self.lm_head(x)
+
+
+class Brain2Text(nn.Module):
+    """Reconstruct text from brain activity.
+
+    Pipeline::
+
+        fMRI → BrainEncoder → BrainTextDecoder (autoregressive) → tokens → text
+
+    Uses byte-level encoding (vocab_size=256) for simplicity. Each byte
+    is one token, so this handles any text without a tokenizer.
+
+    Args:
+        n_voxels: Number of fMRI voxels.
+        max_len: Maximum output text length.
+        hidden_dim: Model dimension.
+        depth: Number of transformer layers.
+        num_heads: Attention heads.
+    """
+
+    def __init__(
+        self,
+        n_voxels: int = 1024,
+        max_len: int = 128,
+        hidden_dim: int = 256,
+        depth: int = 6,
+        num_heads: int = 8,
+    ):
+        super().__init__()
+        cond_dim = hidden_dim
+        self.brain_encoder = BrainEncoder(
+            n_voxels=n_voxels, cond_dim=cond_dim, n_tokens=16,
+        )
+        self.decoder = BrainTextDecoder(
+            vocab_size=256,  # byte-level
+            max_len=max_len,
+            hidden_dim=hidden_dim,
+            depth=depth,
+            num_heads=num_heads,
+            cond_dim=cond_dim,
+        )
+        self.max_len = max_len
+        self.bos_token = 0  # null byte as BOS
+
+    @staticmethod
+    def text_to_tokens(text: str) -> torch.Tensor:
+        """Encode text to byte-level token IDs."""
+        return torch.tensor(list(text.encode("utf-8")), dtype=torch.long)
+
+    @staticmethod
+    def tokens_to_text(tokens: torch.Tensor) -> str:
+        """Decode byte-level token IDs to text."""
+        byte_list = tokens.cpu().tolist()
+        # Stop at first null byte
+        if 0 in byte_list:
+            byte_list = byte_list[: byte_list.index(0)]
+        return bytes(byte_list).decode("utf-8", errors="replace")
+
+    def training_loss(
+        self,
+        text_tokens: torch.Tensor,
+        brain_data: BrainData,
+    ) -> torch.Tensor:
+        """Compute cross-entropy loss for text generation.
+
+        Args:
+            text_tokens: ``(B, T)`` target token IDs (byte-level).
+            brain_data: Corresponding fMRI data.
+
+        Returns:
+            Scalar cross-entropy loss.
+        """
+        _, brain_tokens = self.brain_encoder(brain_data.voxels)
+
+        # Input: [BOS, t1, t2, ..., t_{n-1}], Target: [t1, t2, ..., t_n]
+        input_tokens = text_tokens[:, :-1]
+        target_tokens = text_tokens[:, 1:]
+
+        logits = self.decoder(input_tokens, brain_tokens)
+        return F.cross_entropy(
+            logits.reshape(-1, logits.shape[-1]),
+            target_tokens.reshape(-1),
+            ignore_index=0,
+        )
+
+    @torch.no_grad()
+    def reconstruct(
+        self,
+        brain_data: BrainData,
+        max_len: int | None = None,
+        temperature: float = 0.8,
+        top_k: int = 50,
+    ) -> ReconstructionResult:
+        """Reconstruct text from brain activity via autoregressive decoding.
+
+        Args:
+            brain_data: fMRI data to decode.
+            max_len: Maximum generation length.
+            temperature: Sampling temperature.
+            top_k: Top-k filtering.
+
+        Returns:
+            ReconstructionResult with generated text as metadata.
+        """
+        B = brain_data.batch_size
+        device = brain_data.voxels.device
+        gen_len = max_len or self.max_len
+
+        _, brain_tokens = self.brain_encoder(brain_data.voxels)
+
+        # Start with BOS token
+        generated = torch.full((B, 1), self.bos_token, dtype=torch.long, device=device)
+
+        for _ in range(gen_len - 1):
+            logits = self.decoder(generated, brain_tokens)
+            next_logits = logits[:, -1, :] / temperature
+
+            # Top-k filtering
+            if top_k > 0:
+                topk_vals, _ = next_logits.topk(top_k, dim=-1)
+                threshold = topk_vals[:, -1].unsqueeze(-1)
+                next_logits = next_logits.masked_fill(next_logits < threshold, float("-inf"))
+
+            probs = F.softmax(next_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            generated = torch.cat([generated, next_token], dim=1)
+
+            # Stop if all sequences produced a null byte
+            if (next_token == 0).all():
+                break
+
+        # Decode to text
+        texts = [self.tokens_to_text(generated[i]) for i in range(B)]
+
+        return ReconstructionResult(
+            modality=Modality.TEXT,
+            output=generated,
+            brain_condition=brain_tokens.mean(dim=1),
+            metadata={"texts": texts},
+        )
+
+
+def build_brain2text(
+    n_voxels: int = 1024,
+    max_len: int = 128,
+    hidden_dim: int = 256,
+    depth: int = 6,
+) -> Brain2Text:
+    """Build a Brain2Text model with sensible defaults."""
+    return Brain2Text(
+        n_voxels=n_voxels, max_len=max_len,
+        hidden_dim=hidden_dim, depth=depth,
+    )