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)"""