refactor(virtual-stain): rewrite CLI on cytoland predict_sliding_windows

[ieivanov]

Summary

Rewrites biahub virtual-stain to run virtual staining in-process via cytoland (modular VisCy) instead of shelling out to the viscy predict CLI through a separate conda env. Normalization and test-time augmentation are delegated to VisCy so behavior stays in sync with viscy predict.

• Add cytoland under the stain extra (sourced from VisCy's feat/vs-nonsquare-tta-and-helpers branch via [tool.uv.sources]).
• Validate the predict config against VisCy's own VSUNet / HCSDataModule signatures using jsonargparse — biahub keeps no copy of the model schema, so it can't drift from VisCy.
• Dispatch one GPU job per position; each job runs AugmentedPredictionVSUNet.predict_sliding_windows, looping over timepoints on a single GPU (GPU-bound, no per-timepoint parallelism).
• Output store contains only the predicted virtual-stain channels.
• Normalization: reuse the config's instantiated NormalizeSampled transforms + viscy_data.read_norm_meta (reads precomputed stats from the store; errors with a "run viscy preprocess first" message if missing).
• TTA: reuse AugmentedPredictionVSUNet.with_rotation_tta (rotation TTA, correct for non-square FOVs); configurable, default on.
• Adopt the deskew --cluster {slurm,local,debug} / --init pattern: cluster() + init_only() parsing decorators, get_submitit_cluster(cluster=...), RESOURCES: emission, and --cluster debug running jobs in-process in the foreground.
• Add a minimal example config (settings/example_virtual_stain_settings.yml) pointing at the vscyto3d finetune checkpoint.

Dependencies

Depends on mehta-lab/VisCy#450 (non-square rotation TTA + rotation_tta_transforms / with_rotation_tta / read_norm_meta). Pins iohub 0.3.6 (PyPI). Intentionally not built on refactor-deskew-clean, whose iohub fork (old zarr/tensorstore codec architecture) is runtime-incompatible with cytoland's modern iohub; the --cluster/--init infra was instead replicated locally. Once VisCy#450 merges and releases, repoint [tool.uv.sources] cytoland to a stable ref.

Test plan

• ruff clean; biahub virtual-stain --help shows --cluster/--init.
• uv lock / uv sync --extra all resolve (cytoland + viscy-* from the VisCy branch).
• End-to-end --cluster debug on a non-square (64×48) 2-position plate with TTA on → correct (T, 2, Z, 64, 48) finite output; --init initializes the store and exits; RESOURCES: emitted.

🤖 Generated with Claude Code

[ieivanov]

Ready for review. First tests in /hpc/projects/intracellular_dashboard/refactor_biahub/phase2_dev/mantis_v2_dev/refactor-tests/test_biahub_vs/, we should still do a side-by-side comparison with viscy predict.

[ieivanov]

Based on rough first speed measurements, Claude says this is 2x faster than viscy predict. Worth looking closer

  ~2.2× faster. This confirms the earlier prediction: the viscy path was ~87% I/O-bound (num_workers=0, re-reading each z-window), whereas biahub
  loads the full ZYX volume once per timepoint and runs all 72 windows on-GPU from memory — eliminating the redundant per-window reads.

[ieivanov]

FYI @aofei-liu this workflow will significantly simplify how we run virtual staining. For now, we'll do viscy preprocess followed by biahub virtual-stain. Next, I'll try to eliminate the preprocess steps by computing the needed metrics after phase recon, likely optional, on by default.

[ieivanov]

P.S. When reviewing i'd suggest just looking at the new biahub/virtual_stain.py code, the diff makes no sense since this is a complete rewrite (leveraging @tayllatheodoro predict_sliding_windows function in VisCy)

[Copilot]

Pull request overview

This PR rewrites the biahub virtual-stain command to run virtual staining in-process using cytoland/VisCy (via predict_sliding_windows) instead of shelling out to the viscy predict CLI in a separate environment. It also adopts the --cluster {slurm,local,debug} / --init orchestration pattern and adds an example VisCy-compatible predict config.

Changes:

• Replace the legacy VisCy CLI subprocess workflow with an in-process cytoland-based implementation and per-position GPU job dispatch.
• Add jsonargparse-based config validation against VisCy’s VSUNet / HCSDataModule, plus a new example virtual-stain config file.
• Add stain extra dependency plumbing (cytoland source pinning via uv sources) and update example-settings tests to include the new YAML.

Reviewed changes

Copilot reviewed 6 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
biahub/virtual_stain.py	Replaces subprocess-based workflow with cytoland/VisCy in-process prediction, output store init, and submitit-based per-position job dispatch.
biahub/cli/utils.py	Extends submitit cluster selection to support explicit cluster= selection (while still forcing debug in CI).
biahub/cli/parsing.py	Adds --cluster and --init click parsing decorators used by the new virtual-stain CLI.
tests/test_example_settings.py	Ensures the new example YAML is accounted for and validates it via jsonargparse when cytoland is installed.
settings/example_virtual_stain_settings.yml	Adds a minimal example VisCy-style predict configuration for virtual staining.
pyproject.toml	Adds cytoland to the stain extra and configures an uv git source for cytoland.
uv.lock	Locks cytoland + VisCy branch dependencies and related transitive dependency updates.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ieivanov]

Suggestion on Slack from @mattersoflight:

After thinking about our conversation at lab meeting and brainstorming with claude, I land on adding a tightly scoped viscy.deploy subpackage that uses torch.export that explicitly supports dynamic_shapes for specific arguments. Here is an example:

• Pick one model (VSCyto2D) and read the nn.Module + state_dict out of its Lightning checkpoint.
• Write the deployment wrapper: normalize → pad or augment if needed → run inference → crop if needed.
• Export with torch.export (dynamic H/W) and validate against the viscy predict CLI.

With a verifiable test set of (inputs, predictions) for VSCyto2D and DynaCLR, I think we can have a package. This is preferable over ONNX, because ONNX is its own runtime that can drift from features of core pytorch library.

[ieivanov]

I have no strong opinion on this. I asked Claude to compare the two approaches.

What each approach actually is

This PR (#267 + VisCy #450) — "library reuse / in-process." biahub depends on cytoland (modular VisCy) as a library. A viscy predict YAML is validated against VisCy's own VSUNet/HCSDataModule classes via jsonargparse, the Lightning checkpoint's state_dict is loaded into the nn.Module, and all the hard parts are delegated to VisCy: NormalizeSampled + read_norm_meta for normalization, with_rotation_tta for TTA, predict_sliding_windows for the windowed blend. biahub owns only orchestration and keeps no copy of the model schema. A model = (config + checkpoint).

The proposal — "viscy.deploy + torch.export." Produce a self-contained exported artifact per model: a wrapper (normalize → pad/augment → infer → crop) traced with torch.export, dynamic_shapes on H/W, shipped with a golden (inputs, predictions) test set, validated against the viscy predict CLI.

Key framing: #267 changes where orchestration lives; the proposal changes how the model is packaged. They overlap as answers to "how should biahub run VS models," but optimize for different goals — fast iteration inside the VisCy ecosystem vs. portable, frozen, contract-bound distribution.

A subtlety on the ONNX reasoning: "torch.export avoids a second runtime that drifts" is a correct argument for torch.export over ONNX (it keeps real ATen ops, lowers via AOTInductor, tracks PyTorch numerics far more faithfully than the ONNXRuntime op-set). But that's a comparison between two export technologies — it is not an argument against the library-reuse approach in this PR, which has zero drift precisely because it runs the real VisCy code.

Pros and cons

	Library reuse (#267)	torch.export (viscy.deploy)
Numeric fidelity to viscy predict	Exact — it is the same code	Close, but a second path to validate/keep in sync
Preprocessing/TTA/sliding-window	Reused for free from VisCy	Must be re-expressed; the loop + data-dependent windows are the hard part to trace
Adding/switching models	Add a YAML + checkpoint, zero code	Per-model export + dynamic_shapes spec + golden set + re-export on every retrain
Reproducibility	Tied to the installed VisCy version (note the recurring "bump cytoland pin" commits)	Tied to the artifact — strong, versioned contract
Portability / env weight	Heavy: full Lightning/MONAI/jsonargparse/viscy stack everywhere	Light: artifact runs without VisCy installed
Coupling / churn risk	Tight — VisCy internal changes can break biahub	Decoupled once exported, but the export API itself still matures across PyTorch versions
Engineering cost now	Already done, working, ~2.2× speedup	Real per-model work, most of it in the parts torch.export traces worst

The honest weak point of the proposal: the difficulty of VS is not the forward pass — it's the sliding-window blend, the per-FOV data-dependent normalization statistics (norm_meta read at runtime), and rotation TTA. Those are Python loops, reductions, and data-dependent shapes — exactly what torch.export handles poorly. You realistically end up either (a) exporting only the per-window forward and keeping the windowing/TTA/normalization glue in Python — which undercuts the "self-contained" claim — or (b) fighting the tracer to bake them in. Normalization stats would also need to become inputs to the exported program rather than baked constants.

Switching between VSCyto2D / VSCyto3D / VSNeuroMast

Library reuse wins decisively for model-switching. Those three differ only in architecture hyperparameters, stack depth, and source/target channels — all in the YAML and validated against VSUNet. Switching or adding a model is a config + checkpoint swap with no code; this PR already handles 2D vs 3D transparently via out_stack_depth.

For torch.export, switching is expensive and brittle in proportion to how the zoo grows:

• VSCyto2D vs VSCyto3D differ in dimensionality — 3D adds a dynamic Z plus stack-depth constraints, so the dynamic_shapes spec is genuinely different per model, not a parameter tweak.
• VSNeuroMast has its own stack depth/channel layout → its own wrapper and golden set.
• Every retrain or architecture tweak forces a re-export + re-validation.

So export is least suited to exactly the multi-model-switching scenario, and most suited to its opposite (a small frozen flagship set distributed widely).

Recommendation

Keep this PR as biahub's production path — for HPC batch inference over plates with an evolving zoo of VS models inside the Biohub Python ecosystem, library reuse is correct and per-model export buys nothing against the real bottleneck. The one cost it carries (the VisCy version pin) is a dependency-hygiene problem, not an architecture problem.

Pursue viscy.deploy/torch.export as a separate productization effort: packaging one or two flagship models (VSCyto2D, DynaCLR) as portable, contract-bound artifacts for external consumers who can't carry the VisCy stack and need numerics frozen across PyTorch upgrades. There the golden (inputs, predictions) test set is the right deliverable and torch.export-over-ONNX is the right call — just scope it honestly: start with the 2D single-window forward + normalization-as-input, prove dynamic-H/W export and golden-set validation end-to-end, and treat sliding-window/TTA as an explicit design decision rather than an afterthought, because that's where the effort actually is.

_{🤖 Summary generated with Claude Code.}

[ieivanov]

P.S. This PR will be useful right away. I'll merge tomorrow, and then we can refactor again if needed.

[mattersoflight]

It makes sense to use core VisCy to re-use the test-time-augmentation and blending. Please merge.

[ieivanov]

Speed: biahub virtual-stain vs viscy predict

Benchmarked on mantis v2 data (2026_03_26_A549_CAAX_H2B_DENV_ZIKV, 4 positions, each 5 timepoints, volume (Z,Y,X)=(86,1600,1370)). Both paths do identical work — 360 window-inferences per position (5 timepoints × 72 sliding windows) — with the same checkpoint, fcmae model, fov_statistics normalization, and median TTA. The only variable is the inference path: biahub's in-process predict_sliding_windows vs viscy's DataLoader-driven HCSPredictionWriter.

As-run (⚠️ different GPUs)

	biahub virtual-stain	viscy predict
GPU	H100 / H200	A6000
Time/position	~124 s (predict loop) · ~138 s (job)	~1006 s · ~1038 s

~8× — but not apples-to-apples: biahub set no --constraint and landed on H100/H200; the viscy control was constrained to a100\|a6000\|a40 and landed on A6000.

Same GPU class (A6000 vs A6000)

Comparing a prior biahub A6000 run against the viscy A6000 run (total job elapsed):

	biahub virtual-stain (A6000)	viscy predict (A6000)
Time/position	~368 s	~960–1070 s

→ ~2.5–2.8× faster on equal hardware.

Bottom line

biahub virtual-stain is ~2.5× faster than viscy predict on equal hardware. It loads the full ZYX volume once per timepoint and runs all 72 windows on-GPU from memory, eliminating viscy's per-window DataLoader round-trips.

Note on GPU constraints: the ~8× wall-clock gap in the as-run numbers came almost entirely from hardware — biahub grabbed H100/H200 because it was unconstrained. We should not pin reconstruction/prediction jobs to a specific (weaker) GPU when a more powerful one is available; leaving them unconstrained lets SLURM schedule the fastest free GPU. (The --constraint on the viscy run was only there to set up a controlled comparison.)

_{Caveats: the A6000 biahub number is n=1 and from a slightly older revision (pre per-timepoint logging), so it's total job elapsed rather than the isolated predict loop — if anything that understates biahub, since viscy's predict loop (~1006 s) is nearly its full job time.}

[ieivanov]

Side-by-side comparison here:
biahub virtual-stain: /hpc/projects/intracellular_dashboard/refactor_biahub/phase2_dev/mantis_v2_dev/refactor-tests/test_biahub_vs/
viscy predict: /hpc/projects/intracellular_dashboard/refactor_biahub/phase2_dev/mantis_v2_dev/refactor-tests/test_biahub_vs/viscy_predict/

Other clear wins - biahub virtual-stain writes zarr stores chunked in XYZ. viscy predict chunks in XY only currently

Effect of TTA bug fixed in mehta-lab/VisCy#465

green - biahub virtual-stain nuclei and membrane through predict_sliding_windows coda path
magenta - viscy predict nuclei and membrane without the TTA fix as proposed in mehta-lab/VisCy#451

@edyoshikun It looks like the TTA fix helps reduce artifacts at the edge of the images.

[ieivanov]

One more example