merge repos by mzouink · Pull Request #11 · janelia-cellmap/cellmap-models

mzouink · 2026-02-27T05:12:19Z

No description provided.

Copilot

Pull request overview

This PR introduces a new model_export package to export PyTorch models into multiple deployable formats (pt/ts/onnx/pt2) with associated metadata/README generation, plus utilities to push exports to the Hugging Face Hub. It also refactors dependencies/extras and updates CI installation to include all extras.

Changes:

Add model export utilities (export_torch_model), metadata schema/prompting, and a versioned CellmapModel loader for exported model folders.
Add DaCapo export entrypoint and example scripts for exporting/pushing models.
Restructure dependencies into extras (cellpose/dacapo/export/huggingface) and update CI to install .[all].

Reviewed changes

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

Show a summary per file

File	Description
src/cellmap_models/model_export/generate_metadata.py	Adds `ModelMetadata`, README generation, and metadata export/prompting logic.
src/cellmap_models/model_export/export_model.py	Adds multi-format export and Hugging Face upload helper.
src/cellmap_models/model_export/dacapo_model.py	Adds DaCapo run export helper + CLI entrypoint.
src/cellmap_models/model_export/config.py	Adds export folder configuration.
src/cellmap_models/model_export/cellmap_model_v1.py	Adds runtime loader/checker for exported model directories.
src/cellmap_models/model_export/cellmap_model.py	Adds format-version factory + Hugging Face download helper.
src/cellmap_models/model_export/init.py	Exposes export/public API symbols.
src/cellmap_models/init.py	Exposes `model_export` at top-level.
pyproject.toml	Adds deps/extras for export + new CLI script entrypoint.
examples/export_models/*	Adds example scripts for exporting and pushing models.
README.md	Documents export/load/push workflows.
.github/workflows/tests.yaml	Installs `.[all]` in CI.

Comments suppressed due to low confidence (5)

.github/workflows/tests.yaml:23

CI now installs .[all], which pulls in heavy/optional deps (e.g. dacapo-ml, onnxruntime, huggingface-hub). This increases install time and can break the test matrix on OSes where some extras aren’t available. Prefer installing only what the tests need (e.g. .[dev] plus specific extras per job) or split the workflow into separate jobs for optional integrations.

        pip install ".[all]"

src/cellmap_models/model_export/dacapo_model.py:38

Typo: infernece_output_shape should be inference_output_shape. Even though it’s only a local variable, the misspelling makes this code harder to follow and increases the chance of future mistakes.

    infernece_output_shape = run.model.compute_output_shape(inference_input_shape)[1]
    metadata = ModelMetadata(
        model_name=run_name,
        iteration=iteration,
        model_type=run.model.architecture.__class__.__name__,
        framework="dacapo/torch",
        in_channels=run.model.num_in_channels,
        out_channels=run.model.num_out_channels,
        channels_names=run.task.channels,
        inference_input_shape=inference_input_shape,
        inference_output_shape=infernece_output_shape,

src/cellmap_models/model_export/generate_metadata.py:263

The except Exception as e: raise Exception(...) pattern discards the original traceback. Use exception chaining (raise ... from e) so callers can see the real source of the failure while still adding context.

            except Exception as e:
                raise Exception(f"Error prompting for field {field_name}: {e}")

src/cellmap_models/model_export/export_model.py:6

omnx_version looks like a typo for onnx_version. Since this constant is used as opset_version=omnx_version, the misspelling makes the code harder to read and easier to misuse later—please rename it consistently.

omnx_version = 17

src/cellmap_models/model_export/cellmap_model_v1.py:5

BaseModel and Field are imported but unused in this file, which will trigger linting/type-checker warnings. Please remove the unused imports (or use them if intended).

from typing import Any, Optional
from pydantic import BaseModel, Field
from .generate_metadata import ModelMetadata, CURRENT_FORMAT_VERSION

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

Copilot · 2026-02-27T21:18:41Z

+    return module.CellmapModel(folder_path)
+
+
+def get_huggingface_model(repo_id: str, revision: str | None = None):


This uses the PEP 604 union syntax (str | None), which requires Python 3.10+. The project metadata currently allows Python >=3.7, and CI runs 3.9, so this will fail to import on supported versions. Use Optional[str] (or bump requires-python and CI) to keep versions consistent.

Copilot · 2026-02-27T21:18:42Z

 dependencies = [
  'torch',
  'torchvision',
  'numpy',
  'tqdm',
-  'cellpose',
  'ml-collections',
+  'pydantic',
+  'lazy_loader',
 ]


The base dependency list now includes pydantic, which (in current major versions) requires Python >=3.8. Additionally, new code uses Python 3.10-only syntax (str | None). If this project intends to support older versions, update the code accordingly; otherwise, bump requires-python to match the actual minimum supported version and adjust CI to stop testing unsupported versions.

Copilot · 2026-02-27T21:18:42Z

+    if metadata is not None:
+        from .generate_metadata import export_metadata
+
+        export_metadata(metadata)
+    pt_file = os.path.join(folder_result, "model.pt")


When metadata is provided, export_torch_model() calls export_metadata(metadata) without passing folder_result. export_metadata() writes to get_export_folder()/metadata.model_name, which can differ from folder_result, leaving model artifacts and metadata in different directories. Consider either: (1) writing metadata/readme into folder_result, or (2) making export_torch_model() derive folder_result from export_metadata()’s resolved output path and use that consistently.

Copilot · 2026-02-27T21:18:42Z

+                #   model.load_state_dict(torch.load(pt_path))
+                #   self._pt_model = model
+                # Instead of just torch.load().
+                self._pt_model = torch.load(pt_path, weights_only=True)


export_torch_model() saves the full model object via torch.save(model, pt_file), but CellmapModel.pytorch_model later loads it with weights_only=True, which will not reconstruct the saved nn.Module (and may fail outright depending on PyTorch version). Align the save/load strategy: either save a state_dict and rebuild the module, or load the pickled module without weights_only (with appropriate security notes).

Suggested change

self._pt_model = torch.load(pt_path, weights_only=True)

self._pt_model = torch.load(pt_path)

Copilot · 2026-02-27T21:18:43Z

@@ -4,3 +4,4 @@

 from .utils import download_url_to_file
 from .pytorch import cosem, cellpose, untrained_models


cellpose is now an optional extra in pyproject.toml, but cellmap_models/__init__.py still imports cellpose unconditionally. Because cellmap_models.pytorch.cellpose.load_model imports the external cellpose package, importing cellmap_models will raise ImportError unless users install the optional extra. Make this import conditional (try/except) or remove it from the top-level imports so base installs work without cellpose.

Suggested change

from .pytorch import cosem, cellpose, untrained_models

from .pytorch import cosem, untrained_models

Copilot · 2026-02-27T21:18:43Z

+    result_folder = os.path.join(export_folder, metadata.model_name or "unnamed")
+    if os.path.exists(result_folder) and not overwrite:
+        answer = input(f"Folder {result_folder} already exists. Overwrite? [y/N]: ")
+        if answer.lower() not in ("y", "yes"):
+            return
+    metadata = prompt_for_missing_fields(metadata)


export_metadata() computes result_folder before calling prompt_for_missing_fields(). If model_name is missing and the user enters it during prompting, the metadata will be saved under .../unnamed while metadata.model_name becomes the provided name (folder/name mismatch). Compute result_folder after prompting (or re-compute it if model_name changes) so the output directory matches the finalized metadata.

Suggested change

result_folder = os.path.join(export_folder, metadata.model_name or "unnamed")

if os.path.exists(result_folder) and not overwrite:

answer = input(f"Folder {result_folder} already exists. Overwrite? [y/N]: ")

if answer.lower() not in ("y", "yes"):

return

metadata = prompt_for_missing_fields(metadata)

metadata = prompt_for_missing_fields(metadata)

result_folder = os.path.join(export_folder, metadata.model_name or "unnamed")

if os.path.exists(result_folder) and not overwrite:

answer = input(f"Folder {result_folder} already exists. Overwrite? [y/N]: ")

if answer.lower() not in ("y", "yes"):

return

Copilot · 2026-02-27T21:18:43Z

+        # Export via torch.export (ExportedProgram)
+        dummy_input = torch.rand(input_shape)
+        exported = torch.export.export(model, (dummy_input,))
+        torch.export.save(exported, ep_file)
+        print(f"Model saved to {ep_file}")
+    except Exception as e:
+        print(f"Error exporting with torch.export: {e}")
+
+    try:
+        dummy_input = torch.rand(input_shape)
+        scripted_model = torch.jit.trace(model, dummy_input)
+        scripted_model.save(ts_file)


dummy_input = torch.rand(input_shape) is always created on CPU, but many callers move the model to CUDA (see examples). Tracing/exporting will error if the model and inputs are on different devices. Create dummy_input on the same device/dtype as the model parameters (or accept an optional device argument).

Removed Python 3.12 from the testing matrix.

Rhoades scholar patch 1

rhoadesScholar · 2026-02-27T22:34:52Z

@mzouink Check out Copilot's review comments and either ignore or address them (or tell it to implement its own suggestions). Either way, all should get marked as resolved, and then we merge. 👍🏼

mzouink added 3 commits February 27, 2026 00:11

merge repos

c85f2e9

black format

60bd989

cellpose optional

54110f5

mzouink commented Feb 27, 2026

View reviewed changes

Comment thread examples/export_marwan_dacapo_models/to_be_exported copy.py Outdated

rhoadesScholar requested changes Feb 27, 2026

View reviewed changes

mzouink added 7 commits February 27, 2026 11:56

clean

b04614f

fix

3557681

versionning and huggingface

24ab6c2

fix err

d0fe776

black format

da443da

fix tests

fa610f1

mypy

8a873a6

rhoadesScholar requested a review from Copilot February 27, 2026 21:10

Copilot started reviewing on behalf of rhoadesScholar February 27, 2026 21:11 View session

rhoadesScholar added 2 commits February 27, 2026 16:11

Update Python version requirement and version number

2069681

Update Python versions in CI workflow

2e1dfe3

Copilot AI reviewed Feb 27, 2026

View reviewed changes

rhoadesScholar added 3 commits February 27, 2026 16:59

Merge branch 'export_model' into rhoadesScholar-patch-1

ef22ae9

Update Python versions in tests workflow

a8c3958

Removed Python 3.12 from the testing matrix.

Merge pull request #12 from janelia-cellmap/rhoadesScholar-patch-1

4e118ae

Rhoades scholar patch 1

mzouink merged commit 8ab3af9 into main Mar 2, 2026
8 of 10 checks passed

mzouink deleted the export_model branch March 2, 2026 14:17

Navigation Menu

merge repos - #11

merge repos#11
mzouink merged 15 commits into
mainfrom
export_model

mzouink commented Feb 27, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

rhoadesScholar commented Feb 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return module.CellmapModel(folder_path)


		def get_huggingface_model(repo_id: str, revision: str \| None = None):

	self._pt_model = torch.load(pt_path, weights_only=True)
	self._pt_model = torch.load(pt_path)

		@@ -4,3 +4,4 @@

		from .utils import download_url_to_file
		from .pytorch import cosem, cellpose, untrained_models

	from .pytorch import cosem, cellpose, untrained_models
	from .pytorch import cosem, untrained_models

Conversation

mzouink commented Feb 27, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

rhoadesScholar commented Feb 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants