ZeroGPU-adapted DynaCell demo, repointed to biohub repos

README.md

@@ -1,13 +1,41 @@
 ---
-title: Dynacell
-emoji: 💻
-colorFrom: yellow
-colorTo: blue
+title: DynaCell Virtual Staining Demo
+emoji: 🔬
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.16.0
-python_version: '3.12'
+sdk_version: "5.29.0"
 app_file: app.py
 pinned: false
+suggested_hardware: zero-a10g
+models:
+  - biohub/dynacell-checkpoints
+datasets:
+  - biohub/dynacell-demo-data
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# DynaCell Virtual Staining Demo
+
+Predict fluorescence channels (membrane, nuclei, or organelle structure) from phase-contrast OME-Zarr using three models:
+
+- **CELL-Diff** — flow-matching diffusion model
+- **FNet3D** — 3-D U-Net (FNet architecture)
+- **VSCyto3D** — masked-autoencoder pretrained U-Net
+
+## Quick start
+
+1. Select an organelle from the dropdown.
+2. Click **Load Demo Data** to fetch the matching A549-cell demo dataset directly into the Space — no download/upload needed.
+3. Run predictions in **Tab 1** or generate the CELL-Diff ODE trajectory in **Tab 2**.
+
+## Using your own data
+
+The input must be an OME-Zarr HCS store zipped into a single `.zip` file, with layout:
+
+```
+your_data.zarr/
+  0/0/fov0000/0      # array shape (T, C, Z, Y, X)
+                     # C[0] = Phase3D, Z = 16, YX = 512×512
+```
+
+Use [iohub](https://github.com/czbiohub-sf/iohub) to create compatible zarr stores.

app.py

@@ -0,0 +1,522 @@
+"""DynaCell Virtual Staining Demo — Gradio Space.
+
+Upload a zipped OME-Zarr HCS store once; then:
+  Tab 1 — run CELL-Diff / FNet3D / VSCyto3D predictions on a selected timepoint,
+           view a chosen Z slice, and see Spectral PCC metrics.
+  Tab 2 — visualize the CELL-Diff ODE denoising trajectory as an animated GIF,
+           with a Phase | Exp reference panel at the selected timepoint and Z slice.
+           Changing the Z-slice slider re-renders the GIF instantly from cached data.
+"""
+
+from __future__ import annotations
+
+import sys
+import tempfile
+import zipfile
+from pathlib import Path
+
+import gradio as gr
+import matplotlib.pyplot as plt
+import numpy as np
+from iohub.ngff import open_ome_zarr
+
+sys.path.insert(0, str(Path(__file__).parent))
+from predict_runner import (
+    ORGANELLE_LABELS, TARGET_CHANNELS,
+    preprocess_zarr, run_prediction,
+    compute_trajectory, render_trajectory_gif,
+)
+
+from cubic.metrics.bandlimited import spectral_pcc
+
+ORGANELLES   = ["CAAX", "H2B", "SEC61B", "TOMM20"]
+MODEL_KEYS   = ["celldiff", "fnet3d", "vscyto3d"]
+MODEL_LABELS = {"celldiff": "CELL-Diff", "fnet3d": "FNet3D", "vscyto3d": "VSCyto3D"}
+PHASE_CH = 0
+FLUOR_CH = 2
+_DEMO_REPO = "biohub/dynacell-demo-data"
+
+PATCH_D = 8  # fixed Z window used by all trajectory models
+
+SPACING = [0.174, 0.1494, 0.1494]
+SPECTRAL_KWARGS = dict(bin_delta=1.0, tail_fraction=0.2, apodization="tukey", nbins_low=3)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def extract_zarr_zip(zip_path: str) -> str:
+    """Extract uploaded zip to a fresh temp dir; return the HCS zarr root path."""
+    import json
+    tmpdir = Path(tempfile.mkdtemp())
+    with zipfile.ZipFile(zip_path, "r") as z:
+        z.extractall(tmpdir)
+    for candidate in sorted(tmpdir.rglob(".zattrs")):
+        root = candidate.parent
+        try:
+            zattrs = json.loads((root / ".zattrs").read_text())
+            if "plate" in zattrs:
+                return str(root)
+        except Exception:
+            pass
+    for d in sorted(tmpdir.iterdir()):
+        if d.is_dir():
+            return str(d)
+    raise ValueError("No zarr store found in zip.")
+
+
+def get_data_shape(data_path: str) -> tuple[int, int]:
+    """Return (n_timepoints, n_z_slices) from the first position in the plate."""
+    with open_ome_zarr(data_path, mode="r") as plate:
+        _, pos = next(plate.positions())
+        return pos.data.shape[0], pos.data.shape[2]
+
+
+def percentile_norm(img: np.ndarray, lo: float = 0.5, hi: float = 99.5) -> np.ndarray:
+    p_lo, p_hi = np.percentile(img, [lo, hi])
+    if p_hi == p_lo:
+        return np.zeros_like(img, dtype=np.float32)
+    return np.clip((img - p_lo) / (p_hi - p_lo), 0, 1).astype(np.float32)
+
+
+def compute_spectral_pcc(pred_zarr_path: str, gt_fluor_vol: np.ndarray) -> float | None:
+    """Spectral PCC between the prediction (t=0) and the GT fluorescence volume."""
+    try:
+        with open_ome_zarr(pred_zarr_path, mode="r") as pred_plate:
+            _, pred_pos = next(pred_plate.positions())
+            pred_vol = np.array(pred_pos.data[0, 0], dtype=np.float32)
+        return float(spectral_pcc(pred_vol, gt_fluor_vol, spacing=SPACING, **SPECTRAL_KWARGS))
+    except Exception as e:
+        print(f"spectral_pcc failed: {e}")
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Data loaders
+# ---------------------------------------------------------------------------
+
+def _make_slider_updates(data_path: str, organelle: str) -> tuple:
+    """Read data shape and return slider updates + Phase|Exp figure."""
+    n_tp, n_z = get_data_shape(data_path)
+    z_mid = n_z // 2
+    fig = render_phase_exp_traj(data_path, 0, PATCH_D // 2, organelle)
+    return (
+        gr.Slider(minimum=0, maximum=n_tp - 1, step=1, value=0),    # timepoint_slider
+        gr.Slider(minimum=0, maximum=n_z - 1, step=1, value=z_mid), # z_slice_slider
+        gr.Slider(minimum=0, maximum=n_tp - 1, step=1, value=0),    # traj_timepoint
+        fig,                                                           # traj_static
+        n_tp, n_z,
+    )
+
+
+def load_demo_data(organelle: str, progress=gr.Progress()) -> tuple:
+    """Download the demo zarr, extract it, and return updated UI state."""
+    from huggingface_hub import hf_hub_download
+    filename = f"{organelle}_mock.zarr.zip"
+    progress(0.1, desc=f"Downloading {organelle} demo data...")
+    zip_path = hf_hub_download(repo_id=_DEMO_REPO, filename=filename, repo_type="dataset")
+    progress(0.8, desc="Extracting zarr...")
+    data_path = extract_zarr_zip(zip_path)
+    tp_sl, z_sl, traj_tp, fig, n_tp, n_z = _make_slider_updates(data_path, organelle)
+    progress(1.0, desc="Ready.")
+    status = f"**Loaded:** {filename} (A549 cells, {n_tp} timepoints, {n_z} Z slices)"
+    return data_path, status, tp_sl, z_sl, traj_tp, fig
+
+
+def on_upload(file, organelle: str) -> tuple:
+    """Handle zarr zip upload: extract, read shape, update UI state."""
+    if file is None:
+        raise gr.Error("No file uploaded.")
+    zip_path = file if isinstance(file, str) else file.name
+    data_path = extract_zarr_zip(zip_path)
+    tp_sl, z_sl, traj_tp, fig, n_tp, n_z = _make_slider_updates(data_path, organelle)
+    status = f"**Uploaded:** {Path(zip_path).name} ({n_tp} timepoints, {n_z} Z slices)"
+    return data_path, status, tp_sl, z_sl, traj_tp, fig
+
+
+# ---------------------------------------------------------------------------
+# Tab 2: Phase | Exp reference panel
+# ---------------------------------------------------------------------------
+
+def render_phase_exp(
+    zarr_state: str | None,
+    timepoint: int,
+    z_slice: int,
+    organelle: str,
+) -> plt.Figure | None:
+    """Render Phase and Experimental fluorescence side by side at (timepoint, z_slice)."""
+    if zarr_state is None:
+        return None
+    with open_ome_zarr(zarr_state, mode="r") as plate:
+        _, pos = next(plate.positions())
+        n_tp = pos.data.shape[0]
+        n_z  = pos.data.shape[2]
+        tp = min(timepoint, n_tp - 1)
+        z  = min(z_slice,   n_z  - 1)
+        phase_img = np.array(pos.data[tp, PHASE_CH, z])
+        fluor_img = np.array(pos.data[tp, FLUOR_CH, z])
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6.4, 3.2))
+    ax1.imshow(percentile_norm(phase_img), cmap="gray")
+    ax1.set_title("Phase", fontsize=10)
+    ax1.axis("off")
+    ax2.imshow(percentile_norm(fluor_img), cmap="gray")
+    ax2.set_title(f"Exp ({TARGET_CHANNELS[organelle]})", fontsize=10)
+    ax2.axis("off")
+    fig.suptitle(
+        f"{ORGANELLE_LABELS[organelle]}  |  t={tp}  |  z={z}",
+        fontsize=11, y=1.01,
+    )
+    fig.tight_layout()
+    return fig
+
+
+def render_phase_exp_traj(
+    zarr_state: str | None,
+    timepoint: int,
+    z_patch: int,
+    organelle: str,
+) -> plt.Figure | None:
+    """Phase | Exp panel for the trajectory tab.
+
+    z_patch is a patch-relative index (0 … PATCH_D-1); converted to the
+    absolute Z using z_start = (n_z - PATCH_D) // 2.
+    """
+    if zarr_state is None:
+        return None
+    with open_ome_zarr(zarr_state, mode="r") as plate:
+        _, pos = next(plate.positions())
+        n_tp = pos.data.shape[0]
+        n_z  = pos.data.shape[2]
+        tp      = min(timepoint, n_tp - 1)
+        z_start = (n_z - PATCH_D) // 2
+        z_abs   = z_start + max(0, min(z_patch, PATCH_D - 1))
+        phase_img = np.array(pos.data[tp, PHASE_CH, z_abs])
+        fluor_img = np.array(pos.data[tp, FLUOR_CH, z_abs])
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6.4, 3.2))
+    ax1.imshow(percentile_norm(phase_img), cmap="gray")
+    ax1.set_title("Phase", fontsize=10)
+    ax1.axis("off")
+    ax2.imshow(percentile_norm(fluor_img), cmap="gray")
+    ax2.set_title(f"Exp ({TARGET_CHANNELS[organelle]})", fontsize=10)
+    ax2.axis("off")
+    fig.suptitle(
+        f"{ORGANELLE_LABELS[organelle]}  |  t={tp}  |  z={z_abs} (patch slice {z_patch})",
+        fontsize=11, y=1.01,
+    )
+    fig.tight_layout()
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# Tab 1: Virtual Staining
+# ---------------------------------------------------------------------------
+
+def render_from_z(
+    pred_info: dict | None,
+    z_slice: int,
+    zarr_state: str | None,
+) -> plt.Figure | None:
+    """Re-render the prediction comparison at a different Z slice."""
+    if pred_info is None or zarr_state is None:
+        return None
+
+    organelle       = pred_info["organelle"]
+    timepoint       = pred_info["timepoint"]
+    selected_models = pred_info["selected_models"]
+    pred_paths      = pred_info["paths"]
+    pred_pccs       = pred_info["pccs"]
+    n_z             = pred_info["n_z"]
+    z = min(z_slice, n_z - 1)
+
+    with open_ome_zarr(zarr_state, mode="r") as gt_plate:
+        _, gt_pos = next(gt_plate.positions())
+        phase_img = np.array(gt_pos.data[timepoint, PHASE_CH, z])
+        fluor_img = np.array(gt_pos.data[timepoint, FLUOR_CH, z])
+
+    cols = ["Phase", f"Exp ({TARGET_CHANNELS[organelle]})"] + [MODEL_LABELS[m] for m in selected_models]
+    fig, axes = plt.subplots(1, len(cols), figsize=(3.0 * len(cols), 3.2))
+    if len(cols) == 1:
+        axes = [axes]
+
+    axes[0].imshow(percentile_norm(phase_img), cmap="gray")
+    axes[0].set_title("Phase", fontsize=10)
+    axes[1].imshow(percentile_norm(fluor_img), cmap="gray")
+    axes[1].set_title(f"Exp ({TARGET_CHANNELS[organelle]})", fontsize=10)
+
+    for col_idx, model_key in enumerate(selected_models, start=2):
+        label     = MODEL_LABELS[model_key]
+        pred_path = pred_paths.get(model_key)
+        pcc       = pred_pccs.get(model_key)
+        if pred_path is not None:
+            try:
+                with open_ome_zarr(pred_path, mode="r") as pred_plate:
+                    _, pred_pos = next(pred_plate.positions())
+                    img = percentile_norm(np.array(pred_pos.data[0, 0, z]))
+                title = f"{label}\nSpectral PCC={pcc:.3f}" if pcc is not None else label
+            except Exception as e:
+                img   = np.zeros_like(phase_img, dtype=np.float32)
+                title = f"{label}\n(failed)"
+                print(f"Render failed for {model_key}: {e}")
+        else:
+            img   = np.zeros_like(phase_img, dtype=np.float32)
+            title = f"{label}\n(failed)"
+
+        axes[col_idx].imshow(img, cmap="gray")
+        axes[col_idx].set_title(title, fontsize=9)
+
+    for ax in axes:
+        ax.axis("off")
+    fig.suptitle(
+        f"{ORGANELLE_LABELS[organelle]}  |  t={timepoint}  |  z={z}",
+        fontsize=11, y=1.01,
+    )
+    fig.tight_layout()
+    return fig
+
+
+def run_demo(
+    zarr_zip,
+    organelle: str,
+    selected_models: list[str],
+    timepoint: int,
+    z_slice: int,
+    zarr_state: str | None,
+    progress=gr.Progress(),
+) -> tuple[plt.Figure | None, list[list], str, dict]:
+    if zarr_zip is None and not zarr_state:
+        raise gr.Error("Please load demo data or upload a zarr zip file.")
+    if not selected_models:
+        raise gr.Error("Select at least one model.")
+
+    if zarr_state:
+        data_path = zarr_state
+    else:
+        progress(0.05, desc="Extracting zarr...")
+        zip_path = zarr_zip if isinstance(zarr_zip, str) else zarr_zip.name
+        data_path = extract_zarr_zip(zip_path)
+
+    progress(0.10, desc="Computing normalization statistics...")
+    preprocess_zarr(data_path)
+
+    with open_ome_zarr(data_path, mode="r") as gt_plate:
+        _, gt_pos = next(gt_plate.positions())
+        n_tp, n_z = gt_pos.data.shape[0], gt_pos.data.shape[2]
+        tp = min(timepoint, n_tp - 1)
+        gt_fluor_vol = np.array(gt_pos.data[tp, FLUOR_CH], dtype=np.float32)
+
+    pred_paths: dict[str, str | None] = {}
+    pred_pccs:  dict[str, float | None] = {}
+    n_models = len(selected_models)
+    for i, model_key in enumerate(selected_models):
+        progress(0.15 + 0.60 * i / n_models, desc=f"Running {MODEL_LABELS[model_key]}...")
+        try:
+            path = run_prediction(model_key, organelle, data_path, tp)
+            pred_paths[model_key] = path
+            pred_pccs[model_key]  = compute_spectral_pcc(path, gt_fluor_vol)
+        except Exception as e:
+            pred_paths[model_key] = None
+            pred_pccs[model_key]  = None
+            print(f"Prediction failed for {model_key}: {e}")
+
+    pred_info = {
+        "timepoint": tp, "organelle": organelle,
+        "selected_models": selected_models,
+        "paths": pred_paths, "pccs": pred_pccs, "n_z": n_z,
+    }
+
+    progress(0.80, desc="Rendering figure...")
+    fig = render_from_z(pred_info, min(z_slice, n_z - 1), data_path)
+
+    metrics_rows = [
+        [MODEL_LABELS[m], "failed" if pred_paths.get(m) is None
+         else (f"{pred_pccs[m]:.4f}" if pred_pccs.get(m) is not None else "N/A")]
+        for m in selected_models
+    ]
+
+    progress(1.0, desc="Done.")
+    return fig, metrics_rows, data_path, pred_info
+
+
+# ---------------------------------------------------------------------------
+# Tab 2: CellDiff Trajectory
+# ---------------------------------------------------------------------------
+
+def run_trajectory_demo(
+    zarr_zip,
+    organelle: str,
+    timepoint: int,
+    num_steps: int,
+    z_slice: int,
+    zarr_state: str | None,
+    progress=gr.Progress(),
+) -> tuple[str, str, dict]:
+    """Run ODE trajectory, render GIF, cache trajectory data for Z-slice re-renders."""
+    if zarr_zip is None and not zarr_state:
+        raise gr.Error("Please load demo data or upload a zarr zip file.")
+
+    if zarr_state:
+        data_path = zarr_state
+    else:
+        progress(0.03, desc="Extracting zarr...")
+        zip_path = zarr_zip if isinstance(zarr_zip, str) else zarr_zip.name
+        data_path = extract_zarr_zip(zip_path)
+
+    progress(0.08, desc="Computing normalization statistics...")
+    preprocess_zarr(data_path)
+
+    traj_info = compute_trajectory(organelle, data_path, timepoint, num_steps, progress)
+    gif_path  = render_trajectory_gif(traj_info, z_slice)
+    return gif_path, data_path, traj_info
+
+
+def rerender_gif(traj_info: dict | None, z_slice: int) -> str | None:
+    """Re-render the trajectory GIF at a new Z slice without re-running the ODE."""
+    if traj_info is None:
+        return None
+    return render_trajectory_gif(traj_info, z_slice)
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+with gr.Blocks(title="DynaCell Virtual Staining") as demo:
+    gr.Markdown("## DynaCell Virtual Staining Demo")
+    gr.Markdown(
+        "**Tab 1** runs virtual staining predictions (CELL-Diff / FNet3D / VSCyto3D) "
+        "on a phase-contrast OME-Zarr for a selected timepoint, and reports Spectral PCC. "
+        "**Tab 2** visualizes the CELL-Diff ODE denoising trajectory."
+    )
+
+    zarr_state      = gr.State(value=None)
+    pred_info_state = gr.State(value=None)
+    traj_info_state = gr.State(value=None)
+
+    # ---- Data source row -------------------------------------------------
+    with gr.Row():
+        organelle = gr.Dropdown(
+            choices=[(ORGANELLE_LABELS[o], o) for o in ORGANELLES],
+            value="CAAX", label="Organelle",
+            info="Select the target organelle.",
+        )
+        load_demo_btn = gr.Button("Load Demo Data", variant="secondary", scale=1)
+        zarr_upload = gr.File(
+            label="Or upload your own zarr (.zip)",
+            file_types=[".zip"],
+            scale=2,
+        )
+
+    data_status = gr.Markdown("")
+
+    # ---- Tabs ------------------------------------------------------------
+    with gr.Tabs():
+
+        with gr.Tab("Virtual Staining"):
+            with gr.Row():
+                model_selector = gr.CheckboxGroup(
+                    choices=[(MODEL_LABELS[m], m) for m in MODEL_KEYS],
+                    value=MODEL_KEYS,
+                    label="Models to run",
+                )
+            with gr.Row():
+                timepoint_slider = gr.Slider(
+                    minimum=0, maximum=4, step=1, value=0,
+                    label="Timepoint",
+                    info="Range updates after loading data.",
+                )
+                z_slice_slider = gr.Slider(
+                    minimum=0, maximum=99, step=1, value=15,
+                    label="Z slice",
+                    info="Range updates after loading data.",
+                )
+            run_btn = gr.Button("Run Predictions", variant="primary")
+            output_plot = gr.Plot(label="Predictions")
+            metrics_table = gr.Dataframe(
+                headers=["Model", "Spectral PCC"],
+                label="Spectral PCC (volumetric, vs experimental fluorescence)",
+            )
+
+        with gr.Tab("CELL-Diff Trajectory"):
+            gr.Markdown(
+                "Generate the CELL-Diff ODE denoising trajectory. "
+                "T=0 is pure Gaussian noise; T=N is the final predicted fluorescence. "
+                "After generating, change **Z slice** to instantly re-render the GIF "
+                "at a different slice without re-running the ODE."
+            )
+            with gr.Row():
+                traj_timepoint = gr.Slider(
+                    minimum=0, maximum=4, step=1, value=0,
+                    label="Timepoint",
+                    info="Range updates after loading data.",
+                )
+                traj_z_slice = gr.Slider(
+                    minimum=0, maximum=PATCH_D - 1, step=1, value=PATCH_D // 2,
+                    label=f"Z slice (0–{PATCH_D - 1}, middle {PATCH_D} of full volume)",
+                )
+                traj_num_steps = gr.Slider(
+                    minimum=10, maximum=100, step=10, value=50,
+                    label="ODE steps",
+                )
+            traj_static = gr.Plot(label="Phase | Exp (reference)")
+            traj_btn    = gr.Button("Generate Trajectory", variant="primary")
+            traj_gif    = gr.Image(label="Animated trajectory (GIF)", type="filepath")
+
+    # ---- Event wiring ----------------------------------------------------
+
+    _data_outputs = [
+        zarr_state, data_status,
+        timepoint_slider, z_slice_slider,
+        traj_timepoint,
+        traj_static,
+    ]
+
+    load_demo_btn.click(
+        fn=load_demo_data,
+        inputs=[organelle],
+        outputs=_data_outputs,
+    )
+
+    zarr_upload.upload(
+        fn=on_upload,
+        inputs=[zarr_upload, organelle],
+        outputs=_data_outputs,
+    )
+
+    run_btn.click(
+        fn=run_demo,
+        inputs=[zarr_upload, organelle, model_selector, timepoint_slider, z_slice_slider, zarr_state],
+        outputs=[output_plot, metrics_table, zarr_state, pred_info_state],
+    )
+
+    z_slice_slider.change(
+        fn=render_from_z,
+        inputs=[pred_info_state, z_slice_slider, zarr_state],
+        outputs=[output_plot],
+    )
+
+    # Phase | Exp panel updates on any slider or organelle change
+    for _trigger in (traj_timepoint, traj_z_slice, organelle):
+        _trigger.change(
+            fn=render_phase_exp_traj,
+            inputs=[zarr_state, traj_timepoint, traj_z_slice, organelle],
+            outputs=[traj_static],
+        )
+
+    traj_btn.click(
+        fn=run_trajectory_demo,
+        inputs=[zarr_upload, organelle, traj_timepoint, traj_num_steps, traj_z_slice, zarr_state],
+        outputs=[traj_gif, zarr_state, traj_info_state],
+    )
+
+    # Re-render GIF from cached trajectory when Z slice changes (no ODE re-run)
+    traj_z_slice.change(
+        fn=rerender_gif,
+        inputs=[traj_info_state, traj_z_slice],
+        outputs=[traj_gif],
+    )
+
+if __name__ == "__main__":
+    demo.launch()

config_templates/celldiff.yaml

@@ -0,0 +1,50 @@
+model:
+  class_path: dynacell.engine.DynacellFlowMatching
+  init_args:
+    net_config:
+      input_spatial_size: [8, 512, 512]
+      in_channels: 1
+      dims: [64, 128, 256, 256]
+      num_res_block: [2, 2, 2]
+      hidden_size: 512
+      num_heads: 8
+      dim_head: 64
+      num_hidden_layers: 8
+      patch_size: 4
+    transport_config:
+      path_type: Linear
+      prediction: velocity
+    num_generate_steps: 100
+    predict_method: iterative
+    predict_overlap: [4, 256, 256]
+    ckpt_path: {ckpt_path}
+
+data:
+  class_path: viscy_data.hcs.HCSDataModule
+  init_args:
+    data_path: {data_path}
+    source_channel: Phase3D
+    target_channel: {target_channel}
+    z_window_size: 16
+    batch_size: 1
+    num_workers: 0
+    yx_patch_size: [512, 512]
+    normalizations:
+      - class_path: viscy_transforms.MinMaxSampled
+        init_args:
+          keys: [Phase3D]
+          level: timepoint_statistics
+    augmentations: []
+
+trainer:
+  accelerator: gpu
+  strategy: auto
+  devices: 1
+  num_nodes: 1
+  precision: 32-true
+  callbacks:
+    - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter
+      init_args:
+        output_store: {output_store}
+
+return_predictions: false

config_templates/fnet3d.yaml

@@ -0,0 +1,45 @@
+model:
+  class_path: dynacell.engine.DynacellUNet
+  init_args:
+    architecture: FNet3D
+    model_config:
+      in_channels: 1
+      out_channels: 1
+      depth: 4
+      mult_chan: 32
+      in_stack_depth: 16
+    predict_method: full_image
+    predict_overlap: [4, 256, 256]
+    ckpt_path: {ckpt_path}
+
+data:
+  class_path: viscy_data.hcs.HCSDataModule
+  init_args:
+    data_path: {data_path}
+    source_channel: Phase3D
+    target_channel: {target_channel}
+    z_window_size: 16
+    batch_size: 1
+    num_workers: 0
+    yx_patch_size: [512, 512]
+    normalizations:
+      - class_path: viscy_transforms.NormalizeSampled
+        init_args:
+          keys: [Phase3D]
+          level: fov_statistics
+          subtrahend: mean
+          divisor: std
+    augmentations: []
+
+trainer:
+  accelerator: gpu
+  strategy: auto
+  devices: 1
+  num_nodes: 1
+  precision: 32-true
+  callbacks:
+    - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter
+      init_args:
+        output_store: {output_store}
+
+return_predictions: false

config_templates/vscyto3d.yaml

