Update binary_segmentation.py

binary_segmentation.py

@@ -448,6 +448,11 @@ class BinarySegmenter:
         self.model = None
         self.transform = None
         self._load_model()
+        
+        if DEVICE == "cpu":
+            self.model = self.model.float()
+        self.model.to(DEVICE)
+        self.model.eval()
     
     def _load_model(self):
         """Load the specified segmentation model"""
@@ -497,7 +502,9 @@ class BinarySegmenter:
             self.model = AutoModelForImageSegmentation.from_pretrained(
                 'ZhengPeng7/BiRefNet',
                 trust_remote_code=True,
-                cache_dir=str(self.cache_dir)
+                cache_dir=str(self.cache_dir),
+                torch_dtype=torch.float32, 
+                low_cpu_mem_usage=False 
             )
             
             self.transform = transforms.Compose([
@@ -559,6 +566,9 @@ class BinarySegmenter:
         
         # Transform
         input_tensor = self.transform(image_pil).unsqueeze(0).to(DEVICE)
+        if DEVICE == "cpu":
+            input_tensor = input_tensor.float()
+
         
         # Inference
         with torch.no_grad():