Add GPU support for Cellpose

conf/modules.config

@@ -19,4 +19,36 @@ process {
             saveAs: { filename -> filename.equals('versions.yml') ? null : filename },
         ]
     }
+
+    withName: 'DOWNLOAD_CELLPOSE_MODEL|PATCH_SEGMENTATION_CELLPOSE|RESOLVE_CELLPOSE' {
+        ext.cellpose_v4 = { params.cellpose_version?.toString() == '4' }
+
+        conda = {
+            def gpu = workflow.profile.tokenize(',').contains('gpu')
+            def env_name = task.ext.cellpose_v4
+                ? (gpu ? 'environment_gpu_v4.yml' : 'environment_v4.yml')
+                : (gpu ? 'environment_gpu_v3.yml' : 'environment.yml')
+            "${projectDir}/modules/local/patch_segmentation_cellpose/${env_name}"
+        }
+
+        container = {
+            def singularity = workflow.containerEngine in ['singularity', 'apptainer'] && !task.ext.singularity_pull_docker_container
+            def gpu = workflow.profile.tokenize(',').contains('gpu')
+            task.ext.cellpose_v4
+                ? (gpu
+                    ? (singularity
+                        ? 'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/67/671fb3f9bda5ec16374e0e3e205a9c6ffc70ae1c026a545b5a279e7ceb0eb30b/data'
+                        : 'community.wave.seqera.io/library/python_sopa_pytorch-gpu_cuda-version_pruned:f8ad553fec147735')
+                    : (singularity
+                        ? 'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/77/772148940c7b9c9d6534e7171b62e91054f64074b8aa993a00808a63b08c5921/data'
+                        : 'community.wave.seqera.io/library/python_sopa_git_pip_pruned:6cb26eea74b966e9'))
+                : (gpu
+                    ? (singularity
+                        ? 'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/ce/ce822fd366041478434e9969270adc4d7d7e6ce3600c8a20e668d1cd9892792f/data'
+                        : 'community.wave.seqera.io/library/python_sopa_cellpose_pytorch-gpu_cuda-version:27dd8cb2cf96f33b')
+                    : (singularity
+                        ? 'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/22/22d62d6425b70620138ad8764139528c1acaabf6cd06403134c8439caa1c9a31/data'
+                        : 'community.wave.seqera.io/library/python_sopa_cellpose:d098579826bbcf24'))
+        }
+    }
 }

modules/local/download_cellpose_model/main.nf

@@ -2,6 +2,7 @@ process DOWNLOAD_CELLPOSE_MODEL {
     label "process_single"
 
     conda "${moduleDir}/environment.yml"
+
     container "${ workflow.containerEngine in ['singularity', 'apptainer'] && !task.ext.singularity_pull_docker_container
 ?         'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/22/22d62d6425b70620138ad8764139528c1acaabf6cd06403134c8439caa1c9a31/data'
 :         'community.wave.seqera.io/library/python_sopa_cellpose:d098579826bbcf24' }"

modules/local/patch_segmentation_cellpose/environment_gpu_v3.yml

@@ -0,0 +1,10 @@
+# Built with wave --conda-file environment_gpu.yml --freeze --await 30m --platform linux/amd64
+channels:
+  - conda-forge
+  - bioconda
+dependencies:
+  - python=3.12
+  - sopa=2.2.11
+  - cellpose=3.1.1.3
+  - pytorch-gpu=2.13.0
+  - cuda-version=12.9

modules/local/patch_segmentation_cellpose/environment_gpu_v4.yml

@@ -0,0 +1,15 @@
+# Built with wave --conda-file environment_gpu_v4.yml --freeze --await 30m --platform linux/amd64
+channels:
+  - conda-forge
+  - bioconda
+dependencies:
+  - python=3.12
+  - sopa=2.2.11
+  - pytorch-gpu=2.13.0
+  - cuda-version=12.9
+  - torchvision=0.28.0
+  - git
+  - pip
+  - pip:
+      - cellpose==4.2.1.1
+      - git+https://github.com/facebookresearch/dinov3

modules/local/patch_segmentation_cellpose/environment_v4.yml

@@ -0,0 +1,14 @@
+# Built with wave --conda-file environment_v4.yml --freeze --await 30m --platform linux/amd64
+channels:
+  - conda-forge
+  - bioconda
+dependencies:
+  - python=3.12
+  - sopa=2.2.11
+  - pytorch=2.13.0
+  - torchvision=0.28.0
+  - git
+  - pip
+  - pip:
+      - cellpose==4.2.1.1
+      - git+https://github.com/facebookresearch/dinov3

modules/local/patch_segmentation_cellpose/main.nf

@@ -1,8 +1,10 @@
 process PATCH_SEGMENTATION_CELLPOSE {
     label "process_single"
+    label "process_gpu"
     tag "${meta.sample}"
 
     conda "${moduleDir}/environment.yml"
+
     container "${ workflow.containerEngine in ['singularity', 'apptainer'] && !task.ext.singularity_pull_docker_container
 ?         'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/22/22d62d6425b70620138ad8764139528c1acaabf6cd06403134c8439caa1c9a31/data'
 :         'community.wave.seqera.io/library/python_sopa_cellpose:d098579826bbcf24' }"

modules/local/resolve_cellpose/main.nf

@@ -3,6 +3,7 @@ process RESOLVE_CELLPOSE {
     tag "${meta.sample}"
 
     conda "${moduleDir}/environment.yml"
+
     container "${ workflow.containerEngine in ['singularity', 'apptainer'] && !task.ext.singularity_pull_docker_container
 ?         'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/22/22d62d6425b70620138ad8764139528c1acaabf6cd06403134c8439caa1c9a31/data'
 :         'community.wave.seqera.io/library/python_sopa_cellpose:d098579826bbcf24' }"

modules/local/utils.nf

@@ -48,7 +48,7 @@ def extractSubArgs(Map args, String group) {
             flow_threshold: args.flow_threshold,
             model_type: args.cellpose_model_type,
             pretrained_model: args.pretrained_model,
-            gpu: args.cellpose_use_gpu,
+            gpu: workflow.profile.contains("gpu"),
             min_area: args.min_area_pixels2,
             clip_limit: args.clip_limit,
             clahe_kernel_size: args.clahe_kernel_size,

nextflow.config

@@ -70,7 +70,7 @@ params {
     cellprob_threshold         = null
     cellpose_model_type        = null
     pretrained_model           = null
-    cellpose_use_gpu           = false
+    cellpose_version           = null
     cellpose_kwargs            = null
 
     // Stardist segmentation options

nextflow_schema.json

@@ -551,9 +551,10 @@
                     "type": "string",
                     "description": "Cellpose `pretrained_model` parameter"
                 },
-                "cellpose_use_gpu": {
-                    "type": "boolean",
-                    "description": "Whether to use GPU for Cellpose segmentation"
+                "cellpose_version": {
+                    "type": ["integer", "string"],
+                    "enum": [3, 4, "3", "4"],
+                    "description": "Cellpose major version (3 or 4). Defaults to Cellpose v3. The 'gpu' profile selects the pytorch/CUDA image of the chosen version."
                 },
                 "cellpose_kwargs": {
                     "type": "string",

subworkflows/local/utils_nfcore_sopa_pipeline/main.nf

@@ -338,6 +338,10 @@ def validateParams(params) {
     if (params.use_proseg && !is_visium_hd && params.patch_width_microns != null && params.patch_width_microns != -1) {
         log.warn("Proseg needs to run on one single patch, but received patch_width_microns=${params.patch_width_microns}: setting it to -1 instead.")
     }
+    def cellpose_v4 = params.cellpose_version != null ? params.cellpose_version.toString() == '4' : workflow.profile.contains('gpu')
+    if (params.use_cellpose && cellpose_v4 && params.cellpose_model_type != null) {
+        log.warn("Cellpose v4 only supports the pretrained_model argument: 'cellpose_model_type=${params.cellpose_model_type}' will be ignored.")
+    }
 
     return params
 }