update

app.py

@@ -100,8 +100,8 @@ with gr.Blocks() as app:
 
 
 vsa = VisionSearchAssistant(
-    ground_device = "cuda:0",
-    vlm_device="cuda:0",
+    ground_device = "cuda",
+    vlm_device="cuda",
     vlm_load_4bit=True,
 )
 # Launch the app

models/vsa_model.py

@@ -278,7 +278,6 @@ class VisionSearchAssistant:
     the $\mathcal{F}_{vlm}(cdot)$ in the paper.
 
     """
-    @spaces.GPU
     def __init__(
         self,
         search_model: str = "internlm/internlm2_5-7b-chat",
@@ -289,19 +288,13 @@ class VisionSearchAssistant:
         vlm_load_4bit: bool = True,
         vlm_load_8bit: bool = False,
     ):
-        self.searcher = WebSearcher(
-            model_path = search_model
-        )
-        self.grounder = VisualGrounder(
-            model_path = ground_model,
-            device = ground_device,
-        )
-        self.vlm = VLM(
-            model_path = vlm_model,
-            device = vlm_device,
-            load_4bit = vlm_load_4bit,
-            load_8bit = vlm_load_8bit
-        )
+        self.search_model = search_model
+        self.ground_model = ground_model
+        self.ground_device = ground_device
+        self.vlm_model = vlm_model
+        self.vlm_device = vlm_device
+        self.vlm_load_4bit = vlm_load_4bit
+        self.vlm_load_8bit = vlm_load_8bit
         self.use_correlate = True
     
     @spaces.GPU
@@ -310,7 +303,21 @@ class VisionSearchAssistant:
         image: Union[str, Image.Image, np.ndarray],
         text: str,
         ground_classes: Union[List[str], None] = None
-    ):
+    ):  
+        self.searcher = WebSearcher(
+            model_path = self.search_model
+        )
+        self.grounder = VisualGrounder(
+            model_path = self.ground_model,
+            device = self.ground_device,
+        )
+        self.vlm = VLM(
+            model_path = self.vlm_model,
+            device = self.vlm_device,
+            load_4bit = self.vlm_load_4bit,
+            load_8bit = self.vlm_load_8bit
+        )
+
         # Create and clear the temporary directory.
         if not os.access('temp', os.F_OK):
             os.makedirs('temp')