Update binary_segmentation.py

binary_segmentation.py

@@ -495,57 +495,57 @@ 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,  # ✅ Keep FP32 for CPU
+                torch_dtype=torch.float32, 
                 low_cpu_mem_usage=False 
             )
-        
-            # ✅ QUANTIZE to INT8 for CPU speedup
-            if DEVICE == "cpu":
-                
-                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((1024, 1024)),
+                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,  # ✅ Keep FP32 for CPU
+    #             low_cpu_mem_usage=False 
+    #         )
+        
+    #         # ✅ QUANTIZE to INT8 for CPU speedup
+    #         if DEVICE == "cpu":
+                
+    #             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((1024, 1024)),
+    #             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_rmbg(self):
         """Load RMBG model (good balance)"""