@@ -0,0 +1,49 @@
+model:
+  class_path: dynacell.engine.DynacellUNet
+  init_args:
+    architecture: fcmae
+    model_config:
+      in_channels: 1
+      out_channels: 1
+      in_stack_depth: 15
+      decoder_conv_blocks: 2
+      dims: [96, 192, 384, 768]
+      encoder_blocks: [3, 3, 9, 3]
+      encoder_drop_path_rate: 0.1
+      pretraining: false
+      stem_kernel_size: [5, 4, 4]
+    predict_method: full_image
+    predict_overlap: [4, 256, 256]
+    ckpt_path: {ckpt_path}
+
+data:
+  class_path: viscy_data.hcs.HCSDataModule
+  init_args:
+    data_path: {data_path}
+    source_channel: Phase3D
+    target_channel: {target_channel}
+    z_window_size: 15
+    batch_size: 1
+    num_workers: 0
+    yx_patch_size: [512, 512]
+    normalizations:
+      - class_path: viscy_transforms.NormalizeSampled
+        init_args:
+          keys: [Phase3D]
+          level: fov_statistics
+          subtrahend: mean
+          divisor: std
+    augmentations: []
+
+trainer:
+  accelerator: gpu
+  strategy: auto
+  devices: 1
+  num_nodes: 1
+  precision: 32-true
+  callbacks:
+    - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter
+      init_args:
+        output_store: {output_store}
+
+return_predictions: false

predict_runner.py

@@ -0,0 +1,286 @@
+"""Download checkpoints from HF Hub, generate configs, run dynacell predict, and generate trajectories."""
+
+from __future__ import annotations
+
+import json
+import shutil
+import subprocess
+import tempfile
+import uuid
+from pathlib import Path
+
+import spaces
+import zarr
+from huggingface_hub import hf_hub_download
+
+CHECKPOINT_REPO = "biohub/dynacell-checkpoints"
+TEMPLATE_DIR = Path(__file__).parent / "config_templates"
+
+# (model, organelle) → filename in the HF checkpoint repo
+CHECKPOINT_FILES: dict[tuple[str, str], str] = {
+    ("celldiff",  "CAAX"):   "celldiff_caax.ckpt",
+    ("celldiff",  "H2B"):    "celldiff_h2b.ckpt",
+    ("celldiff",  "SEC61B"): "celldiff_sec61b.ckpt",
+    ("celldiff",  "TOMM20"): "celldiff_tomm20.ckpt",
+    ("fnet3d",    "CAAX"):   "fnet3d_caax.ckpt",
+    ("fnet3d",    "H2B"):    "fnet3d_h2b.ckpt",
+    ("fnet3d",    "SEC61B"): "fnet3d_sec61b.ckpt",
+    ("fnet3d",    "TOMM20"): "fnet3d_tomm20.ckpt",
+    ("vscyto3d",  "CAAX"):   "vscyto3d_caax.ckpt",
+    ("vscyto3d",  "H2B"):    "vscyto3d_h2b.ckpt",
+    ("vscyto3d",  "SEC61B"): "vscyto3d_sec61b.ckpt",
+    ("vscyto3d",  "TOMM20"): "vscyto3d_tomm20.ckpt",
+}
+
+TARGET_CHANNELS: dict[str, str] = {
+    "CAAX":   "Membrane",
+    "H2B":    "Nuclei",
+    "SEC61B": "Structure",
+    "TOMM20": "Structure",
+}
+
+ORGANELLE_LABELS: dict[str, str] = {
+    "CAAX":   "Membrane (CAAX)",
+    "H2B":    "Chromatin (H2B)",
+    "SEC61B": "ER (SEC61B)",
+    "TOMM20": "Mitochondria (TOMM20)",
+}
+
+FLUOR_CH = 2  # channel index for fluorescence in the input zarr
+
+# Cache downloaded checkpoints in /tmp so the Space doesn't re-download each run
+_ckpt_cache: dict[str, str] = {}
+
+
+def get_checkpoint(model: str, organelle: str) -> str:
+    """Download (or return cached) checkpoint path for a given model + organelle."""
+    key = (model, organelle)
+    filename = CHECKPOINT_FILES[key]
+    if filename not in _ckpt_cache:
+        print(f"Downloading {filename} from {CHECKPOINT_REPO} ...")
+        local = hf_hub_download(repo_id=CHECKPOINT_REPO, filename=filename)
+        _ckpt_cache[filename] = local
+    return _ckpt_cache[filename]
+
+
+def preprocess_zarr(data_path: str) -> None:
+    """Compute normalization statistics for the uploaded zarr via viscy preprocess."""
+    subprocess.run(
+        ["viscy", "preprocess", f"--data_path={data_path}", "--num_workers=1", "--block_size=32"],
+        check=True,
+    )
+
+
+def create_single_timepoint_zarr(source_path: str, timepoint: int) -> str:
+    """Copy source HCS zarr plate, keeping only the selected timepoint.
+
+    Remaps timepoint_statistics in .zattrs so index "0" carries the selected
+    timepoint's normalization stats (needed by celldiff's MinMaxSampled).
+    """
+    out_path = Path(tempfile.gettempdir()) / f"dynacell_t{timepoint}_{uuid.uuid4().hex[:8]}.zarr"
+    shutil.copytree(source_path, str(out_path))
+
+    src_store = zarr.open(source_path, mode="r")
+    dst_store = zarr.open(str(out_path), mode="r+")
+
+    def _trim(src_grp: zarr.Group, dst_grp: zarr.Group) -> None:
+        for key in list(src_grp.keys()):
+            item = src_grp[key]
+            if isinstance(item, zarr.Array) and key == "0":
+                # Write selected timepoint into index 0, then resize to T=1
+                dst_arr = dst_grp[key]
+                dst_arr[0] = item[timepoint]
+                dst_arr.resize((1,) + item.shape[1:])
+            elif isinstance(item, zarr.Group):
+                _trim(item, dst_grp[key])
+
+    _trim(src_store, dst_store)
+
+    # Remap timepoint_statistics["<timepoint>"] → ["0"] in each FOV's .zattrs
+    def _remap_tp_stats(zattrs_path: Path) -> None:
+        if not zattrs_path.exists():
+            return
+        zattrs = json.loads(zattrs_path.read_text())
+        norm = zattrs.get("normalization", {})
+        changed = False
+        for ch_data in norm.values():
+            if "timepoint_statistics" in ch_data:
+                tp_stats = ch_data["timepoint_statistics"]
+                t_key = str(timepoint)
+                if t_key in tp_stats:
+                    ch_data["timepoint_statistics"] = {"0": tp_stats[t_key]}
+                    changed = True
+        if changed:
+            zattrs_path.write_text(json.dumps(zattrs))
+
+    for row in out_path.iterdir():
+        if not row.is_dir():
+            continue
+        for col in row.iterdir():
+            if not col.is_dir():
+                continue
+            for fov in col.iterdir():
+                if fov.is_dir():
+                    _remap_tp_stats(fov / ".zattrs")
+
+    return str(out_path)
+
+
+@spaces.GPU(duration=120)
+def run_prediction(model: str, organelle: str, data_path: str, timepoint: int) -> str:
+    """Run prediction for a single timepoint; return the output zarr path.
+
+    Creates a single-timepoint subset of the source zarr, runs prediction on it,
+    and returns the path to the output zarr (which has T=1). The `dynacell predict`
+    subprocess inherits the ZeroGPU allocation from this decorated frame.
+    """
+    subset_path = create_single_timepoint_zarr(data_path, timepoint)
+
+    ckpt_path = get_checkpoint(model, organelle)
+    output_dir = Path(tempfile.gettempdir()) / f"dynacell_pred_{uuid.uuid4().hex[:8]}"
+    output_store = str(output_dir / f"{organelle}_{model}.zarr")
+
+    template = (TEMPLATE_DIR / f"{model}.yaml").read_text()
+    config_text = template.format(
+        ckpt_path=ckpt_path,
+        data_path=subset_path,
+        output_store=output_store,
+        target_channel=TARGET_CHANNELS[organelle],
+    )
+
+    config_path = Path(tempfile.gettempdir()) / f"dynacell_cfg_{uuid.uuid4().hex[:8]}.yaml"
+    config_path.write_text(config_text)
+
+    print(f"Running dynacell predict: {model} / {organelle} / t={timepoint}")
+    subprocess.run(["dynacell", "predict", "-c", str(config_path)], check=True)
+    config_path.unlink(missing_ok=True)
+
+    return output_store
+
+
+@spaces.GPU(duration=120)
+def compute_trajectory(
+    organelle: str,
+    data_path: str,
+    timepoint: int = 0,
+    num_steps: int = 50,
+    progress=None,
+) -> dict:
+    """Run the CELL-Diff ODE; save trajectory to /tmp as .npy; return metadata dict.
+
+    The returned dict contains everything needed to call render_trajectory_gif
+    without re-running the ODE.
+    """
+    import numpy as np
+    import torch
+    from iohub.ngff import open_ome_zarr
+    from dynacell.engine import DynacellFlowMatching
+    from viscy_data._utils import _read_norm_meta
+
+    if progress is not None:
+        progress(0.05, desc="Downloading CELL-Diff checkpoint...")
+    ckpt_path = get_checkpoint("celldiff", organelle)
+
+    if progress is not None:
+        progress(0.15, desc="Loading model...")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = DynacellFlowMatching.load_from_checkpoint(ckpt_path, map_location=device)
+    model.eval()
+    patch_d, patch_h, patch_w = model.model.net.input_spatial_size  # (8, 512, 512)
+
+    if progress is not None:
+        progress(0.25, desc="Reading phase data...")
+    with open_ome_zarr(data_path, mode="r") as plate:
+        _, pos = next(plate.positions())
+        phase_ch  = pos.get_channel_index("Phase3D")
+        phase_raw = np.array(pos.data[timepoint, phase_ch])
+        norm_meta = _read_norm_meta(pos)
+
+    tp_stats = norm_meta["Phase3D"]["timepoint_statistics"][str(timepoint)]
+    lo = tp_stats["p1"].item()
+    hi = tp_stats["p99"].item()
+    phase_norm = np.clip(phase_raw.astype(np.float32), lo, hi)
+    phase_norm = 2.0 * (phase_norm - lo) / (hi - lo + 1e-8) - 1.0
+
+    z_total = phase_norm.shape[0]
+    z_start = (z_total - patch_d) // 2
+    phase_crop = phase_norm[z_start:z_start + patch_d, :patch_h, :patch_w]
+
+    if progress is not None:
+        progress(0.35, desc=f"Generating {num_steps}-step ODE trajectory...")
+    phase_tensor = (
+        torch.from_numpy(phase_crop).float()
+        .unsqueeze(0).unsqueeze(0)
+        .to(device)
+    )
+    with torch.no_grad():
+        trajectory = model.model.generate_trajectory(phase_tensor, num_steps=num_steps)
+    traj_np = trajectory[:, 0].cpu().numpy().astype(np.float32)  # (num_steps, 1, D, H, W)
+
+    if progress is not None:
+        progress(0.90, desc="Saving trajectory to disk...")
+    traj_path = str(Path(tempfile.gettempdir()) / f"traj_np_{uuid.uuid4().hex[:8]}.npy")
+    np.save(traj_path, traj_np)
+
+    if progress is not None:
+        progress(1.0, desc="Done.")
+
+    return {
+        "traj_path": traj_path,
+        "z_start": z_start,
+        "patch_d": patch_d,
+        "organelle": organelle,
+        "timepoint": timepoint,
+        "num_steps": num_steps,
+    }
+
+
+def render_trajectory_gif(traj_info: dict, z_patch: int) -> str:
+    """Render a GIF from a cached trajectory at the given patch-relative Z slice.
+
+    Fast — only loads the saved .npy and runs matplotlib; does not re-run the ODE.
+    """
+    import numpy as np
+    import matplotlib.pyplot as plt
+    from matplotlib.animation import FuncAnimation, PillowWriter
+
+    traj_np   = np.load(traj_info["traj_path"])
+    z_start   = traj_info["z_start"]
+    patch_d   = traj_info["patch_d"]
+    organelle = traj_info["organelle"]
+    timepoint = traj_info["timepoint"]
+    num_steps = traj_info["num_steps"]
+
+    z_patch = max(0, min(z_patch, patch_d - 1))
+    z_abs   = z_start + z_patch
+
+    def pnorm(img: np.ndarray) -> np.ndarray:
+        lo_p, hi_p = np.percentile(img, [0.5, 99.5])
+        if hi_p == lo_p:
+            return np.zeros_like(img, dtype=np.float32)
+        return np.clip((img - lo_p) / (hi_p - lo_p), 0, 1).astype(np.float32)
+
+    frame_idx = np.linspace(0, num_steps - 1, min(50, num_steps), dtype=int)
+    fig_a, ax_a = plt.subplots(figsize=(4, 4))
+    ax_a.axis("off")
+    im = ax_a.imshow(
+        pnorm(traj_np[0, 0, z_patch]), cmap="gray", vmin=0, vmax=1, interpolation="nearest"
+    )
+    ttl = ax_a.set_title(
+        f"{ORGANELLE_LABELS[organelle]}  t={timepoint}  z={z_abs}\nStep 0  (noise → prediction)",
+        fontsize=9,
+    )
+
+    def update(frame: int):
+        s = frame_idx[frame]
+        im.set_data(pnorm(traj_np[s, 0, z_patch]))
+        ttl.set_text(
+            f"{ORGANELLE_LABELS[organelle]}  t={timepoint}  z={z_abs}\nStep {s}  (noise → prediction)"
+        )
+        return im, ttl
+
+    anim = FuncAnimation(fig_a, update, frames=len(frame_idx), interval=80, blit=True)
+    gif_path = str(Path(tempfile.gettempdir()) / f"traj_{uuid.uuid4().hex[:8]}.gif")
+    anim.save(gif_path, writer=PillowWriter(fps=12))
+    plt.close(fig_a)
+    return gif_path

requirements.txt

@@ -0,0 +1,14 @@
+gradio>=5.0
+spaces>=0.30
+zarr>=2.16
+numpy>=1.24
+matplotlib>=3.7
+torch>=2.1
+huggingface_hub>=0.20
+iohub>=0.1
+cubic>=0.7.0a2
+git+https://github.com/mehta-lab/VisCy.git@dynacell-models#subdirectory=packages/viscy-data
+git+https://github.com/mehta-lab/VisCy.git@dynacell-models#subdirectory=packages/viscy-models
+git+https://github.com/mehta-lab/VisCy.git@dynacell-models#subdirectory=packages/viscy-transforms
+git+https://github.com/mehta-lab/VisCy.git@dynacell-models#subdirectory=packages/viscy-utils
+git+https://github.com/mehta-lab/VisCy.git@dynacell-models#subdirectory=applications/dynacell