Add files via upload

Author: 1038lab

AILab_RMBG.py

@@ -312,15 +312,8 @@ def load_model(self, model_name):
                 import transparent_background
                 self.model = transparent_background.Remover()
                 self.current_model_version = model_name
-            except ImportError:
-                try:
-                    import pip
-                    pip.main(['install', 'transparent_background'])
-                    import transparent_background
-                    self.model = transparent_background.Remover()
-                    self.current_model_version = model_name
-                except Exception as e:
-                    handle_model_error(f"Failed to install transparent_background: {str(e)}")
+            except Exception as e:
+                handle_model_error(f"Failed to initialize transparent_background: {str(e)}")
 
     def process_image(self, image, model_name, params):
         try:

AILab_SegmentV2.py

@@ -210,7 +210,9 @@ def segment_v2(self, image, prompt, sam_model, dino_model, threshold=0.30,
             sam_key = (sam_info["model_type"], sam_ckpt_path, device)
             if sam_key not in self.sam_model_cache:
                 try:
-                    sam = sam_model_registry[sam_info["model_type"]](checkpoint=sam_ckpt_path)
+                    sam = sam_model_registry[sam_info["model_type"]]()
+                    state_dict = torch.load(sam_ckpt_path, map_location="cpu")
+                    sam.load_state_dict(state_dict, strict=False)
                     sam.to(device)
                     self.sam_model_cache[sam_key] = SamPredictor(sam)
                 except RuntimeError as e:

requirements.txt

@@ -1,16 +1,9 @@
 # Base dependencies
-torch>=2.0.0
-torchvision>=0.15.0
-Pillow>=9.0.0
-numpy>=1.22.0
 huggingface-hub>=0.19.0
-transformers>=4.35.0
-safetensors>=0.3.0
-transparent-background>=1.2.4
-tqdm>=4.65.0
+transparent-background>=1.1.2
 segment-anything>=1.0
 groundingdino-py>=0.4.0
 opencv-python>=4.7.0
-scipy>=1.10.0
 onnxruntime>=1.15.0
-onnxruntime-gpu>=1.15.0
+onnxruntime-gpu>=1.15.0
+protobuf>=3.20.2,<6.0.0