From 644250171f7b7f7a19740f5dfc6fb8e8f9c7406d Mon Sep 17 00:00:00 2001
From: Marwan Zouinkhi <zouinkhi.marwan@gmail.com>
Date: Fri, 31 Jul 2026 13:21:08 -0400
Subject: [PATCH] fix cellpose deps

---
 .../pytorch/cellpose/__init__.py              | 18 ++++++++++---
 .../pytorch/cellpose/add_model.py             |  8 +++---
 .../pytorch/cellpose/load_model.py            |  6 +++--
 tests/test_cellpose_optional.py               | 27 +++++++++++++++++++
 4 files changed, 50 insertions(+), 9 deletions(-)
 create mode 100644 tests/test_cellpose_optional.py

diff --git a/src/cellmap_models/pytorch/cellpose/__init__.py b/src/cellmap_models/pytorch/cellpose/__init__.py
index 153f67f..4c0497d 100644
--- a/src/cellmap_models/pytorch/cellpose/__init__.py
+++ b/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",
diff --git a/src/cellmap_models/pytorch/cellpose/add_model.py b/src/cellmap_models/pytorch/cellpose/add_model.py
index 969fa01..b8d1ba3 100644
--- a/src/cellmap_models/pytorch/cellpose/add_model.py
+++ b/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
diff --git a/src/cellmap_models/pytorch/cellpose/load_model.py b/src/cellmap_models/pytorch/cellpose/load_model.py
index 3b77ae4..d6f76e8 100644
--- a/src/cellmap_models/pytorch/cellpose/load_model.py
+++ b/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"
diff --git a/tests/test_cellpose_optional.py b/tests/test_cellpose_optional.py
new file mode 100644
index 0000000..b59ff76
--- /dev/null
+++ b/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)