cortexflowx 0.1.0__py3-none-any.whl

cortexflow/flow_matching.py

@@ -0,0 +1,236 @@
+"""Rectified Flow Matching for training and sampling.
+
+Implements the rectified flow framework from Lipman et al. (2022) and
+the improvements from Esser et al. (2024, Stable Diffusion 3):
+
+- **Linear interpolation** paths: x_t = (1 - t) · x_0 + t · x_1
+- **Velocity prediction**: v = x_1 - x_0  (the model learns the vector
+  field that transports noise to data)
+- **Logit-normal timestep sampling**: biases training toward perceptually
+  relevant noise levels (SD3 improvement over uniform sampling)
+- **Euler / Midpoint ODE solvers** for inference
+
+The training loss is simply MSE between predicted and target velocity:
+    L = E_{t, x_0, x_1} [ || v_θ(x_t, t, c) - (x_1 - x_0) ||² ]
+
+Reference: "Flow Matching Guide and Code" (arXiv:2412.06264)
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cortexflow._types import FlowConfig
+
+
+class RectifiedFlowMatcher:
+    """Rectified flow matching training and sampling.
+
+    Usage::
+
+        fm = RectifiedFlowMatcher()
+
+        # Training step
+        loss = fm.compute_loss(model, x_clean, brain_global, brain_tokens)
+        loss.backward()
+
+        # Sampling
+        x_gen = fm.sample(model, shape=(B, C, H, W), brain_global=bg, brain_tokens=bt)
+    """
+
+    def __init__(self, config: FlowConfig | None = None) -> None:
+        self.config = config or FlowConfig()
+
+    # ── Timestep Sampling ────────────────────────────────────────────
+
+    def sample_timesteps(
+        self, batch_size: int, device: torch.device
+    ) -> torch.Tensor:
+        """Sample timesteps t ∈ (0, 1).
+
+        If ``logit_normal`` is enabled (SD3-style), samples from a
+        logit-normal distribution which biases toward intermediate noise
+        levels where the perceptual signal is strongest.
+        """
+        cfg = self.config
+
+        if cfg.logit_normal:
+            # Logit-normal: sample u ~ N(mean, std²), then t = sigmoid(u)
+            u = torch.randn(batch_size, device=device)
+            u = u * cfg.logit_normal_std + cfg.logit_normal_mean
+            t = torch.sigmoid(u)
+        else:
+            t = torch.rand(batch_size, device=device)
+
+        # Clamp away from exact 0 and 1 for numerical stability
+        return t.clamp(cfg.sigma_min, 1.0 - cfg.sigma_min)
+
+    # ── Training ─────────────────────────────────────────────────────
+
+    def compute_loss(
+        self,
+        model: nn.Module,
+        x_1: torch.Tensor,
+        brain_global: torch.Tensor,
+        brain_tokens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Compute rectified flow matching loss.
+
+        Args:
+            model: DiT model that predicts velocity v(x_t, t, condition).
+            x_1: Clean data (target) ``(B, C, H, W)``.
+            brain_global: Global brain embedding ``(B, D)``.
+            brain_tokens: Brain token sequence ``(B, T, D)`` for cross-attention.
+
+        Returns:
+            Scalar MSE loss: E[ ||v_pred - v_target||² ]
+        """
+        B = x_1.shape[0]
+        device = x_1.device
+
+        # Sample timesteps
+        t = self.sample_timesteps(B, device)
+
+        # Sample noise
+        x_0 = torch.randn_like(x_1)
+
+        # Linear interpolation: x_t = (1-t) * noise + t * data
+        t_expand = t.view(B, *([1] * (x_1.ndim - 1)))
+        x_t = (1.0 - t_expand) * x_0 + t_expand * x_1
+
+        # Target velocity: points from noise to data
+        v_target = x_1 - x_0
+
+        # Model prediction
+        v_pred = model(x_t, t, brain_global, brain_tokens)
+
+        return F.mse_loss(v_pred, v_target)
+
+    # ── Sampling ─────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def sample(
+        self,
+        model: nn.Module,
+        shape: tuple[int, ...],
+        brain_global: torch.Tensor,
+        brain_tokens: torch.Tensor | None = None,
+        num_steps: int | None = None,
+        cfg_scale: float | None = None,
+        brain_global_uncond: torch.Tensor | None = None,
+        brain_tokens_uncond: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Generate samples by solving the ODE from noise to data.
+
+        Args:
+            model: Trained DiT model.
+            shape: Output shape ``(B, C, H, W)``.
+            brain_global: Global brain condition.
+            brain_tokens: Brain token sequence.
+            num_steps: Override number of ODE steps.
+            cfg_scale: Classifier-free guidance scale. 1.0 = no guidance.
+            brain_global_uncond: Unconditional brain embedding for CFG.
+            brain_tokens_uncond: Unconditional brain tokens for CFG.
+
+        Returns:
+            Generated latent ``(B, C, H, W)``.
+        """
+        cfg = self.config
+        steps = num_steps or cfg.num_steps
+        guidance = cfg_scale if cfg_scale is not None else cfg.cfg_scale
+        do_cfg = guidance > 1.0 and brain_global_uncond is not None
+        device = brain_global.device
+
+        # Start from pure noise
+        x = torch.randn(shape, device=device)
+        dt = 1.0 / steps
+
+        for i in range(steps):
+            t_val = i / steps
+            t = torch.full((shape[0],), t_val, device=device)
+
+            if cfg.solver == "midpoint":
+                # Midpoint method: evaluate at t + dt/2
+                v = self._get_velocity(
+                    model, x, t, brain_global, brain_tokens,
+                    guidance, do_cfg, brain_global_uncond, brain_tokens_uncond,
+                )
+                x_mid = x + v * (dt / 2)
+                t_mid = torch.full((shape[0],), t_val + dt / 2, device=device)
+                v_mid = self._get_velocity(
+                    model, x_mid, t_mid, brain_global, brain_tokens,
+                    guidance, do_cfg, brain_global_uncond, brain_tokens_uncond,
+                )
+                x = x + v_mid * dt
+            else:
+                # Euler method
+                v = self._get_velocity(
+                    model, x, t, brain_global, brain_tokens,
+                    guidance, do_cfg, brain_global_uncond, brain_tokens_uncond,
+                )
+                x = x + v * dt
+
+        return x
+
+    def _get_velocity(
+        self,
+        model: nn.Module,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        brain_global: torch.Tensor,
+        brain_tokens: torch.Tensor | None,
+        guidance: float,
+        do_cfg: bool,
+        brain_global_uncond: torch.Tensor | None,
+        brain_tokens_uncond: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """Get velocity with optional classifier-free guidance."""
+        if do_cfg and brain_global_uncond is not None:
+            # Conditional prediction
+            v_cond = model(x, t, brain_global, brain_tokens)
+            # Unconditional prediction
+            v_uncond = model(x, t, brain_global_uncond, brain_tokens_uncond)
+            # CFG interpolation
+            return v_uncond + guidance * (v_cond - v_uncond)
+        else:
+            return model(x, t, brain_global, brain_tokens)
+
+
+class EMAModel:
+    """Exponential Moving Average of model parameters.
+
+    Maintains a shadow copy of model weights updated as:
+        θ_ema = decay · θ_ema + (1 - decay) · θ_model
+    """
+
+    def __init__(self, model: nn.Module, decay: float = 0.9999) -> None:
+        self.decay = decay
+        self.shadow: dict[str, torch.Tensor] = {}
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = param.data.clone()
+
+    @torch.no_grad()
+    def update(self, model: nn.Module) -> None:
+        for name, param in model.named_parameters():
+            if name in self.shadow:
+                self.shadow[name].lerp_(param.data, 1.0 - self.decay)
+
+    def apply_to(self, model: nn.Module) -> dict[str, torch.Tensor]:
+        """Replace model params with EMA params. Returns original params."""
+        originals: dict[str, torch.Tensor] = {}
+        for name, param in model.named_parameters():
+            if name in self.shadow:
+                originals[name] = param.data.clone()
+                param.data.copy_(self.shadow[name])
+        return originals
+
+    def restore(self, model: nn.Module, originals: dict[str, torch.Tensor]) -> None:
+        """Restore original params after EMA evaluation."""
+        for name, param in model.named_parameters():
+            if name in originals:
+                param.data.copy_(originals[name])

cortexflow/training.py

@@ -0,0 +1,283 @@
+"""Training utilities for cortexflow models.
+
+Provides a generic training loop, learning rate schedulers, and data
+utilities for training brain decoders on fMRI datasets.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Callable, Iterator
+
+import torch
+import torch.nn as nn
+
+from cortexflow._types import TrainingConfig
+from cortexflow.flow_matching import EMAModel
+
+
+class WarmupCosineScheduler:
+    """Learning rate schedule: linear warmup → cosine decay.
+
+    Args:
+        optimizer: PyTorch optimizer.
+        warmup_steps: Linear warmup duration.
+        total_steps: Total training steps.
+        min_lr_ratio: Minimum LR as a fraction of peak.
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        warmup_steps: int,
+        total_steps: int,
+        min_lr_ratio: float = 0.01,
+    ) -> None:
+        self.optimizer = optimizer
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.min_lr_ratio = min_lr_ratio
+        self.base_lrs = [pg["lr"] for pg in optimizer.param_groups]
+        self._step = 0
+
+    def step(self) -> None:
+        self._step += 1
+        lr_scale = self._get_scale(self._step)
+        for pg, base_lr in zip(self.optimizer.param_groups, self.base_lrs):
+            pg["lr"] = base_lr * lr_scale
+
+    def _get_scale(self, step: int) -> float:
+        if step <= self.warmup_steps:
+            return step / max(1, self.warmup_steps)
+        progress = (step - self.warmup_steps) / max(
+            1, self.total_steps - self.warmup_steps
+        )
+        progress = min(progress, 1.0)
+        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+        return self.min_lr_ratio + (1.0 - self.min_lr_ratio) * cosine
+
+    @property
+    def current_lr(self) -> float:
+        return self.optimizer.param_groups[0]["lr"]
+
+
+class SyntheticBrainDataset:
+    """Synthetic dataset for testing and development.
+
+    Yields (stimulus, fmri) pairs with configurable modality.
+    """
+
+    def __init__(
+        self,
+        n_samples: int = 1000,
+        n_voxels: int = 1024,
+        modality: str = "image",
+        img_size: int = 64,
+        n_mels: int = 80,
+        audio_len: int = 128,
+        max_text_len: int = 64,
+    ) -> None:
+        self.n_samples = n_samples
+        self.n_voxels = n_voxels
+        self.modality = modality
+        self.img_size = img_size
+        self.n_mels = n_mels
+        self.audio_len = audio_len
+        self.max_text_len = max_text_len
+
+    def __len__(self) -> int:
+        return self.n_samples
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        """Generate a single synthetic sample."""
+        g = torch.Generator().manual_seed(idx)
+        fmri = torch.randn(self.n_voxels, generator=g)
+
+        if self.modality == "image":
+            stimulus = torch.rand(3, self.img_size, self.img_size, generator=g)
+        elif self.modality == "audio":
+            stimulus = torch.randn(self.n_mels, self.audio_len, generator=g).abs()
+        elif self.modality == "text":
+            # Random byte-level tokens (printable ASCII range)
+            tokens = torch.randint(32, 127, (self.max_text_len,), generator=g)
+            stimulus = tokens
+        else:
+            raise ValueError(f"Unknown modality: {self.modality}")
+
+        return {"fmri": fmri, "stimulus": stimulus}
+
+    def to_loader(
+        self, batch_size: int = 8, shuffle: bool = True
+    ) -> Iterator[dict[str, torch.Tensor]]:
+        """Simple batch iterator (no DataLoader dependency)."""
+        indices = list(range(self.n_samples))
+        if shuffle:
+            g = torch.Generator()
+            perm = torch.randperm(self.n_samples, generator=g).tolist()
+            indices = perm
+
+        for start in range(0, self.n_samples, batch_size):
+            batch_indices = indices[start : start + batch_size]
+            items = [self[i] for i in batch_indices]
+            batch: dict[str, torch.Tensor] = {}
+            for key in items[0]:
+                batch[key] = torch.stack([item[key] for item in items])
+            yield batch
+
+
+class Trainer:
+    """Generic training loop for cortexflow models.
+
+    Handles optimizer setup, LR scheduling, EMA, gradient clipping,
+    and logging.
+
+    Args:
+        model: A cortexflow model with a ``training_loss`` method.
+        config: Training configuration.
+        device: Device to train on.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        config: TrainingConfig | None = None,
+        device: str | torch.device = "cpu",
+    ) -> None:
+        self.model = model.to(device)
+        self.config = config or TrainingConfig()
+        self.device = torch.device(device)
+
+        cfg = self.config
+        self.optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=cfg.learning_rate,
+            weight_decay=cfg.weight_decay,
+            betas=(0.9, 0.95),
+        )
+
+        self.ema: EMAModel | None = None
+        if cfg.ema_decay > 0:
+            self.ema = EMAModel(model, decay=cfg.ema_decay)
+
+        self.global_step = 0
+        self.scheduler: WarmupCosineScheduler | None = None
+        self._log_fn: Callable[[dict[str, Any]], None] | None = None
+
+    def set_logger(self, fn: Callable[[dict[str, Any]], None]) -> None:
+        """Set a logging callback: fn({"step": ..., "loss": ..., ...})."""
+        self._log_fn = fn
+
+    def _log(self, metrics: dict[str, Any]) -> None:
+        if self._log_fn is not None:
+            self._log_fn(metrics)
+
+    def train_step(
+        self,
+        batch: dict[str, torch.Tensor],
+        loss_fn: Callable[[nn.Module, dict[str, torch.Tensor]], torch.Tensor],
+    ) -> float:
+        """Execute a single training step.
+
+        Args:
+            batch: Dict of tensors (moved to device automatically).
+            loss_fn: Computes loss from (model, batch) → scalar.
+
+        Returns:
+            Loss value as float.
+        """
+        cfg = self.config
+        self.model.train()
+
+        # Move data to device
+        batch_dev = {
+            k: v.to(self.device) if isinstance(v, torch.Tensor) else v
+            for k, v in batch.items()
+        }
+
+        loss = loss_fn(self.model, batch_dev)
+        loss.backward()
+
+        # Gradient clipping
+        if cfg.grad_clip > 0:
+            nn.utils.clip_grad_norm_(self.model.parameters(), cfg.grad_clip)
+
+        self.optimizer.step()
+        self.optimizer.zero_grad()
+
+        # EMA update
+        if self.ema is not None:
+            self.ema.update(self.model)
+
+        # LR schedule
+        if self.scheduler is not None:
+            self.scheduler.step()
+
+        self.global_step += 1
+
+        # Logging
+        if self.global_step % cfg.log_every == 0:
+            lr = (
+                self.scheduler.current_lr
+                if self.scheduler
+                else cfg.learning_rate
+            )
+            self._log({
+                "step": self.global_step,
+                "loss": loss.item(),
+                "lr": lr,
+            })
+
+        return loss.item()
+
+    def fit(
+        self,
+        dataset: SyntheticBrainDataset,
+        loss_fn: Callable[[nn.Module, dict[str, torch.Tensor]], torch.Tensor],
+        n_epochs: int | None = None,
+    ) -> list[float]:
+        """Full training loop over a dataset.
+
+        Args:
+            dataset: Training data.
+            loss_fn: Loss function (model, batch) → scalar.
+            n_epochs: Override number of epochs.
+
+        Returns:
+            List of per-step losses.
+        """
+        cfg = self.config
+        epochs = n_epochs or cfg.n_epochs
+        steps_per_epoch = max(1, len(dataset) // cfg.batch_size)
+        total_steps = epochs * steps_per_epoch
+
+        self.scheduler = WarmupCosineScheduler(
+            self.optimizer, cfg.warmup_steps, total_steps
+        )
+
+        losses: list[float] = []
+        for epoch in range(epochs):
+            for batch in dataset.to_loader(batch_size=cfg.batch_size, shuffle=True):
+                loss = self.train_step(batch, loss_fn)
+                losses.append(loss)
+
+        return losses
+
+    def save_checkpoint(self, path: str) -> None:
+        """Save model + optimizer + EMA state."""
+        state = {
+            "model": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "global_step": self.global_step,
+        }
+        if self.ema is not None:
+            state["ema"] = self.ema.shadow
+        torch.save(state, path)
+
+    def load_checkpoint(self, path: str) -> None:
+        """Load model + optimizer + EMA state."""
+        state = torch.load(path, map_location=self.device, weights_only=False)
+        self.model.load_state_dict(state["model"])
+        self.optimizer.load_state_dict(state["optimizer"])
+        self.global_step = state.get("global_step", 0)
+        if self.ema is not None and "ema" in state:
+            self.ema.shadow = state["ema"]