Upload app.py with huggingface_hub

app.py

@@ -1,9 +1,9 @@
 #!/usr/bin/env python3
 """
-UNIStainNet Interactive Demo — Hugging Face Spaces (ZeroGPU)
+UNIStainNet Interactive Demo — Hugging Face Spaces
 
 Virtual staining of H&E histopathology images to IHC (HER2, Ki67, ER, PR).
-Uses @spaces.GPU for on-demand GPU allocation on ZeroGPU.
+Supports ZeroGPU (HF Pro) for live inference, falls back to gallery-only on CPU.
 """
 
 import json
@@ -14,89 +14,82 @@ from pathlib import Path
 import gradio as gr
 import numpy as np
 import torch
+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
 
-# ZeroGPU support: use @spaces.GPU if available, otherwise no-op
+# ── ZeroGPU support ──────────────────────────────────────────────────
 try:
     import spaces
-    GPU_AVAILABLE = torch.cuda.is_available()
+    HAS_SPACES = True
 except ImportError:
     spaces = None
-    GPU_AVAILABLE = torch.cuda.is_available()
+    HAS_SPACES = False
 
+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"):
+    if HAS_SPACES 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) ────────────────
+NO_GPU_MSG = (
+    "GPU is not available on this Space (requires HF Pro for ZeroGPU). "
+    "Please use the **Gallery** tab to browse pre-computed results, "
+    "or run the app locally with a GPU: `python app.py`"
+)
+
+# ── Lazy imports (avoid crash if no GPU) ─────────────────────────────
 _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
+    """Load models onto GPU. Only called when GPU is confirmed available."""
+    from src.models.trainer import UNIStainNetTrainer
+    import timm
 
+    if _model_cache["model"] is None:
         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")
+        _model_cache["spatial_pool_size"] = getattr(model.hparams, "uni_spatial_size", 32)
 
         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,
+            "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")
+        print("  Models 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 ────────────────────────────────────────────
+# ── Preprocessing ────────────────────────────────────────────────────
 
 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
@@ -108,52 +101,40 @@ def preprocess_he(pil_image, target_size=TARGET_SIZE):
 
 
 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]
-    )
+    from src.data.mist_dataset import STAIN_TO_LABEL
+    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
+    num_crops, patches_per_side = 4, 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 = 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
-    )
+    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
-    )
+    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:
@@ -162,74 +143,82 @@ def extract_uni_features(uni_model, he_tensor_01, spatial_pool_size=32):
         result = pooled.permute(0, 2, 3, 1)
     else:
         result = full_grid
-
     return result.reshape(B, S * S, 1024)
 
 
-# ── GPU-accelerated inference functions ──────────────────────────────
+# ── 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()
+def _generate_single_gpu(image, stain, guidance_scale):
+    """GPU inference for single stain."""
+    from src.data.mist_dataset import STAIN_TO_LABEL
     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)
+    labels = torch.tensor([STAIN_TO_LABEL[stain]], 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"
+    return tensor_to_pil(gen)
 
 
 @_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()
+def _generate_all_gpu(image, guidance_scale):
+    """GPU inference for all 4 stains."""
+    from src.data.mist_dataset import STAIN_TO_LABEL
     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)
+        labels = torch.tensor([STAIN_TO_LABEL[stain]], device="cuda", dtype=torch.long)
         with torch.no_grad():
-            gen = model.generate(
-                he_tensor, uni_feats, labels, guidance_scale=guidance_scale
-            )
+            gen = model.generate(he_tensor, uni_feats, labels, guidance_scale=guidance_scale)
         results[stain] = tensor_to_pil(gen)
+    return he_pil, results
 
-    elapsed = time.time() - t0
-    return (
-        he_pil,
-        results["HER2"],
-        results["Ki67"],
-        results["ER"],
-        results["PR"],
-        f"{elapsed:.2f}s",
-    )
+
+def generate_single_stain(image, stain, guidance_scale):
+    """Wrapper with GPU availability check."""
+    if image is None:
+        return None, "No image uploaded"
+    if not GPU_AVAILABLE and not HAS_SPACES:
+        return None, NO_GPU_MSG
+    try:
+        t0 = time.time()
+        result = _generate_single_gpu(image, stain, guidance_scale)
+        return result, f"{time.time() - t0:.2f}s"
+    except RuntimeError as e:
+        if "NVIDIA" in str(e) or "CUDA" in str(e) or "cuda" in str(e):
+            return None, NO_GPU_MSG
+        raise
+
+
+def generate_all_stains(image, guidance_scale):
+    """Wrapper with GPU availability check."""
+    if image is None:
+        return None, None, None, None, None, "No image uploaded"
+    if not GPU_AVAILABLE and not HAS_SPACES:
+        return None, None, None, None, None, NO_GPU_MSG
+    try:
+        t0 = time.time()
+        he_pil, results = _generate_all_gpu(image, guidance_scale)
+        elapsed = f"{time.time() - t0:.2f}s"
+        return he_pil, results["HER2"], results["Ki67"], results["ER"], results["PR"], elapsed
+    except RuntimeError as e:
+        if "NVIDIA" in str(e) or "CUDA" in str(e) or "cuda" in str(e):
+            return None, None, None, None, None, NO_GPU_MSG
+        raise
 
 
-# ── Gallery helpers ──────────────────────────────────────────────────
+# ── Gallery ──────────────────────────────────────────────────────────
 
 def load_gallery():
     meta_path = GALLERY_DIR / "metadata.json"
@@ -270,18 +259,20 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
 
     # ── Tab 1: Single Stain ──────────────────────────────────────
     with gr.Tab("Virtual Staining"):
+        if not GPU_AVAILABLE and not HAS_SPACES:
+            gr.Markdown(
+                f"> **Note:** {NO_GPU_MSG}"
+            )
         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"
-                )
+                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)
+                gen_time = gr.Textbox(label="Status", interactive=False)
             with gr.Column(scale=1):
                 output_image = gr.Image(type="pil", label="Generated IHC", height=400)
 
@@ -293,6 +284,10 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
 
     # ── Tab 2: Cross-Stain ───────────────────────────────────────
     with gr.Tab("Cross-Stain Comparison"):
+        if not GPU_AVAILABLE and not HAS_SPACES:
+            gr.Markdown(
+                f"> **Note:** {NO_GPU_MSG}"
+            )
         gr.Markdown(
             "Generate **all 4 IHC stains** from a single H&E input. "
             "Demonstrates the unified multi-stain capability."
@@ -300,11 +295,10 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
         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",
+                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)
+        cross_time = gr.Textbox(label="Status", interactive=False)
 
         with gr.Row():
             cross_he_out = gr.Image(type="pil", label="H&E Input", height=300)
@@ -326,7 +320,7 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
         else:
             gr.Markdown(
                 "Pre-computed examples — no GPU required. "
-                "Select an example to view the H&E input and generated IHC stains."
+                "Select an example to view the H&E input, ground truth, and generated IHC stains."
             )
             gallery_dropdown = gr.Dropdown(
                 choices=gallery_names,