Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,99 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00013_test_0_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00013_test_0_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00013_test_0_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00013_test_0_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00013_test_0_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00013_test_0_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00198_test_0_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00198_test_0_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00198_test_0_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00198_test_0_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00198_test_0_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_0_00198_test_0_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00276_test_1+_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00276_test_1+_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00276_test_1+_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00276_test_1+_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00276_test_1+_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00276_test_1+_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00791_test_1+_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00791_test_1+_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00791_test_1+_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00791_test_1+_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00791_test_1+_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_1+_00791_test_1+_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00259_test_2+_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00259_test_2+_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00259_test_2+_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00259_test_2+_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00259_test_2+_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00259_test_2+_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00293_test_2+_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00293_test_2+_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00293_test_2+_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00293_test_2+_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00293_test_2+_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_2+_00293_test_2+_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00220_test_3+_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00220_test_3+_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00220_test_3+_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00220_test_3+_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00220_test_3+_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00220_test_3+_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00277_test_3+_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00277_test_3+_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00277_test_3+_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00277_test_3+_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00277_test_3+_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/BCI_HER2_3+_00277_test_3+_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_35M2101733_7_3_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_35M2101733_7_3_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_35M2101733_7_3_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_35M2101733_7_3_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_35M2101733_7_3_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_35M2101733_7_3_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_40M2101566_28_8_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_40M2101566_28_8_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_40M2101566_28_8_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_40M2101566_28_8_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_40M2101566_28_8_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_ER_40M2101566_28_8_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_19M2102438_35_28_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_19M2102438_35_28_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_19M2102438_35_28_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_19M2102438_35_28_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_19M2102438_35_28_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_19M2102438_35_28_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_67M2100642_15_18_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_67M2100642_15_18_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_67M2100642_15_18_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_67M2100642_15_18_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_67M2100642_15_18_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_HER2_67M2100642_15_18_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_10M2102916_10_20_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_10M2102916_10_20_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_10M2102916_10_20_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_10M2102916_10_20_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_10M2102916_10_20_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_10M2102916_10_20_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_80M2100377_14_29_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_80M2100377_14_29_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_80M2100377_14_29_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_80M2100377_14_29_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_80M2100377_14_29_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_Ki67_80M2100377_14_29_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_17M2102569_15_14_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_17M2102569_15_14_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_17M2102569_15_14_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_17M2102569_15_14_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_17M2102569_15_14_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_17M2102569_15_14_he.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_28M2101987_14_22_gen_er.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_28M2101987_14_22_gen_her2.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_28M2101987_14_22_gen_ki67.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_28M2101987_14_22_gen_pr.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_28M2101987_14_22_gt.png filter=lfs diff=lfs merge=lfs -text
+gallery/images/MIST_PR_28M2101987_14_22_he.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,28 @@
 ---
-title: UNIStainNet
-emoji: 👁
-colorFrom: yellow
+title: UNIStainNet - Virtual IHC Staining
+emoji: 🔬
+colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 6.10.0
+sdk_version: "6.10.0"
 app_file: app.py
 pinned: false
+license: mit
+hardware: cpu-basic
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# UNIStainNet: Foundation-Model-Guided Virtual Staining
+
+Virtual staining of H&E histopathology images to IHC (HER2, Ki67, ER, PR) using a single unified 42M-parameter SPADE-UNet conditioned on dense spatial tokens from a frozen UNI pathology foundation model.
+
+## Features
+- **Upload** an H&E image and generate IHC stains in real-time
+- **Cross-stain comparison**: Generate all 4 stains from a single input
+- **Gallery**: Browse pre-computed examples (no GPU needed)
+
+## Architecture
+| Component | Details |
+|-----------|---------|
+| Generator | SPADE-UNet with UNI spatial conditioning + FiLM stain embeddings |
+| UNI Features | 4x4 sub-crop tiling → UNI ViT-L/16 → 32x32 spatial tokens (1024-dim) |
+| Parameters | 42M (generator), UNI frozen (303M) |

app.py

@@ -0,0 +1,387 @@
+#!/usr/bin/env python3
+"""
+UNIStainNet Interactive Demo — Hugging Face Spaces (ZeroGPU)
+
+Virtual staining of H&E histopathology images to IHC (HER2, Ki67, ER, PR).
+Uses @spaces.GPU for on-demand GPU allocation on ZeroGPU.
+"""
+
+import json
+import os
+import time
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import torch
+
+# ZeroGPU support: use @spaces.GPU if available, otherwise no-op
+try:
+    import spaces
+    GPU_AVAILABLE = torch.cuda.is_available()
+except ImportError:
+    spaces = None
+    GPU_AVAILABLE = torch.cuda.is_available()
+
+
+def _gpu_decorator(duration=60):
+    """Apply @spaces.GPU if available, otherwise return identity decorator."""
+    if spaces is not None and hasattr(spaces, "GPU"):
+        return spaces.GPU(duration=duration)
+    return lambda fn: fn
+import torch.nn.functional as F
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+from PIL import Image
+from huggingface_hub import hf_hub_download
+
+from src.models.trainer import UNIStainNetTrainer
+from src.data.mist_dataset import STAIN_TO_LABEL, LABEL_TO_STAIN
+
+# ── Constants ────────────────────────────────────────────────────────
+STAIN_NAMES = ["HER2", "Ki67", "ER", "PR"]
+GALLERY_DIR = Path(__file__).parent / "gallery"
+TARGET_SIZE = 512
+
+# Model repo where checkpoint is stored (uploaded separately)
+MODEL_REPO = os.environ.get("MODEL_REPO", "faceless-void/UNIStainNet")
+CHECKPOINT_FILENAME = "mist_multistain_last.ckpt"
+
+# ── Global model cache (loaded lazily on GPU request) ────────────────
+_model_cache = {"model": None, "uni_model": None, "spatial_pool_size": 32}
+
+
+def _get_checkpoint_path():
+    """Download checkpoint from HF Hub if not local."""
+    local_path = Path(__file__).parent / "checkpoints" / CHECKPOINT_FILENAME
+    if local_path.exists():
+        return str(local_path)
+    # Download from HF model repo
+    return hf_hub_download(repo_id=MODEL_REPO, filename=CHECKPOINT_FILENAME)
+
+
+def _load_models():
+    """Load UNIStainNet + UNI (called inside @spaces.GPU function)."""
+    if _model_cache["model"] is None:
+        import timm
+
+        ckpt_path = _get_checkpoint_path()
+        print(f"Loading UNIStainNet from {ckpt_path} ...")
+        model = UNIStainNetTrainer.load_from_checkpoint(ckpt_path, strict=False)
+        model = model.cuda().eval()
+        _model_cache["model"] = model
+        _model_cache["spatial_pool_size"] = getattr(
+            model.hparams, "uni_spatial_size", 32
+        )
+        print("  Generator loaded")
+
+        print("Loading UNI ViT-L/16 ...")
+        uni_model = timm.create_model(
+            "hf-hub:MahmoodLab/uni",
+            pretrained=True,
+            init_values=1e-5,
+            dynamic_img_size=True,
+        )
+        uni_model = uni_model.cuda().eval()
+        _model_cache["uni_model"] = uni_model
+        print("  UNI loaded")
+    else:
+        # Models already loaded — move to current GPU device
+        _model_cache["model"] = _model_cache["model"].cuda()
+        _model_cache["uni_model"] = _model_cache["uni_model"].cuda()
+
+    return _model_cache["model"], _model_cache["uni_model"], _model_cache["spatial_pool_size"]
+
+
+# ── Preprocessing helpers ────────────────────────────────────────────
+
+def preprocess_he(pil_image, target_size=TARGET_SIZE):
+    """Center-crop and resize H&E to target_size x target_size."""
+    w, h = pil_image.size
+    short = min(w, h)
+    left = (w - short) // 2
+    top = (h - short) // 2
+    pil_image = pil_image.crop((left, top, left + short, top + short))
+    if short != target_size:
+        pil_image = pil_image.resize((target_size, target_size), Image.BICUBIC)
+    return pil_image
+
+
+def pil_to_tensor(pil_image):
+    """PIL → [1, 3, H, W] in [-1, 1]."""
+    t = TF.to_tensor(pil_image)
+    t = TF.normalize(t, [0.5] * 3, [0.5] * 3)
+    return t.unsqueeze(0)
+
+
+def tensor_to_pil(tensor):
+    """[1, 3, H, W] in [-1, 1] → PIL."""
+    t = ((tensor[0].cpu() + 1) / 2).clamp(0, 1)
+    return TF.to_pil_image(t)
+
+
+def extract_uni_features(uni_model, he_tensor_01, spatial_pool_size=32):
+    """Extract UNI spatial features from H&E crop ([1,3,H,W] in [0,1])."""
+    uni_transform = T.Normalize(
+        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+    )
+    B = he_tensor_01.shape[0]
+    num_crops = 4
+    patches_per_side = 14
+    crop_h = he_tensor_01.shape[2] // num_crops
+    crop_w = he_tensor_01.shape[3] // num_crops
+
+    sub_crops = []
+    for i in range(num_crops):
+        for j in range(num_crops):
+            sub = he_tensor_01[
+                :, :, i * crop_h : (i + 1) * crop_h, j * crop_w : (j + 1) * crop_w
+            ]
+            sub = F.interpolate(sub, size=(224, 224), mode="bicubic", align_corners=False)
+            sub = torch.stack([uni_transform(s) for s in sub])
+            sub_crops.append(sub)
+
+    all_crops = torch.stack(sub_crops, dim=1).reshape(B * 16, 3, 224, 224).cuda()
+
+    with torch.no_grad():
+        all_feats = uni_model.forward_features(all_crops)
+        patch_tokens = all_feats[:, 1:, :]
+
+    patch_tokens = patch_tokens.reshape(
+        B, num_crops, num_crops, patches_per_side, patches_per_side, 1024
+    )
+    full_size = num_crops * patches_per_side
+    full_grid = patch_tokens.permute(0, 1, 3, 2, 4, 5).reshape(
+        B, full_size, full_size, 1024
+    )
+
+    S = spatial_pool_size
+    if S < full_size:
+        grid_bchw = full_grid.permute(0, 3, 1, 2)
+        pooled = F.adaptive_avg_pool2d(grid_bchw, S)
+        result = pooled.permute(0, 2, 3, 1)
+    else:
+        result = full_grid
+
+    return result.reshape(B, S * S, 1024)
+
+
+# ── GPU-accelerated inference functions ──────────────────────────────
+
+@_gpu_decorator(duration=60)
+def generate_single_stain(image, stain, guidance_scale):
+    """Generate a single IHC stain from an H&E image (GPU)."""
+    if image is None:
+        return None, "No image uploaded"
+
+    t0 = time.time()
+    model, uni_model, spatial_pool_size = _load_models()
+
+    he_pil = preprocess_he(image)
+    he_tensor = pil_to_tensor(he_pil).cuda()
+    he_01 = ((he_tensor + 1) / 2).clamp(0, 1)
+
+    uni_feats = extract_uni_features(uni_model, he_01, spatial_pool_size).cuda()
+    label = STAIN_TO_LABEL[stain]
+    labels = torch.tensor([label], device="cuda", dtype=torch.long)
+
+    with torch.no_grad():
+        gen = model.generate(he_tensor, uni_feats, labels, guidance_scale=guidance_scale)
+
+    result = tensor_to_pil(gen)
+    elapsed = time.time() - t0
+    return result, f"{elapsed:.2f}s"
+
+
+@_gpu_decorator(duration=120)
+def generate_all_stains(image, guidance_scale):
+    """Generate all 4 IHC stains from one H&E image (GPU)."""
+    if image is None:
+        return None, None, None, None, None, "No image uploaded"
+
+    t0 = time.time()
+    model, uni_model, spatial_pool_size = _load_models()
+
+    he_pil = preprocess_he(image)
+    he_tensor = pil_to_tensor(he_pil).cuda()
+    he_01 = ((he_tensor + 1) / 2).clamp(0, 1)
+
+    uni_feats = extract_uni_features(uni_model, he_01, spatial_pool_size).cuda()
+
+    results = {}
+    for stain in STAIN_NAMES:
+        label = STAIN_TO_LABEL[stain]
+        labels = torch.tensor([label], device="cuda", dtype=torch.long)
+        with torch.no_grad():
+            gen = model.generate(
+                he_tensor, uni_feats, labels, guidance_scale=guidance_scale
+            )
+        results[stain] = tensor_to_pil(gen)
+
+    elapsed = time.time() - t0
+    return (
+        he_pil,
+        results["HER2"],
+        results["Ki67"],
+        results["ER"],
+        results["PR"],
+        f"{elapsed:.2f}s",
+    )
+
+
+# ── Gallery helpers ──────────────────────────────────────────────────
+
+def load_gallery():
+    meta_path = GALLERY_DIR / "metadata.json"
+    if not meta_path.exists():
+        return None
+    with open(meta_path) as f:
+        return json.load(f)
+
+
+def show_gallery(name, gallery):
+    if not name or not gallery or name not in gallery:
+        return None, None, None, None, None, None
+    entry = gallery[name]
+    base = GALLERY_DIR / "images"
+    he = Image.open(base / entry["he"]).convert("RGB") if "he" in entry else None
+    gt = Image.open(base / entry["gt"]).convert("RGB") if "gt" in entry else None
+    gen_her2 = Image.open(base / entry["gen_her2"]).convert("RGB") if "gen_her2" in entry else None
+    gen_ki67 = Image.open(base / entry["gen_ki67"]).convert("RGB") if "gen_ki67" in entry else None
+    gen_er = Image.open(base / entry["gen_er"]).convert("RGB") if "gen_er" in entry else None
+    gen_pr = Image.open(base / entry["gen_pr"]).convert("RGB") if "gen_pr" in entry else None
+    return he, gt, gen_her2, gen_ki67, gen_er, gen_pr
+
+
+# ── Build Gradio App ─────────────────────────────────────────────────
+
+gallery = load_gallery()
+gallery_names = list(gallery.keys()) if gallery else []
+
+with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
+    gr.Markdown(
+        """
+        # UNIStainNet: Foundation-Model-Guided Virtual Staining
+        **H&E → IHC (HER2, Ki67, ER, PR)** &nbsp;|&nbsp;
+        Single unified model &nbsp;|&nbsp; 42M parameters &nbsp;|&nbsp;
+        UNI spatial conditioning
+        """
+    )
+
+    # ── Tab 1: Single Stain ──────────────────────────────────────
+    with gr.Tab("Virtual Staining"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(type="pil", label="Upload H&E Image", height=400)
+                stain_choice = gr.Radio(
+                    choices=STAIN_NAMES, value="HER2", label="Target IHC Stain"
+                )
+                guidance_slider = gr.Slider(
+                    minimum=1.0, maximum=3.0, step=0.1, value=1.0,
+                    label="Guidance Scale (1.0 = no CFG)",
+                )
+                generate_btn = gr.Button("Generate", variant="primary")
+                gen_time = gr.Textbox(label="Time", interactive=False)
+            with gr.Column(scale=1):
+                output_image = gr.Image(type="pil", label="Generated IHC", height=400)
+
+        generate_btn.click(
+            fn=generate_single_stain,
+            inputs=[input_image, stain_choice, guidance_slider],
+            outputs=[output_image, gen_time],
+        )
+
+    # ── Tab 2: Cross-Stain ───────────────────────────────────────
+    with gr.Tab("Cross-Stain Comparison"):
+        gr.Markdown(
+            "Generate **all 4 IHC stains** from a single H&E input. "
+            "Demonstrates the unified multi-stain capability."
+        )
+        with gr.Row():
+            cross_input = gr.Image(type="pil", label="Upload H&E Image", height=350)
+            cross_guidance = gr.Slider(
+                minimum=1.0, maximum=3.0, step=0.1, value=1.0,
+                label="Guidance Scale",
+            )
+        cross_btn = gr.Button("Generate All Stains", variant="primary")
+        cross_time = gr.Textbox(label="Time", interactive=False)
+
+        with gr.Row():
+            cross_he_out = gr.Image(type="pil", label="H&E Input", height=300)
+            cross_her2 = gr.Image(type="pil", label="HER2", height=300)
+            cross_ki67 = gr.Image(type="pil", label="Ki67", height=300)
+            cross_er = gr.Image(type="pil", label="ER", height=300)
+            cross_pr = gr.Image(type="pil", label="PR", height=300)
+
+        cross_btn.click(
+            fn=generate_all_stains,
+            inputs=[cross_input, cross_guidance],
+            outputs=[cross_he_out, cross_her2, cross_ki67, cross_er, cross_pr, cross_time],
+        )
+
+    # ── Tab 3: Gallery ───────────────────────────────────────────
+    with gr.Tab("Gallery"):
+        if not gallery_names:
+            gr.Markdown("No pre-computed gallery available.")
+        else:
+            gr.Markdown(
+                "Pre-computed examples — no GPU required. "
+                "Select an example to view the H&E input and generated IHC stains."
+            )
+            gallery_dropdown = gr.Dropdown(
+                choices=gallery_names,
+                value=gallery_names[0] if gallery_names else None,
+                label="Select Example",
+            )
+            with gr.Row():
+                gal_he = gr.Image(type="pil", label="H&E Input", height=300)
+                gal_gt = gr.Image(type="pil", label="Ground Truth IHC", height=300)
+            with gr.Row():
+                gal_her2 = gr.Image(type="pil", label="Generated HER2", height=300)
+                gal_ki67 = gr.Image(type="pil", label="Generated Ki67", height=300)
+                gal_er = gr.Image(type="pil", label="Generated ER", height=300)
+                gal_pr = gr.Image(type="pil", label="Generated PR", height=300)
+
+            gallery_dropdown.change(
+                fn=lambda name: show_gallery(name, gallery),
+                inputs=[gallery_dropdown],
+                outputs=[gal_he, gal_gt, gal_her2, gal_ki67, gal_er, gal_pr],
+            )
+
+    # ── Tab 4: About ─────────────────────────────────────────────
+    with gr.Tab("About"):
+        gr.Markdown(
+            """
+            ## UNIStainNet
+
+            A SPADE-UNet generator conditioned on dense spatial tokens from a frozen
+            [UNI](https://github.com/mahmoodlab/UNI) pathology foundation model (ViT-L/16).
+
+            **Key features:**
+            - Dense UNI spatial conditioning (32x32 = 1,024 tokens)
+            - Misalignment-aware loss suite for consecutive-section training pairs
+            - Single unified model serves 4 IHC markers (HER2, Ki67, ER, PR)
+            - 42M generator parameters, single forward pass inference
+
+            ### Architecture
+
+            | Component | Details |
+            |-----------|---------|
+            | Generator | SPADE-UNet with UNI spatial conditioning + FiLM stain embeddings |
+            | Discriminator | Multi-scale PatchGAN (512 + 256) with spectral norm |
+            | UNI Features | 4x4 sub-crop tiling → UNI ViT-L/16 → 32x32 spatial tokens |
+            | Parameters | 42M (generator), UNI frozen (303M) |
+
+            ### Results (MIST, unified model)
+
+            | Stain | FID ↓ | KID×1k ↓ | Pearson-R ↑ | DAB KL ↓ |
+            |-------|-------|-----------|-------------|----------|
+            | HER2  | 34.5  | 2.2       | 0.929       | 0.166    |
+            | Ki67  | 27.2  | 1.8       | 0.927       | 0.119    |
+            | ER    | 29.2  | 1.8       | 0.949       | 0.182    |
+            | PR    | 29.0  | 1.1       | 0.943       | 0.171    |
+            """
+        )
+
+if __name__ == "__main__":
+    demo.launch(theme=gr.themes.Soft())