Update binary_segmentation.py

binary_segmentation.py

@@ -494,19 +494,50 @@ class BinarySegmenter:
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
         ])
     
+    # def _load_birefnet(self):
+    #     """Load BiRefNet model (best accuracy, larger)"""
+    #     try:
+    #         from transformers import AutoModelForImageSegmentation
+            
+    #         self.model = AutoModelForImageSegmentation.from_pretrained(
+    #             'ZhengPeng7/BiRefNet',
+    #             trust_remote_code=True,
+    #             cache_dir=str(self.cache_dir),
+    #             torch_dtype=torch.float32, 
+    #             low_cpu_mem_usage=False 
+    #         )
+            
+    #         self.transform = transforms.Compose([
+    #             transforms.Resize((320, 320)),
+    #             transforms.ToTensor(),
+    #             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    #         ])
+    #     except ImportError:
+    #         raise ImportError("BiRefNet requires: pip install transformers")
+
     def _load_birefnet(self):
         """Load BiRefNet model (best accuracy, larger)"""
         try:
             from transformers import AutoModelForImageSegmentation
-            
+        
             self.model = AutoModelForImageSegmentation.from_pretrained(
                 'ZhengPeng7/BiRefNet',
                 trust_remote_code=True,
                 cache_dir=str(self.cache_dir),
-                torch_dtype=torch.float32, 
+                torch_dtype=torch.float32,  # ✅ Keep FP32 for CPU
                 low_cpu_mem_usage=False 
             )
-            
+        
+            # ✅ QUANTIZE to INT8 for CPU speedup
+            if DEVICE == "cpu":
+                import torch.quantization
+                self.model = torch.quantization.quantize_dynamic(
+                    self.model,
+                    {torch.nn.Linear, torch.nn.Conv2d},
+                    dtype=torch.qint8
+                )
+                logger.info("✅ BiRefNet quantized to INT8")
+        
             self.transform = transforms.Compose([
                 transforms.Resize((320, 320)),
                 transforms.ToTensor(),