Update image_segmenter.py

image_segmenter.py

@@ -7,23 +7,17 @@ import spaces
 
 class ImageSegmenter:
     def __init__(self, model_type="yolov8s-seg") -> None:
-        self.model_type = model_type
-        self.device = 'cuda'  # ZeroGPU will always use CUDA
-        self.is_show_bounding_boxes = True
-        self.is_show_segmentation_boundary = False
-        self.is_show_segmentation = False
-        self.confidence_threshold = 0.5
-        self.cls_clr = {}
-
-        # params
-        self.bb_thickness = 2
-        self.bb_clr = (255, 0, 0)
-
-        # variables
-        self.masks = {}
-        
-        # Model will be loaded in predict to work with ZeroGPU
-        self.model = None
+    # Don't initialize CUDA here
+    self.model_type = model_type
+    self.is_show_bounding_boxes = True
+    self.is_show_segmentation_boundary = False
+    self.is_show_segmentation = False
+    self.confidence_threshold = 0.5
+    self.cls_clr = {}
+    self.bb_thickness = 2
+    self.bb_clr = (255, 0, 0)
+    self.masks = {}
+    self.model = None  # Model will be loaded in predict
 
     def get_cls_clr(self, cls_id):
         if cls_id in self.cls_clr: