Modularize VisCy

@edyoshikun

Modularize VisCy

Restructure VisCy into a uv workspace monorepo to enable reusing transforms, dataloaders, and models in other projects without pulling in all of VisCy.

Implementation Details

Subpackages

Here are the proposed subpackages:

Package	Contents	Workspace Deps	External Deps
`viscy-transforms`	All augmentation transforms (GPU + CPU)	none	kornia, monai, torch
`viscy-data`	DataLoaders and Lightning DataModules	none	lightning, iohub, torch
`viscy-models`	unet/, representation/, translation/, preprocessing/, trainer.py	none	lightning, timm, torch
`viscy-airtable`	Airtable integration	none	pyairtable

graph BT
    transforms[viscy-transforms<br/><i>independent</i>]
    data[viscy-data<br/><i>independent</i>]
    models[viscy-models<br/><i>independent</i>]
    viscy[viscy<br/><i>meta-package + CLI</i>]
    airtable[viscy-airtable<br/><i>independent</i>]

    transforms --> viscy
    data --> viscy
    models --> viscy
    airtable --> viscy

One of the components we'd have to figure out a bit more is the utils/ submodule. I'm thinking this can be partitioned out with the relevant utility files and methods moving to those subpackages.

viscy/
├── applications/               # Publication-based packages
│   ├── cytoland/
│   │   └── examples/           # Quickstarts, demos, notebooks
│   ├── dynaclr/
│   │   ├── examples/           # Jupyter notebooks, demos
│   │   └── evaluation/         # Evaluation metrics code
│   └── dynacell/
│       ├── examples/
│       └── evaluation/
├── packages/
│   ├── viscy-data/
│   │   ├── src/viscy_data/
│   │   ├── tests/              # Unit tests
│   │   └── pyproject.toml
│   ├── viscy-models/
│   │   ├── src/viscy_models/
│   │   ├── tests/              # Unit tests
│   │   └── pyproject.toml
│   ├── viscy-transforms/
│   │   ├── src/viscy_transforms/
│   │   ├── tests/              # Unit tests
│   │   └── pyproject.toml
│   └── viscy-airtable/
│       ├── src/viscy_airtable/
│       ├── tests/              # Unit tests
│       └── pyproject.toml
├── scripts/                    # Standalone scripts (not part of packages)
├── src/viscy/                  # Meta-package source (CLI, re-exports)
├── tests/                      # Integration tests (cross-package)
├── pyproject.toml              # Workspace config + meta-package
└── uv.lock                     # Shared lockfile

There is a way to enforce module boundaries (i.e. ensuring that viscy-models doesn't import from viscy-data) with tach which checks imports at lint-time and has a pre-commit hook, but I don't believe we need it.

Testing

This can get a bit trickier, but from what I've read, we should have tests for each of the packages, and then run integration tests for the VisCy "meta" package.

Each package's tests can be ran independently with uv run --package <name> pytest or pytest packages/<name>/tests/ as well.

Integration tests should instantiate DataModules, LightningModules, and run one iteration of training with fast_dev_run=True. Required configs:

CytoLand pre-training with 3D images
CytoLand virtual staining with 2D images
Time-aware contrastive training with 3D movies
Classical contrastive training with 2D movies

Versioning and Releases

There are a couple of ways to do this:

We can sync / coordinate versioning and releases with all subpackages and the VisCy meta package. This is useful for when packages are very tightly coupled and when users install all the packages together.
Each package has its own version lifecycle, this is good if we expect the subpackages to have different release cadences, and if people will only like to install a 1-2 of the subpackages.

Decision: Option 2.

Open Questions

Applications: Publication-based packages (CytoLand, DynaCLR, DynaCell) that depend on viscy or individual subpackages. Reproducible examples can serve as integration tests.

To execute this somewhat smoothly, We will create a new target branch, maybe call it modular-viscy or something, and split each of the subpackages into their own sub-issue. We'll go through and implement each subpackage one at a time in their own PRs so we won't be reviewing too much at any given time, then a final check before modular-viscy gets merged into main.

Here's an order from what seems the least coupled to the most coupled:

Workspace and scaffolding: Setting up the structure and creating an empty meta package to add to.
viscy-transforms: Extract the transforms
viscy-data: Extract the dataloaders/dataset and define a minimal base dataset/dataloader that can be subclassed (inspired by flowbench's base_model)
viscy-airtable: Extract airtable functionality into separate package.
viscy-models: Extract the models (I think this might be the most work tbh)
Finalize viscy: Reincorporate the CLI, re-exports, and add integration-tests.
Hydra integration: Add Hydra with Lightning (@edyoshikun + @alxndrkalinin).
Update downstream: Fix anything downstream that ends up breaking (maybe Biahub).

Some references I looked at

UV Workspaces Documentation
UV Workspace Dependencies
LlamaIndex Monorepo Case Study
LiveKit Agents - Example of ~50-package uv workspace
Tweag: Python Monorepo CI - Monorepo testing
Managing Python workspaces with uv

@srivarra

Hi @srivarra @edyoshikun @alxndrkalinin
Thanks for writing up these plans! These make sense to me.

├── tests/ # Integration tests (cross-package)
Our integration tests should "exercise" all the LightningModules and DataModules to see that we can do one iteration of model training as the code is refactored. I suggest finding following four configs and wrapping them in an integration test. The integration test should instantiate the data module, lightning module, and run one iteration of training with fast_dev_run=True option.

CytoLand pre-training with 3D images,
CytoLand virtual staining with 2D images.
Time-aware contrastive training with 3D movies.
Classical contrastive training with 2D movies.

@mattersoflight

@mattersoflight

Our integration tests should "exercise" all the LightningModules and DataModules...

Makes sense to me, I'll add that in.

@alxndrkalinin

I would like to add hydra with lightning. This can happen after abstracting the airtable, batch transforms, and dataloaders since it does break the CLI. For reference, it looks like it is a good way for us to keep track of models and easily do any sweeps like in flowbench. If we decide to add this, I already have a PR that I was going to test soon. @alxndrkalinin and I can take care of this part

I'll let you decide if you want to merge airtable PR to this PR or subsequent one.

@srivarra

Otherwise this plan looks great to me! thanks @srivarra

I agree that the plan looks good, and also in favor for adding hydra.

For applications/ I imagine it also being a set of separate modules that depend on the main viscy meta-package (or on individual viscy- subpackages) and implement project-specific analyses.

And I'd also vote for individual package versioning, some could be stable for a while, while other require more frequent changes.

Overall, LGTM! 🚀

@srivarra

@srivarra I agree with your and Alex's point of view of releasing each package independently.
It make sense to me that applications are packages that depend on specific sub-pacakges. I'd map applications to units of publications, e.g., CytoLand, DynaCLR, and DynaCell. What do you think @alxndrkalinin @edyoshikun?

Makes sense to me. If applications at some point can start sharing lots of code, then it can be refactored out to one of the viscy- packages. It'd be nice to have reproducible examples for each such application so we can use them as integration tests like we discussed.

@srivarra

@srivarra for the applications folder I think we will do the following:

# DynaCLR project
/applications/DynaCLR/examples/{jupter_notebooks or demos}
/applications/DynaCLR/evaluation/{code for evaluation metrics.py}

# CytoLand project
/applications/CytoLand/examples/quickstart.py
/applications/CytoLand/examples/quickstart.py

We will get rid of the examples folder at the current root folder and have them one per project (it's ugly and unecessary).

@edyoshikun

@edyoshikun Sounds good, I will adjust the design doc with this feedback!

@alxndrkalinin

@alxndrkalinin and I just discussed. we will make a separate packages/viscy-utils for the shared dependencies between different applications.

@edyoshikun

@edyoshikun @alxndrkalinin Sure that makes sense.

Navigation Menu

Repository navigation

Modularize VisCy #353

Modularize VisCy

Implementation Details

Subpackages

Testing

Versioning and Releases

Open Questions

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Modularize VisCy #353

Description

Modularize VisCy

Implementation Details

Subpackages

Testing

Versioning and Releases

Open Questions

Activity

mattersoflight commented on Jan 21, 2026

srivarra commented on Jan 21, 2026

edyoshikun commented on Jan 22, 2026

edyoshikun commented on Jan 22, 2026

alxndrkalinin commented on Jan 22, 2026

mattersoflight commented on Jan 23, 2026

alxndrkalinin commented on Jan 23, 2026

edyoshikun commented on Feb 3, 2026

srivarra commented on Feb 3, 2026

edyoshikun commented on Feb 12, 2026

srivarra commented on Feb 12, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions