fix cellpose deps

src/cellmap_models/pytorch/cellpose/__init__.py

@@ -1,7 +1,17 @@
-from .add_model import add_model
-from .load_model import load_model
-from .get_model import get_model
-from .download_checkpoint import download_checkpoint
+import lazy_loader as lazy
+
+
+# Keep Cellpose itself out of the import path until one of these helpers is
+# actually used. Cellpose is an optional dependency of cellmap-models.
+__getattr__, __dir__, __all__ = lazy.attach(
+    __name__,
+    submod_attrs={
+        "add_model": ["add_model"],
+        "load_model": ["load_model"],
+        "get_model": ["get_model"],
+        "download_checkpoint": ["download_checkpoint"],
+    },
+)
 
 models_dict = {
     "jrc_fly-mb-z0419-20_nuc_cp": "https://github.com/janelia-cellmap/cellmap-models/releases/download/2024.08.06/jrc_fly-mb-z0419-20_nuc_cp",

src/cellmap_models/pytorch/cellpose/add_model.py

@@ -1,9 +1,6 @@
 import os
 import sys
 from typing import Optional
-from cellpose.io import add_model as _add_model
-from cellpose.models import MODEL_DIR
-from .get_model import get_model
 
 
 def add_model(model_name: Optional[str] = None):
@@ -12,6 +9,11 @@ def add_model(model_name: Optional[str] = None):
     Args:
         model_name (str): model name
     """
+    from cellpose.io import add_model as _add_model
+    from cellpose.models import MODEL_DIR
+
+    from .get_model import get_model
+
     if model_name is None:
         model_name = sys.argv[1]
     base_path = MODEL_DIR

src/cellmap_models/pytorch/cellpose/load_model.py

@@ -2,8 +2,6 @@
 from pathlib import Path
 from typing import Union
 import torch
-from .get_model import get_model
-from cellpose.models import CellposeModel
 
 
 def load_model(
@@ -21,6 +19,10 @@ def load_model(
     Returns:
         model: model
     """
+    from cellpose.models import CellposeModel
+
+    from .get_model import get_model
+
     model_path = get_model(model_name, base_path)
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"

tests/test_cellpose_optional.py

@@ -0,0 +1,27 @@
+import subprocess
+import sys
+
+
+def test_cellpose_is_not_required_to_import_cellmap_models():
+    code = """
+import sys
+sys.modules[\'cellpose\'] = None
+
+import cellmap_models
+import cellmap_models.pytorch
+from cellmap_models.pytorch.cellpose import (
+    add_model,
+    download_checkpoint,
+    get_model,
+    load_model,
+    models_dict,
+    models_list,
+)
+
+assert models_list == list(models_dict)
+assert all(callable(function) for function in (
+    add_model, download_checkpoint, get_model, load_model
+))
+assert \'cellpose\' not in sys.modules or sys.modules[\'cellpose\'] is None
+"""
+    subprocess.run([sys.executable, "-c", code], check=True)