Add RTX 50-series CUDA support

INSTALL.md

@@ -11,6 +11,7 @@ It is recommended to install the dependencies of mrSegmentation through a virtua
 
 > [!NOTE]  
 > By default, some of the dependencies in the `requirements.txt` file install a CPU-only version of PyTorch. The plugin will run smoothly on the CPU, but for better performances, it is advised to install the GPU-enabled version of PyTorch.
+> The optional `torch-requirements.txt` file installs PyTorch with CUDA 12.8 support, which is required for NVIDIA RTX 50-series GPUs.
 
 - On Linux:
 ```
@@ -75,4 +76,4 @@ export MESHROOM_PLUGINS_PATH=/path/to/mrSegmentation:$MESHROOM_PLUGINS_PATH
 - On Windows:
 ```
 set MESHROOM_PLUGINS_PATH=/path/to/mrSegmentation;%MESHROOM_PLUGINS_PATH%
-```
+```

docker/Dockerfile_rocky

@@ -28,7 +28,6 @@ RUN python -m venv ${MRPLUGIN_BUNDLE}
 WORKDIR ${MRPLUGIN_BUNDLE}
 
 RUN ${MRPLUGIN_BUNDLE}/bin/pip install -r ${MRPLUGIN_DEV}/requirements.txt
-RUN sed -i -e 's/cu124/cu121/g' ${MRPLUGIN_DEV}/torch-requirements.txt
 RUN ${MRPLUGIN_BUNDLE}/bin/pip install -r ${MRPLUGIN_DEV}/torch-requirements.txt --upgrade --force-reinstall
 
 
@@ -59,4 +58,4 @@ ENV RDS_RECOGNITION_MODEL_PATH="${MRPLUGIN_BUNDLE}/models/ram_plus_swin_large_14
 ## Clean-up
 ## ========
 
-RUN rm -rf share && rm pyvenv.cfg
+RUN rm -rf share && rm pyvenv.cfg

docker/Dockerfile_ubuntu

@@ -34,7 +34,6 @@ RUN python -m venv ${MRPLUGIN_BUNDLE}
 WORKDIR ${MRPLUGIN_BUNDLE}
 
 RUN ${MRPLUGIN_BUNDLE}/bin/pip install -r ${MRPLUGIN_DEV}/requirements.txt
-RUN sed -i -e 's/cu124/cu121/g' ${MRPLUGIN_DEV}/torch-requirements.txt
 RUN ${MRPLUGIN_BUNDLE}/bin/pip install -r ${MRPLUGIN_DEV}/torch-requirements.txt --upgrade --force-reinstall
 
 
@@ -65,4 +64,4 @@ ENV RDS_RECOGNITION_MODEL_PATH="${MRPLUGIN_BUNDLE}/models/ram_plus_swin_large_14
 ## Clean-up
 ## ========
 
-RUN rm -rf share && rm pyvenv.cfg
+RUN rm -rf share && rm pyvenv.cfg

docker/build-rocky.sh

@@ -2,7 +2,7 @@
 set -ex
 
 test -z "$MRSEGMENTATION_VERSION" && MRSEGMENTATION_VERSION="$(git rev-parse --abbrev-ref HEAD)-$(git rev-parse --short HEAD)"
-test -z "$CUDA_VERSION" && CUDA_VERSION=12.1.1
+test -z "$CUDA_VERSION" && CUDA_VERSION=12.8.1
 test -z "$ROCKY_VERSION" && ROCKY_VERSION=9
 
 test -d docker || (

docker/build-ubuntu.sh

@@ -2,7 +2,7 @@
 set -ex
 
 test -z "$MRSEGMENTATION_VERSION" && MRSEGMENTATION_VERSION="$(git rev-parse --abbrev-ref HEAD)-$(git rev-parse --short HEAD)"
-test -z "$CUDA_VERSION" && CUDA_VERSION=12.1.1
+test -z "$CUDA_VERSION" && CUDA_VERSION=12.8.1
 test -z "$UBUNTU_VERSION" && UBUNTU_VERSION=22.04
 
 test -d docker || (

segmentationRDS/segmentation.py

@@ -44,12 +44,35 @@ def cleanstr(s: str) -> str:
     return sclean
 
 
+def get_device(useGPU: bool) -> str:
+    if not useGPU:
+        return 'cpu'
+    if not torch.cuda.is_available():
+        print("Cannot execute on GPU, fallback to CPU execution mode")
+        return 'cpu'
+
+    device_index = torch.cuda.current_device()
+    major, minor = torch.cuda.get_device_capability(device_index)
+    device_arch = f"sm_{major}{minor}"
+    device_compute = f"compute_{major}{minor}"
+    supported_arches = torch.cuda.get_arch_list()
+    if supported_arches and device_arch not in supported_arches and device_compute not in supported_arches:
+        device_name = torch.cuda.get_device_name(device_index)
+        print(
+            f"CUDA device '{device_name}' requires {device_arch}, but this PyTorch "
+            f"build supports: {', '.join(supported_arches)}. "
+            "Fallback to CPU execution mode. For RTX 50-series GPUs, install a "
+            "PyTorch CUDA 12.8+ build."
+        )
+        return 'cpu'
+
+    return 'cuda'
+
+
 class SegmentationRDS:
 
     def __init__(self, RAM_CHECKPOINT_PATH:str, GD_CONFIG_PATH:str, GD_CHECKPOINT_PATH:str, SAM_CHECKPOINT_PATH:str, RAM_VIT:str='swin_l', RAM_IMAGE_SIZE:int=384, SAM_ENCODER_VERSION:str='vit_h', useGPU:bool=True):
-        self.DEVICE = 'cuda' if useGPU and torch.cuda.is_available() else 'cpu'
-        if self.DEVICE == 'cpu' and useGPU:
-            print("Cannot execute on GPU, fallback to CPU execution mode")
+        self.DEVICE = get_device(useGPU)
         self.RAM_IMAGE_SIZE = RAM_IMAGE_SIZE
         # Load models
         # Recognize Anything
@@ -191,9 +214,7 @@ def process(self, image: np.ndarray, prompt: str, synonyms: str = '', invert: bo
 class SegmentAnything:
 
     def __init__(self, SAM_CHECKPOINT_PATH:str, SAM_ENCODER_VERSION:str='vit_h', useGPU:bool=True):
-        self.DEVICE = 'cuda' if useGPU and torch.cuda.is_available() else 'cpu'
-        if self.DEVICE == 'cpu' and useGPU:
-            print("Cannot execute on GPU, fallback to CPU execution mode")
+        self.DEVICE = get_device(useGPU)
         # Load models
         sam = sam_model_registry[SAM_ENCODER_VERSION](checkpoint=SAM_CHECKPOINT_PATH)
         sam.to(self.DEVICE)
@@ -262,9 +283,7 @@ def process(self, image: np.ndarray, bboxes = [], clicksIn: np.ndarray = [], cli
 class RecognizeAnything:
 
     def __init__(self, RAM_CHECKPOINT_PATH:str, RAM_VIT:str='swin_l', RAM_IMAGE_SIZE:int=384, useGPU:bool=True):
-        self.DEVICE = 'cuda' if useGPU and torch.cuda.is_available() else 'cpu'
-        if self.DEVICE == 'cpu' and useGPU:
-            print("Cannot execute on GPU, fallback to CPU execution mode")
+        self.DEVICE = get_device(useGPU)
         self.RAM_IMAGE_SIZE = RAM_IMAGE_SIZE
         # Load models
         # Recognize Anything
@@ -294,9 +313,7 @@ def get_tags(self, image:np.ndarray) -> list[str]:
 class DetectAnything:
 
     def __init__(self, RAM_CHECKPOINT_PATH:str, GD_CONFIG_PATH:str, GD_CHECKPOINT_PATH:str, RAM_VIT:str='swin_l', RAM_IMAGE_SIZE:int=384, useGPU:bool=True):
-        self.DEVICE = 'cuda' if useGPU and torch.cuda.is_available() else 'cpu'
-        if self.DEVICE == 'cpu' and useGPU:
-            print("Cannot execute on GPU, fallback to CPU execution mode")
+        self.DEVICE = get_device(useGPU)
         self.RAM_IMAGE_SIZE = RAM_IMAGE_SIZE
         # Load models
         # Recognize Anything
@@ -420,9 +437,7 @@ class BiRefNetSeg:
     def __init__(self, modelType:str, useGPU:bool=True):
         from birefnet.models.birefnet import BiRefNet
 
-        self.DEVICE = 'cuda' if useGPU and torch.cuda.is_available() else 'cpu'
-        if self.DEVICE == 'cpu' and useGPU:
-            print("Cannot execute on GPU, fallback to CPU execution mode")
+        self.DEVICE = get_device(useGPU)
         # Load models
         pretrainedModel = 'ZhengPeng7/BiRefNet_HR-matting'
         if modelType == 'BiRefNet LR':

torch-requirements.txt

@@ -1,3 +1,3 @@
--i https://download.pytorch.org/whl/cu124
-torch==2.5.1
-torchvision==0.20.*
+-i https://download.pytorch.org/whl/cu128
+torch==2.12.0
+torchvision==0.27.0