pipeline.py

"""
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})"
        )