Upload folder using huggingface_hub

README.md

@@ -0,0 +1,215 @@
+---
+license: apache-2.0
+tags:
+  - biology
+  - genomics
+  - diffusion
+  - single-cell
+  - transcriptomics
+  - microscopy
+  - image-to-rna
+  - conditional-generation
+  - scgpt
+  - vit
+datasets:
+  - altoslabs/scGeneScope
+pipeline_tag: feature-extraction
+model-index:
+  - name: PhenoSeq
+    results:
+      - task:
+          type: feature-extraction
+          name: Image-conditioned RNA-seq generation
+        metrics:
+          - type: loss
+            value: 0.1683
+            name: Validation MSE Loss (epoch 87)
+---
+
+# PhenoSeq: Image-Conditioned Diffusion for Single-Cell Transcriptomics
+
+[![HuggingFace](https://img.shields.io/badge/🤗%20HuggingFace-Sentinal4D/PhenoSeq-yellow)](https://huggingface.co/Sentinal4D/PhenoSeq)
+
+PhenoSeq is a Gaussian diffusion model that **generates scGPT RNA-seq embeddings conditioned on ViT-L microscopy imaging features**. Given fluorescence microscopy images of a cell or well, it predicts a 512-dimensional scGPT embedding representing the transcriptomic state of individual cells — enabling image-to-transcriptome translation at single-cell resolution.
+
+## Model summary
+
+| Property | Value |
+|---|---|
+| Input | ViT-L imaging features — `(N, 5120)` per sample (5 fluorescence channels × 1,024 dims) |
+| Output | scGPT embeddings — `(n_cells, 512)` |
+| Architecture | Cross-attention diffusion denoiser |
+| Diffusion steps | 1,000 (cosine schedule) |
+| Inference steps | 50 (DDIM, default) |
+| Model parameters | ~168 M |
+| Training dataset | [scGeneScope](https://huggingface.co/datasets/altoslabs/scGeneScope) |
+| Best val MSE loss | 0.1683 (epoch 87) |
+
+## Architecture
+
+The denoiser uses a **cross-attention transformer** stack:
+
+1. **Imaging encoder** — 2-layer self-attention transformer projects `(B, N, 5120)` → `(B, N, 1024)` context.
+2. **RNA + time input** — noisy scGPT embedding and sinusoidal time embedding are projected to `model_dim`.
+3. **Cross-attention blocks** (×6) — RNA queries attend to imaging context, with self-attention and adaptive layer normalization conditioned on the timestep.
+4. **Output projection** — predicts noise `ε ∈ ℝ^{512}` for the denoising objective.
+
+The diffusion process uses a cosine beta schedule over T=1,000 steps with EMA weight averaging (decay 0.9999).
+
+## Quick start
+
+```bash
+pip install torch numpy huggingface_hub
+python example.py
+```
+
+```python
+import numpy as np
+from pipeline import PhenoSeqPipeline
+
+pipe = PhenoSeqPipeline.from_pretrained("Sentinal4D/PhenoSeq")
+
+# img_features: ViT-L embeddings — (n_cells, n_imaging_cells=16, 5120)
+img_features = np.random.randn(8, 16, 5120).astype(np.float32)
+rna_predictions = pipe(img_features)   # → (8, 512)
+```
+
+## Inputs and outputs
+
+### Input: imaging features
+
+`img_features` must be ViT-L image embeddings extracted from **5 fluorescence channels** using a ViT-L/14 backbone, resulting in 5 × 1,024 = 5,120 dimensions per imaging cell. Each sample/well typically provides **N = 16** imaging cells (evenly spaced from the available pool) that form the conditioning context.
+
+Shape: `(B, N, 5120)` where `B` is the number of target RNA cells to predict.
+
+> **Imaging normalisation** — `img_norm.npz` contains the per-feature mean and std computed from the training split of scGeneScope. These are applied automatically by `PhenoSeqPipeline`. If you work with a different dataset you will need to recompute and supply your own normalisation stats.
+
+### Output: RNA embeddings
+
+Shape: `(B, 512)` — scGPT-space embeddings un-normalized back to the original scGPT embedding scale. These can be used directly for downstream tasks such as cell-type classification (see `classify_improved.py`), clustering, or trajectory inference.
+
+## Imaging normalisation stats
+
+The pipeline requires `img_norm.npz` (per-feature mean and std from the training split). This file is distributed alongside `best_model.pt` in this repo. If you retrain or use different data, regenerate it:
+
+```bash
+python save_img_norm.py --config config.yaml --output img_norm.npz
+```
+
+## Full inference on scGeneScope data
+
+For large-scale inference over the cached scGeneScope data (`.npz` per sample):
+
+```bash
+# Fast (DDIM, 50 steps)
+python infer.py --checkpoint best_model.pt --ddim_steps 50
+
+# Val split only
+python infer.py --checkpoint best_model.pt --split val --output_dir results/predictions
+
+# Full DDPM sampling (slower, ~1000 steps)
+python infer.py --checkpoint best_model.pt --ddim_steps 0
+```
+
+Output: one `{Sample_ID}.npz` per sample under `results/predictions/`, with key `X` of shape `(n_cells, 512)`.
+
+## Downstream: cell-type classification
+
+Predicted RNA embeddings can be evaluated with the included classifier:
+
+```bash
+python classify_improved.py
+```
+
+See [PERFORMANCE_ANALYSIS.md](PERFORMANCE_ANALYSIS.md) for benchmark results.
+
+## Training
+
+The model was trained from scratch on scGeneScope using:
+
+```bash
+python train.py --config config.yaml
+```
+
+Key hyperparameters (see [config.yaml](config.yaml)):
+
+| Hyperparameter | Value |
+|---|---|
+| Batch size | 256 |
+| Learning rate | 1e-4 (cosine with warmup) |
+| Epochs | 5,000 (best at epoch 87) |
+| Model dim | 1,024 |
+| Attention heads | 8 |
+| Cross-attention layers | 6 |
+| Diffusion steps | 1,000 |
+| Schedule | Cosine |
+| EMA decay | 0.9999 |
+
+## Data
+
+Training and evaluation data come from the [scGeneScope](https://huggingface.co/datasets/altoslabs/scGeneScope) dataset (Altos Labs):
+
+- **Imaging features**: ViT-L/14 embeddings extracted from 5-channel fluorescence microscopy images, stored as `.h5ad` files.
+- **RNA-seq features**: scGPT cell embeddings (512-dim) from paired single-cell RNA-seq, stored as `.h5ad` files.
+- Samples are matched by `Sample_ID` at well level.
+
+Prepare the local cache before training:
+
+```bash
+python prepare_data.py
+```
+
+## Repository structure
+
+```
+PhenoSeq/
+├── pipeline.py          ← self-contained inference pipeline (start here)
+├── best_model.pt        ← trained checkpoint with EMA weights & RNA norm stats
+├── img_norm.npz         ← imaging normalisation stats (mean/std, training split)
+├── config.yaml          ← full training configuration
+├── infer.py             ← batch inference over cached scGeneScope data
+├── save_img_norm.py     ← helper to (re)generate img_norm.npz
+├── train.py             ← training entry point
+├── prepare_data.py      ← extract imaging/RNA features → .npz cache
+├── models/
+│   ├── denoiser.py      ← cross-attention denoiser
+│   ├── diffusion.py     ← Gaussian diffusion process (forward + reverse)
+│   └── lit_module.py    ← PyTorch Lightning training wrapper
+└── data/
+    └── dataset.py       ← paired imaging-RNA dataset & dataloader
+```
+
+## Requirements
+
+```
+torch>=2.0
+anndata>=0.10
+numpy>=1.24
+scipy>=1.10
+PyYAML>=6.0
+tqdm>=4.65
+# Optional: huggingface_hub (for from_pretrained with a Hub repo id)
+```
+
+Install:
+```bash
+pip install -r requirements.txt
+```
+
+## License
+
+Apache 2.0 — see [LICENSE](LICENSE).
+
+## Citation
+
+If you use this model, please cite the scGeneScope dataset:
+
+```bibtex
+@dataset{scgenescope,
+  author    = {Altos Labs},
+  title     = {scGeneScope},
+  year      = {2024},
+  publisher = {HuggingFace},
+  url       = {https://huggingface.co/datasets/altoslabs/scGeneScope}
+}
+```

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4031b1d130e05558cc60027fc05ff5538162f2e12e62eeafa19565e213c8d50
+size 2696336068

example.py

@@ -0,0 +1,36 @@
+"""
+PhenoSeq — minimal inference example.
+
+Downloads the pretrained model from HuggingFace and generates
+scGPT RNA-seq embeddings from synthetic ViT-L imaging features.
+
+Usage:
+    pip install torch numpy huggingface_hub
+    python example.py
+"""
+
+import numpy as np
+from pipeline import PhenoSeqPipeline
+
+# ── Load model from the Hub ────────────────────────────────────────────────────
+pipe = PhenoSeqPipeline.from_pretrained("Sentinal4D/PhenoSeq")
+print(pipe)
+
+# ── Prepare imaging features ───────────────────────────────────────────────────
+# Real use: extract ViT-L/14 embeddings from 5-channel fluorescence microscopy.
+# Shape: (n_cells, n_imaging_cells, 5120)
+#   n_cells        — number of single cells to predict RNA for
+#   n_imaging_cells — imaging cells sampled per well (16 during training)
+#   5120           — 5 fluorescence channels × 1024 ViT-L dims
+n_cells        = 8
+n_imaging_cells = 16
+img_features = np.random.randn(n_cells, n_imaging_cells, 5120).astype(np.float32)
+
+# ── Run inference ──────────────────────────────────────────────────────────────
+# Returns scGPT-space embeddings in the original (denormalized) scale.
+# DDIM with 50 steps by default; pass ddim_steps=0 for full 1000-step DDPM.
+rna_predictions = pipe(img_features)
+
+print(f"\nInput  imaging features : {img_features.shape}")    # (8, 16, 5120)
+print(f"Output RNA embeddings   : {rna_predictions.shape}")  # (8, 512)
+print(f"Output range            : [{rna_predictions.min():.3f}, {rna_predictions.max():.3f}]")

img_norm.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5767ce8ad1c472ef38e478ebb3beb2f9b21ac4e36f64d044e22e5688757b06ad
+size 37563

models/denoiser.py

@@ -0,0 +1,266 @@
+"""
+Cross-Attention Denoiser for Image-conditioned RNA-seq Diffusion.
+
+The denoiser takes:
+  - noisy RNA-seq embeddings (scGPT, 512-dim)
+  - imaging features (ViT-L, 5120-dim) as conditioning context
+  - diffusion timestep
+
+And predicts the noise ε added to the RNA-seq embeddings.
+
+Architecture:
+  1. Project imaging features → model_dim, apply self-attention to create a
+     rich context representation.
+  2. Embed noisy RNA + sinusoidal time embedding → model_dim.
+  3. Multiple cross-attention transformer blocks where RNA queries attend to
+     imaging context.
+  4. Project back to RNA embedding space and predict noise.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.model_utils import (
+    SinusoidalTimeEmbedding,
+    Mish,
+    FeedForward,
+    AdaLayerNorm,
+)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Cross-Attention Block
+# ──────────────────────────────────────────────────────────────────────────────
+
+class CrossAttentionBlock(nn.Module):
+    """
+    Single cross-attention layer: RNA query attends to imaging context.
+    Uses adaptive layer norm for time-conditioning and pre-norm residual style.
+    """
+
+    def __init__(self, dim: int, num_heads: int, time_dim: int, ff_mult: int = 4, dropout: float = 0.1):
+        super().__init__()
+
+        # Time-conditioned norms
+        self.norm_rna = AdaLayerNorm(dim, time_dim)
+        self.norm_ctx = nn.LayerNorm(dim)
+        self.norm_ff = AdaLayerNorm(dim, time_dim)
+
+        # Cross-attention: Q from RNA, K/V from imaging
+        self.cross_attn = nn.MultiheadAttention(
+            embed_dim=dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+
+        # Self-attention on RNA after cross-attention  
+        self.self_attn_norm = AdaLayerNorm(dim, time_dim)
+        self.self_attn = nn.MultiheadAttention(
+            embed_dim=dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+
+        # Feedforward
+        self.ff = FeedForward(dim, mult=ff_mult, dropout=dropout)
+
+    def forward(
+        self,
+        rna: torch.Tensor,       # (B, 1, D)
+        context: torch.Tensor,    # (B, S, D)  imaging context
+        time_emb: torch.Tensor,   # (B, time_dim)
+    ) -> torch.Tensor:
+        # Cross-attention
+        rna_normed = self.norm_rna(rna, time_emb)
+        ctx_normed = self.norm_ctx(context)
+        rna = rna + self.cross_attn(rna_normed, ctx_normed, ctx_normed, need_weights=False)[0]
+
+        # Self-attention (useful when we have multiple RNA tokens, but also adds
+        # a residual self-refinement step)
+        rna_normed = self.self_attn_norm(rna, time_emb)
+        rna = rna + self.self_attn(rna_normed, rna_normed, rna_normed, need_weights=False)[0]
+
+        # Feedforward
+        rna = rna + self.ff(self.norm_ff(rna, time_emb))
+
+        return rna
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Imaging Context Encoder
+# ──────────────────────────────────────────────────────────────────────────────
+
+class ImagingEncoder(nn.Module):
+    """
+    Encodes a set of imaging cell features into context representations.
+    Uses a small self-attention stack to let imaging cells attend to each other
+    before being used as cross-attention context.
+    """
+
+    def __init__(self, img_dim: int, model_dim: int, num_heads: int = 4, num_layers: int = 2, dropout: float = 0.1):
+        super().__init__()
+
+        self.proj = nn.Sequential(
+            nn.Linear(img_dim, model_dim),
+            Mish(),
+            nn.LayerNorm(model_dim),
+            nn.Dropout(dropout),
+        )
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=model_dim,
+            nhead=num_heads,
+            dim_feedforward=model_dim * 4,
+            dropout=dropout,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+
+    def forward(self, img_features: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            img_features: (B, N, img_dim)
+        Returns:
+            (B, N, model_dim) context representations
+        """
+        x = self.proj(img_features)
+        return self.encoder(x)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Full Denoiser
+# ──────────────────────────────────────────────────────────────────────────────
+
+class Img2RNADenoiser(nn.Module):
+    """
+    Noise prediction network for image-conditioned RNA-seq diffusion.
+
+    Args:
+        img_dim:    input imaging feature dimension (5120)
+        rna_dim:    RNA-seq embedding dimension (512)
+        model_dim:  internal model dimension (1024)
+        num_heads:  number of attention heads (8)
+        num_layers: number of cross-attention blocks (6)
+        time_dim:   time embedding dimension (256)
+        ff_mult:    feedforward multiplier (4)
+        dropout:    dropout rate (0.1)
+    """
+
+    def __init__(
+        self,
+        img_dim: int = 5120,
+        rna_dim: int = 512,
+        model_dim: int = 1024,
+        num_heads: int = 8,
+        num_layers: int = 6,
+        time_dim: int = 256,
+        ff_mult: int = 4,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+
+        self.model_dim = model_dim
+        self.rna_dim = rna_dim
+
+        # ── Time embedding ────────────────────────────────────────────────
+        self.time_embed = nn.Sequential(
+            SinusoidalTimeEmbedding(time_dim),
+            nn.Linear(time_dim, time_dim * 4),
+            Mish(),
+            nn.Linear(time_dim * 4, time_dim),
+        )
+
+        # ── Imaging encoder ───────────────────────────────────────────────
+        self.img_encoder = ImagingEncoder(
+            img_dim=img_dim,
+            model_dim=model_dim,
+            num_heads=min(num_heads, 4),
+            num_layers=2,
+            dropout=dropout,
+        )
+
+        # ── RNA input projection ──────────────────────────────────────────
+        self.rna_proj = nn.Sequential(
+            nn.Linear(rna_dim, model_dim),
+            Mish(),
+            nn.LayerNorm(model_dim),
+        )
+
+        # ── Cross-attention transformer stack ─────────────────────────────
+        self.layers = nn.ModuleList([
+            CrossAttentionBlock(
+                dim=model_dim,
+                num_heads=num_heads,
+                time_dim=time_dim,
+                ff_mult=ff_mult,
+                dropout=dropout,
+            )
+            for _ in range(num_layers)
+        ])
+
+        # ── Output projection: predict noise in RNA embedding space ──────
+        self.out_norm = nn.LayerNorm(model_dim)
+        self.out_proj = nn.Sequential(
+            nn.Linear(model_dim, model_dim),
+            Mish(),
+            nn.Linear(model_dim, rna_dim),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """
+        Initialize weights for stable diffusion training.
+
+        Uses PyTorch defaults (kaiming) for all layers, with a small-scale
+        init on the very last linear so the model starts by predicting
+        near-zero noise while still allowing gradient flow.
+        """
+        # Small (not zero) init for the final output linear
+        nn.init.normal_(self.out_proj[2].weight, std=1e-4)
+        nn.init.zeros_(self.out_proj[2].bias)
+
+    def forward(
+        self,
+        noisy_rna: torch.Tensor,    # (B, rna_dim) — noisy scGPT embedding
+        img_features: torch.Tensor,  # (B, N, img_dim) — imaging features
+        timestep: torch.Tensor,      # (B,) — diffusion timestep
+    ) -> torch.Tensor:
+        """
+        Predict the noise component in the noisy RNA embedding.
+
+        Returns:
+            predicted_noise: (B, rna_dim)
+        """
+        # Time embedding
+        t_emb = self.time_embed(timestep)  # (B, time_dim)
+
+        # Encode imaging context
+        context = self.img_encoder(img_features)  # (B, N, model_dim)
+
+        # Project noisy RNA and add as a single token
+        rna = self.rna_proj(noisy_rna).unsqueeze(1)  # (B, 1, model_dim)
+
+        # Cross-attention layers
+        for layer in self.layers:
+            rna = layer(rna, context, t_emb)
+
+        # Output projection
+        rna = self.out_norm(rna.squeeze(1))  # (B, model_dim)
+        noise_pred = self.out_proj(rna)       # (B, rna_dim)
+
+        return noise_pred
+
+    @torch.no_grad()
+    def count_parameters(self) -> dict:
+        """Count trainable and total parameters."""
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return {"total": total, "trainable": trainable}

models/diffusion.py

@@ -0,0 +1,353 @@
+"""
+Gaussian Diffusion for Image-conditioned RNA-seq Generation.
+
+Implements the forward (noise) and reverse (denoise) diffusion processes
+with support for linear and cosine beta schedules.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Beta schedules
+# ──────────────────────────────────────────────────────────────────────────────
+
+def linear_beta_schedule(num_steps: int, beta_start: float = 1e-4, beta_end: float = 0.02) -> torch.Tensor:
+    return torch.linspace(beta_start, beta_end, num_steps, dtype=torch.float64)
+
+
+def cosine_beta_schedule(num_steps: int, s: float = 0.008) -> torch.Tensor:
+    """
+    Cosine schedule as proposed in "Improved Denoising Diffusion Probabilistic Models"
+    (Nichol & Dhariwal, 2021).
+    """
+    steps = torch.arange(num_steps + 1, dtype=torch.float64)
+    f_t = torch.cos(((steps / num_steps) + s) / (1 + s) * (np.pi / 2)) ** 2
+    alphas_cumprod = f_t / f_t[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return torch.clamp(betas, min=1e-5, max=0.999)
+
+
+def get_beta_schedule(schedule: str, num_steps: int, **kwargs) -> torch.Tensor:
+    if schedule == "linear":
+        return linear_beta_schedule(num_steps, **kwargs)
+    elif schedule == "cosine":
+        return cosine_beta_schedule(num_steps)
+    else:
+        raise ValueError(f"Unknown schedule: {schedule}")
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Gaussian Diffusion
+# ──────────────────────────────────────────────────────────────────────────────
+
+class GaussianDiffusion(nn.Module):
+    """
+    Gaussian diffusion process for generating RNA-seq embeddings
+    conditioned on imaging features.
+
+    The model learns to predict the noise ε given:
+      - noisy RNA embedding x_t
+      - conditioning imaging features
+      - timestep t
+
+    Training loss:
+        L = E[||ε - ε_θ(x_t, img, t)||²]
+
+    Args:
+        denoiser: noise prediction network (Img2RNADenoiser)
+        num_steps: number of diffusion steps T
+        schedule: beta schedule type ("linear" or "cosine")
+        beta_start: start of linear schedule
+        beta_end: end of linear schedule
+        loss_type: "l2", "l1", or "huber"
+    """
+
+    def __init__(
+        self,
+        denoiser: nn.Module,
+        num_steps: int = 1000,
+        schedule: str = "cosine",
+        beta_start: float = 1e-4,
+        beta_end: float = 0.02,
+        loss_type: str = "l2",
+        rna_norm: Optional[dict] = None,
+    ):
+        super().__init__()
+
+        self.denoiser = denoiser
+        self.num_steps = num_steps
+        self.loss_type = loss_type
+
+        # Compute noise schedule
+        betas = get_beta_schedule(
+            schedule, num_steps,
+            beta_start=beta_start, beta_end=beta_end,
+        )
+
+        alphas = 1.0 - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0)
+        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
+
+        # Register as buffers (moved to device automatically)
+        self.register_buffer("betas", betas.float())
+        self.register_buffer("alphas", alphas.float())
+        self.register_buffer("alphas_cumprod", alphas_cumprod.float())
+        self.register_buffer("alphas_cumprod_prev", alphas_cumprod_prev.float())
+
+        # Pre-compute useful quantities
+        self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod).float())
+        self.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod).float())
+        self.register_buffer("sqrt_recip_alphas", torch.sqrt(1.0 / alphas).float())
+
+        # Posterior q(x_{t-1} | x_t, x_0) parameters
+        posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+        self.register_buffer("posterior_variance", posterior_variance.float())
+        self.register_buffer("posterior_log_variance_clipped",
+                             torch.log(torch.clamp(posterior_variance, min=1e-20)).float())
+        self.register_buffer("posterior_mean_coef1",
+                             (betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)).float())
+        self.register_buffer("posterior_mean_coef2",
+                             ((1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod)).float())
+
+        # RNA-seq normalization stats for un-normalizing generated samples
+        if rna_norm is not None:
+            self.register_buffer("rna_mean", torch.from_numpy(rna_norm['mean']).float())
+            self.register_buffer("rna_std", torch.from_numpy(rna_norm['std']).float())
+        else:
+            self.rna_mean = None
+            self.rna_std = None
+
+    def unnormalize_rna(self, x: torch.Tensor) -> torch.Tensor:
+        """Convert normalised diffusion output back to original RNA embedding scale."""
+        if self.rna_mean is not None and self.rna_std is not None:
+            return x * self.rna_std + self.rna_mean
+        return x
+
+    # ── Forward diffusion (add noise) ─────────────────────────────────────
+
+    def q_sample(
+        self,
+        x_start: torch.Tensor,
+        t: torch.Tensor,
+        noise: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Sample from q(x_t | x_0) = N(√ᾱ_t · x_0, (1-ᾱ_t) · I).
+
+        Args:
+            x_start: (B, D) clean RNA embeddings
+            t: (B,) timesteps
+            noise: optional pre-sampled noise
+        Returns:
+            x_t: (B, D) noisy embeddings
+        """
+        if noise is None:
+            noise = torch.randn_like(x_start)
+
+        sqrt_alpha = self.sqrt_alphas_cumprod[t]       # (B,)
+        sqrt_one_minus = self.sqrt_one_minus_alphas_cumprod[t]  # (B,)
+
+        # Reshape for broadcasting over feature dim
+        while sqrt_alpha.dim() < x_start.dim():
+            sqrt_alpha = sqrt_alpha.unsqueeze(-1)
+            sqrt_one_minus = sqrt_one_minus.unsqueeze(-1)
+
+        return sqrt_alpha * x_start + sqrt_one_minus * noise
+
+    # ── Training loss ─────────────────────────────────────────────────────
+
+    def compute_loss(
+        self,
+        rna_embedding: torch.Tensor,   # (B, rna_dim) clean target
+        img_features: torch.Tensor,    # (B, N, img_dim) conditioning
+        noise: Optional[torch.Tensor] = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Compute the diffusion training loss.
+
+        Randomly samples timesteps, adds noise, predicts noise, returns loss.
+
+        Returns:
+            dict with "loss" and other metrics
+        """
+        B = rna_embedding.shape[0]
+        device = rna_embedding.device
+
+        # Sample random timesteps
+        t = torch.randint(0, self.num_steps, (B,), device=device).long()
+
+        # Sample noise
+        if noise is None:
+            noise = torch.randn_like(rna_embedding)
+
+        # Add noise
+        x_t = self.q_sample(rna_embedding, t, noise)
+
+        # Predict noise
+        noise_pred = self.denoiser(
+            noisy_rna=x_t,
+            img_features=img_features,
+            timestep=t,
+        )
+
+        # Compute loss
+        if self.loss_type == "l2":
+            loss = F.mse_loss(noise_pred, noise)
+        elif self.loss_type == "l1":
+            loss = F.l1_loss(noise_pred, noise)
+        elif self.loss_type == "huber":
+            loss = F.smooth_l1_loss(noise_pred, noise)
+        else:
+            raise ValueError(f"Unknown loss type: {self.loss_type}")
+
+        return {
+            "loss": loss,
+            "mse": F.mse_loss(noise_pred, noise).detach(),
+        }
+
+    # ── Reverse diffusion (sampling) ──────────────────────────────────────
+
+    @torch.no_grad()
+    def p_sample(
+        self,
+        x_t: torch.Tensor,
+        t: int,
+        img_features: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Single reverse diffusion step: sample x_{t-1} from p_θ(x_{t-1} | x_t).
+
+        Args:
+            x_t: (B, D) current noisy state
+            t: scalar timestep
+            img_features: (B, N, img_dim) conditioning
+        Returns:
+            x_{t-1}: (B, D) less noisy state
+        """
+        B = x_t.shape[0]
+        t_batch = torch.full((B,), t, device=x_t.device, dtype=torch.long)
+
+        # Predict noise
+        noise_pred = self.denoiser(
+            noisy_rna=x_t,
+            img_features=img_features,
+            timestep=t_batch,
+        )
+
+        # Compute posterior mean
+        sqrt_recip_alpha = self.sqrt_recip_alphas[t]
+        beta = self.betas[t]
+        sqrt_one_minus = self.sqrt_one_minus_alphas_cumprod[t]
+
+        mean = sqrt_recip_alpha * (x_t - beta * noise_pred / sqrt_one_minus)
+
+        if t > 0:
+            noise = torch.randn_like(x_t)
+            sigma = torch.sqrt(self.posterior_variance[t])
+            return mean + sigma * noise
+        else:
+            return mean
+
+    @torch.no_grad()
+    def sample(
+        self,
+        img_features: torch.Tensor,   # (B, N, img_dim)
+        shape: Optional[tuple] = None,
+    ) -> torch.Tensor:
+        """
+        Generate RNA-seq embeddings conditioned on imaging features.
+
+        Args:
+            img_features: (B, N, img_dim) conditioning features
+            shape: output shape (B, rna_dim); inferred if None
+        Returns:
+            x_0: (B, rna_dim) generated RNA-seq embeddings
+        """
+        B = img_features.shape[0]
+        device = img_features.device
+
+        if shape is None:
+            shape = (B, self.denoiser.rna_dim)
+
+        # Start from pure noise
+        x = torch.randn(shape, device=device)
+
+        # Iterative denoising
+        for t in reversed(range(self.num_steps)):
+            x = self.p_sample(x, t, img_features)
+
+        return self.unnormalize_rna(x)
+
+    @torch.no_grad()
+    def sample_ddim(
+        self,
+        img_features: torch.Tensor,
+        num_inference_steps: int = 50,
+        eta: float = 0.0,
+        shape: Optional[tuple] = None,
+    ) -> torch.Tensor:
+        """
+        DDIM sampling for faster inference.
+
+        Args:
+            img_features: (B, N, img_dim) 
+            num_inference_steps: number of denoising steps (< num_steps for speed)
+            eta: controls stochasticity (0 = deterministic DDIM, 1 = DDPM)
+            shape: output shape
+        Returns:
+            x_0: (B, rna_dim) generated embeddings
+        """
+        B = img_features.shape[0]
+        device = img_features.device
+
+        if shape is None:
+            shape = (B, self.denoiser.rna_dim)
+
+        # Create sub-sequence of timesteps
+        step_size = self.num_steps // num_inference_steps
+        timesteps = list(range(0, self.num_steps, step_size))[::-1]
+
+        x = torch.randn(shape, device=device)
+
+        for i, t in enumerate(timesteps):
+            t_batch = torch.full((B,), t, device=device, dtype=torch.long)
+
+            noise_pred = self.denoiser(
+                noisy_rna=x,
+                img_features=img_features,
+                timestep=t_batch,
+            )
+
+            # Predict x_0
+            alpha_t = self.alphas_cumprod[t]
+            sqrt_alpha_t = self.sqrt_alphas_cumprod[t]
+            sqrt_one_minus_t = self.sqrt_one_minus_alphas_cumprod[t]
+
+            x_0_pred = (x - sqrt_one_minus_t * noise_pred) / sqrt_alpha_t
+
+            if i < len(timesteps) - 1:
+                t_prev = timesteps[i + 1]
+                alpha_t_prev = self.alphas_cumprod[t_prev]
+            else:
+                alpha_t_prev = torch.tensor(1.0, device=device)
+
+            # DDIM update
+            sigma = eta * torch.sqrt(
+                (1 - alpha_t_prev) / (1 - alpha_t) * (1 - alpha_t / alpha_t_prev)
+            )
+
+            pred_dir = torch.sqrt(1 - alpha_t_prev - sigma ** 2) * noise_pred
+            x = torch.sqrt(alpha_t_prev) * x_0_pred + pred_dir
+
+            if sigma > 0 and i < len(timesteps) - 1:
+                x = x + sigma * torch.randn_like(x)
+
+        return self.unnormalize_rna(x)

models/lit_module.py

@@ -0,0 +1,288 @@
+"""
+PyTorch Lightning LightningModule for Img2RNA diffusion training.
+
+Wraps the GaussianDiffusion model with:
+  - Automatic optimizer / scheduler setup
+  - Training / validation step definitions
+  - EMA weight averaging
+  - Periodic sample generation for qualitative monitoring
+"""
+
+from __future__ import annotations
+
+import copy
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+
+from models.denoiser import Img2RNADenoiser
+from models.diffusion import GaussianDiffusion
+
+
+class EMA:
+    """Exponential Moving Average of model parameters (Lightning-compatible)."""
+
+    def __init__(self, model: nn.Module, decay: float = 0.9999):
+        self.decay = decay
+        self.shadow = {
+            name: p.clone().detach()
+            for name, p in model.named_parameters()
+            if p.requires_grad
+        }
+        self._backup: dict[str, torch.Tensor] = {}
+
+    @torch.no_grad()
+    def update(self, model: nn.Module):
+        for name, p in model.named_parameters():
+            if p.requires_grad and name in self.shadow:
+                self.shadow[name].lerp_(p.data, 1 - self.decay)
+
+    def apply(self, model: nn.Module):
+        self._backup = {
+            name: p.data.clone()
+            for name, p in model.named_parameters()
+            if name in self.shadow
+        }
+        for name, p in model.named_parameters():
+            if name in self.shadow:
+                p.data.copy_(self.shadow[name])
+
+    def restore(self, model: nn.Module):
+        for name, p in model.named_parameters():
+            if name in self._backup:
+                p.data.copy_(self._backup[name])
+        self._backup.clear()
+
+    def state_dict(self):
+        return {k: v.cpu() for k, v in self.shadow.items()}
+
+    def load_state_dict(self, state_dict: dict):
+        self.shadow = {k: v.clone() for k, v in state_dict.items()}
+
+
+class Img2RNALitModule(pl.LightningModule):
+    """
+    Lightning module for image-conditioned RNA-seq diffusion.
+
+    Handles training, validation, EMA, and optional sample generation.
+    Works with any Lightning logger (TensorBoard, wandb, CSV, etc.).
+    """
+
+    def __init__(
+        self,
+        # Model config
+        img_dim: int = 5120,
+        rna_dim: int = 512,
+        model_dim: int = 1024,
+        num_heads: int = 8,
+        num_layers: int = 6,
+        time_dim: int = 256,
+        ff_mult: int = 4,
+        dropout: float = 0.1,
+        # Diffusion config
+        num_steps: int = 1000,
+        schedule: str = "cosine",
+        beta_start: float = 1e-4,
+        beta_end: float = 0.02,
+        loss_type: str = "l2",
+        # Normalization stats
+        rna_norm: Optional[dict] = None,
+        # Training config
+        lr: float = 1e-4,
+        weight_decay: float = 1e-5,
+        warmup_epochs: int = 5,
+        scheduler_type: str = "cosine",
+        max_grad_norm: float = 1.0,
+        ema_decay: float = 0.9999,
+        # Sampling config (for validation visualisation)
+        val_sample_every_n_epochs: int = 10,
+        ddim_steps: int = 50,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+        # Build denoiser
+        self.denoiser = Img2RNADenoiser(
+            img_dim=img_dim,
+            rna_dim=rna_dim,
+            model_dim=model_dim,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            time_dim=time_dim,
+            ff_mult=ff_mult,
+            dropout=dropout,
+        )
+
+        # Build diffusion wrapper
+        self.diffusion = GaussianDiffusion(
+            denoiser=self.denoiser,
+            num_steps=num_steps,
+            schedule=schedule,
+            beta_start=beta_start,
+            beta_end=beta_end,
+            loss_type=loss_type,
+            rna_norm=rna_norm,
+        )
+
+        # EMA (initialised in on_fit_start so it lives on correct device)
+        self._ema: Optional[EMA] = None
+        self._ema_decay = ema_decay
+
+        # Cache for logging
+        self._val_step_outputs: list[dict] = []
+
+    # ── Lifecycle hooks ───────────────────────────────────────────────────
+
+    def on_fit_start(self):
+        self._ema = EMA(self.diffusion, decay=self._ema_decay)
+        param_info = self.denoiser.count_parameters()
+        self.log_dict({
+            "model/total_params": float(param_info["total"]),
+            "model/trainable_params": float(param_info["trainable"]),
+        })
+
+    # ── Training ──────────────────────────────────────────────────────────
+
+    def training_step(self, batch: dict, batch_idx: int) -> torch.Tensor:
+        img_features = batch["img_features"]   # (B, N, img_dim)
+        rna_embedding = batch["rna_embedding"]  # (B, rna_dim)
+
+        result = self.diffusion.compute_loss(rna_embedding, img_features)
+        loss = result["loss"]
+
+        # Log metrics
+        self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True)
+        self.log("train/mse", result["mse"], on_step=False, on_epoch=True)
+        self.log("train/lr", self.optimizers().param_groups[0]["lr"], on_step=True, on_epoch=False)
+
+        return loss
+
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        if self._ema is not None:
+            self._ema.update(self.diffusion)
+
+    # ── Validation ────────────────────────────────────────────────────────
+
+    def validation_step(self, batch: dict, batch_idx: int) -> dict:
+        img_features = batch["img_features"]
+        rna_embedding = batch["rna_embedding"]
+
+        result = self.diffusion.compute_loss(rna_embedding, img_features)
+
+        self.log("val/loss", result["loss"], on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
+        self.log("val/mse", result["mse"], on_step=False, on_epoch=True, sync_dist=True)
+
+        self._val_step_outputs.append({
+            "loss": result["loss"].detach(),
+            "mse": result["mse"].detach(),
+        })
+
+        return result
+
+    def on_validation_epoch_end(self):
+        # Optional: generate samples for qualitative monitoring
+        if (
+            self.current_epoch > 0
+            and self.current_epoch % self.hparams.val_sample_every_n_epochs == 0
+        ):
+            self._log_generated_samples()
+
+        self._val_step_outputs.clear()
+
+    @torch.no_grad()
+    def _log_generated_samples(self, n_samples: int = 8):
+        """Generate a few samples and log statistics to the logger."""
+        # Use EMA weights for sampling
+        if self._ema is not None:
+            self._ema.apply(self.diffusion)
+
+        try:
+            # Create dummy imaging conditioning (random from normal; in practice
+            # you'd use real validation imaging features)
+            dummy_img = torch.randn(
+                n_samples,
+                self.hparams.get("n_imaging_cells", 16),
+                self.hparams.img_dim,
+                device=self.device,
+            )
+            generated = self.diffusion.sample_ddim(
+                dummy_img,
+                num_inference_steps=self.hparams.ddim_steps,
+            )
+
+            self.log("val/generated_mean", generated.mean())
+            self.log("val/generated_std", generated.std())
+            self.log("val/generated_min", generated.min())
+            self.log("val/generated_max", generated.max())
+
+        finally:
+            if self._ema is not None:
+                self._ema.restore(self.diffusion)
+
+    # ── Optimizer & scheduler ─────────────────────────────────────────────
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            self.parameters(),
+            lr=self.hparams.lr,
+            weight_decay=self.hparams.weight_decay,
+            betas=(0.9, 0.999),
+        )
+
+        if self.hparams.scheduler_type == "cosine":
+            # Total steps will be set by trainer
+            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer,
+                T_max=self.trainer.estimated_stepping_batches,
+                eta_min=self.hparams.lr * 0.01,
+            )
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": scheduler,
+                    "interval": "step",
+                    "frequency": 1,
+                },
+            }
+        else:
+            return optimizer
+
+    def on_before_optimizer_step(self, optimizer):
+        # Gradient clipping
+        if self.hparams.max_grad_norm > 0:
+            nn.utils.clip_grad_norm_(self.parameters(), self.hparams.max_grad_norm)
+
+    # ── Inference helpers ─────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def generate(
+        self,
+        img_features: torch.Tensor,
+        use_ema: bool = True,
+        ddim: bool = True,
+        ddim_steps: int = 50,
+    ) -> torch.Tensor:
+        """
+        Generate RNA-seq embeddings from imaging features.
+
+        Args:
+            img_features: (B, N, img_dim)
+            use_ema: whether to use EMA weights
+            ddim: use DDIM sampling (faster)
+            ddim_steps: number of DDIM steps
+        Returns:
+            (B, rna_dim) generated embeddings
+        """
+        if use_ema and self._ema is not None:
+            self._ema.apply(self.diffusion)
+
+        try:
+            if ddim:
+                return self.diffusion.sample_ddim(img_features, num_inference_steps=ddim_steps)
+            else:
+                return self.diffusion.sample(img_features)
+        finally:
+            if use_ema and self._ema is not None:
+                self._ema.restore(self.diffusion)

models/model_utils.py

@@ -0,0 +1,123 @@
+"""
+Shared model utilities: sinusoidal embeddings, activation functions, blocks.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Activations
+# ──────────────────────────────────────────────────────────────────────────────
+
+class Mish(nn.Module):
+    """Mish activation: x * tanh(softplus(x))."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.tanh(F.softplus(x))
+
+
+class GEGLU(nn.Module):
+    """Gated GELU activation for feedforward blocks."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, gate = x.chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Sinusoidal positional / time embedding
+# ──────────────────────────────────────────────────────────────────────────────
+
+class SinusoidalTimeEmbedding(nn.Module):
+    """
+    Maps scalar timesteps to sinusoidal embeddings.
+
+    Args:
+        dim: output embedding dimension (must be even)
+        max_period: controls frequency range
+    """
+
+    def __init__(self, dim: int, max_period: int = 10000):
+        super().__init__()
+        assert dim % 2 == 0, "dim must be even"
+        self.dim = dim
+        self.max_period = max_period
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            t: (B,) integer or float timesteps
+        Returns:
+            (B, dim) sinusoidal embeddings
+        """
+        half = self.dim // 2
+        freqs = torch.exp(
+            -math.log(self.max_period)
+            * torch.arange(half, device=t.device, dtype=torch.float32)
+            / half
+        )
+        args = t[:, None].float() * freqs[None, :]
+        return torch.cat([args.sin(), args.cos()], dim=-1)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Building blocks
+# ──────────────────────────────────────────────────────────────────────────────
+
+class FeedForward(nn.Module):
+    """
+    Transformer feedforward block with GEGLU gating.
+
+    Args:
+        dim: input/output dimension
+        mult: hidden dimension multiplier  
+        dropout: dropout rate
+    """
+
+    def __init__(self, dim: int, mult: int = 4, dropout: float = 0.1):
+        super().__init__()
+        inner_dim = dim * mult * 2  # ×2 for GEGLU split
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, inner_dim),
+            GEGLU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * mult, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class AdaLayerNorm(nn.Module):
+    """
+    Adaptive Layer Normalization conditioned on time embedding.
+    Applies scale/shift modulation: γ(t) * LayerNorm(x) + β(t)
+
+    Args:
+        dim: feature dimension
+        cond_dim: conditioning (time) embedding dimension
+    """
+
+    def __init__(self, dim: int, cond_dim: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.proj = nn.Sequential(
+            Mish(),
+            nn.Linear(cond_dim, dim * 2),
+        )
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (B, L, D)  or (B, D)
+            cond: (B, cond_dim)
+        """
+        gamma, beta = self.proj(cond).chunk(2, dim=-1)
+        if x.dim() == 3 and gamma.dim() == 2:
+            gamma = gamma.unsqueeze(1)
+            beta = beta.unsqueeze(1)
+        return self.norm(x) * (1 + gamma) + beta

pipeline.py

@@ -0,0 +1,288 @@
+"""
+PhenoSeq inference pipeline.
+
+Self-contained entry point for generating scGPT RNA-seq embeddings from
+ViT-L microscopy imaging features using a pretrained diffusion model.
+
+Quick start
+-----------
+>>> import numpy as np
+>>> from pipeline import PhenoSeqPipeline
+>>>
+>>> # Load from a local directory (or a HuggingFace repo id)
+>>> pipe = PhenoSeqPipeline.from_pretrained(".")
+>>>
+>>> # img_features: (n_cells, n_imaging_cells, 5120) — raw ViT-L embeddings
+>>> # img_norm is loaded automatically from img_norm.npz when present
+>>> rna = pipe(img_features)   # → np.ndarray (n_cells, 512)
+
+Input format
+------------
+img_features : np.ndarray or torch.Tensor
+    Shape (B, N, 5120) where
+      B = number of target RNA cells to generate
+      N = number of imaging cells per well (default 16)
+    Features should be raw (unnormalized) ViT-L embeddings when img_norm
+    is available, or pre-normalized when img_norm is None.
+
+Output format
+-------------
+np.ndarray of shape (B, 512) — scGPT embedding space predictions,
+un-normalized back to the original scGPT scale.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from pathlib import Path
+from typing import Optional, Union
+
+import numpy as np
+import torch
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Pipeline
+# ──────────────────────────────────────────────────────────────────────────────
+
+class PhenoSeqPipeline:
+    """
+    Wraps the trained PhenoSeq diffusion model for single-cell RNA prediction.
+
+    Parameters
+    ----------
+    diffusion : GaussianDiffusion
+        Loaded diffusion model (weights applied, eval mode).
+    img_norm : dict with 'mean' and 'std' arrays of shape (5120,), optional
+        Imaging normalisation statistics computed from the training split.
+        If None, input features are assumed to be pre-normalized.
+    device : str or torch.device
+        Device for inference ('cuda', 'cpu', etc.).
+    ddim_steps : int
+        Number of DDIM denoising steps (50 is a good default; 0 = full DDPM).
+    """
+
+    def __init__(
+        self,
+        diffusion,
+        img_norm: Optional[dict] = None,
+        device: Union[str, torch.device] = "cpu",
+        ddim_steps: int = 50,
+    ):
+        self.diffusion = diffusion
+        self.img_norm = img_norm
+        self.device = torch.device(device)
+        self.ddim_steps = ddim_steps
+        self.diffusion.to(self.device).eval()
+
+    # ── Factory ───────────────────────────────────────────────────────────
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_dir: Union[str, Path] = ".",
+        checkpoint_name: str = "best_model.pt",
+        img_norm_name: str = "img_norm.npz",
+        device: Optional[Union[str, torch.device]] = None,
+        ddim_steps: int = 50,
+        use_ema: bool = True,
+    ) -> "PhenoSeqPipeline":
+        """
+        Load a pipeline from a local directory or HuggingFace Hub repo.
+
+        Parameters
+        ----------
+        model_dir : str or Path
+            Local path or HuggingFace Hub repo id containing the checkpoint.
+        checkpoint_name : str
+            Filename of the PyTorch checkpoint inside model_dir.
+        img_norm_name : str
+            Filename of the imaging normalisation stats (.npz with 'mean', 'std').
+            If the file is not found, img_norm is set to None and a warning is logged.
+        device : str, torch.device, or None
+            Target device; auto-selects CUDA when available if None.
+        ddim_steps : int
+            DDIM sampling steps.
+        use_ema : bool
+            Prefer EMA weights when available in the checkpoint (recommended).
+        """
+        model_dir = Path(model_dir)
+
+        # ── Resolve from HuggingFace Hub if path doesn't exist locally ────
+        if not model_dir.exists():
+            try:
+                from huggingface_hub import snapshot_download
+                model_dir = Path(snapshot_download(str(model_dir)))
+                logger.info(f"Downloaded from HuggingFace Hub → {model_dir}")
+            except Exception as exc:
+                raise FileNotFoundError(
+                    f"Directory '{model_dir}' not found locally and Hub download failed: {exc}"
+                ) from exc
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = torch.device(device)
+
+        # ── Load checkpoint ────────────────────────���──────────────────────
+        ckpt_path = model_dir / checkpoint_name
+        if not ckpt_path.exists():
+            raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
+
+        logger.info(f"Loading checkpoint from {ckpt_path}")
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+
+        # ── Build model from embedded config ──────────────────────────────
+        cfg = ckpt.get("config", {})
+        mc = ckpt.get("model_cfg", cfg.get("model", {}))
+        dc = cfg.get("diffusion", {})
+
+        from models.denoiser import Img2RNADenoiser
+        from models.diffusion import GaussianDiffusion
+
+        # RNA normalisation is stored in the checkpoint
+        rna_norm_raw = ckpt.get("rna_norm")
+        rna_norm = (
+            {"mean": np.array(rna_norm_raw["mean"]), "std": np.array(rna_norm_raw["std"])}
+            if rna_norm_raw is not None else None
+        )
+
+        denoiser = Img2RNADenoiser(
+            img_dim   = mc.get("img_dim",    5120),
+            rna_dim   = mc.get("rna_dim",    512),
+            model_dim = mc.get("model_dim",  1024),
+            num_heads = mc.get("num_heads",  8),
+            num_layers= mc.get("num_layers", 6),
+            time_dim  = mc.get("time_dim",   256),
+            ff_mult   = mc.get("ff_mult",    4),
+            dropout   = 0.0,
+        )
+
+        diffusion = GaussianDiffusion(
+            denoiser   = denoiser,
+            num_steps  = dc.get("num_steps",  1000),
+            schedule   = dc.get("schedule",   "cosine"),
+            beta_start = dc.get("beta_start", 1e-4),
+            beta_end   = dc.get("beta_end",   0.02),
+            rna_norm   = rna_norm,
+        )
+
+        # ── Load weights ──────────────────────────────────────────────────
+        ema_state = ckpt.get("ema_state_dict")
+        if use_ema and ema_state is not None:
+            logger.info("Loading EMA weights")
+            diffusion.denoiser.load_state_dict(
+                {k.removeprefix("denoiser."): v for k, v in ema_state.items()},
+                strict=False,
+            )
+        else:
+            if use_ema and ema_state is None:
+                logger.warning("EMA weights not found; using standard model weights")
+            diffusion.load_state_dict(ckpt["model_state_dict"], strict=True)
+
+        n_params = sum(p.numel() for p in denoiser.parameters())
+        logger.info(f"Model loaded ({n_params:,} parameters) on {device}")
+
+        # ── Imaging normalisation stats ───────────────────────────────────
+        img_norm = None
+        norm_path = model_dir / img_norm_name
+        if norm_path.exists():
+            data = np.load(norm_path)
+            img_norm = {"mean": data["mean"], "std": data["std"]}
+            logger.info(f"Loaded imaging normalisation stats from {norm_path}")
+        else:
+            logger.warning(
+                f"{img_norm_name} not found in {model_dir}. "
+                "Pass pre-normalized imaging features, or provide img_norm manually. "
+                "See save_img_norm.py to generate this file from your training data."
+            )
+
+        return cls(diffusion, img_norm=img_norm, device=device, ddim_steps=ddim_steps)  # type: ignore[return-value]
+
+    # ── Inference ─────────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        img_features: Union[np.ndarray, torch.Tensor],
+        batch_size: int = 256,
+        ddim_steps: Optional[int] = None,
+    ) -> np.ndarray:
+        """
+        Generate scGPT RNA-seq embeddings from imaging features.
+
+        Parameters
+        ----------
+        img_features : array-like, shape (B, N, 5120) or (N, 5120)
+            ViT-L imaging embeddings. When img_norm is available these should
+            be raw (unnormalized); otherwise provide pre-normalized features.
+            If 2D (N, 5120), a batch dimension is added automatically.
+        batch_size : int
+            Number of cells processed per forward pass.
+        ddim_steps : int, optional
+            Override the pipeline's default DDIM steps for this call.
+            Pass 0 to use full DDPM sampling (slower, ~1000 steps).
+
+        Returns
+        -------
+        np.ndarray of shape (B, 512)
+            Predicted scGPT-space RNA embeddings in the original (denormalized) scale.
+        """
+        steps = ddim_steps if ddim_steps is not None else self.ddim_steps
+
+        # Coerce to numpy then torch
+        if isinstance(img_features, torch.Tensor):
+            img_np = img_features.cpu().float().numpy()
+        else:
+            img_np = np.asarray(img_features, dtype=np.float32)
+
+        # Add batch dim if single sample
+        if img_np.ndim == 2:
+            img_np = img_np[np.newaxis]
+
+        if img_np.ndim != 3 or img_np.shape[-1] != 5120:
+            raise ValueError(
+                f"Expected img_features shape (B, N, 5120), got {img_np.shape}"
+            )
+
+        # Normalize imaging features if stats are available
+        if self.img_norm is not None:
+            mean = self.img_norm["mean"].astype(np.float32)  # (5120,)
+            std  = self.img_norm["std"].astype(np.float32)   # (5120,)
+            img_np = (img_np - mean) / std
+
+        # Batch inference
+        all_preds: list[np.ndarray] = []
+        for start in range(0, len(img_np), batch_size):
+            chunk = torch.from_numpy(img_np[start : start + batch_size]).to(self.device)
+
+            if steps > 0:
+                preds = self.diffusion.sample_ddim(
+                    img_features        = chunk,
+                    num_inference_steps = steps,
+                    eta                 = 0.0,
+                )
+            else:
+                preds = self.diffusion.sample(chunk)
+
+            all_preds.append(preds.cpu().float().numpy())
+
+        return np.concatenate(all_preds, axis=0)  # (B, 512)
+
+    # ── Convenience ───────────────────────────────────────────────────────
+
+    @property
+    def rna_dim(self) -> int:
+        return self.diffusion.denoiser.rna_dim
+
+    def __repr__(self) -> str:
+        d = self.diffusion.denoiser
+        return (
+            f"PhenoSeqPipeline("
+            f"rna_dim={d.rna_dim}, "
+            f"model_dim={d.model_dim}, num_layers={len(d.layers)}, "
+            f"ddim_steps={self.ddim_steps}, device={self.device})"
+        )