update

app.py

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

models/vsa_model.py

@@ -213,21 +213,7 @@ class WebSearcher:
         backend_config = PytorchEngineConfig(
             max_batch_size = 1,
         )
-        # llm = LMDeployServer(
-        #     path = model_path,
-        #     model_name = model_name,
-        #     meta_template = INTERNLM2_META,
-        #     top_p = top_p,
-        #     top_k = top_k,
-        #     temperature = temperature,
-        #     max_new_tokens = max_new_tokens,
-        #     repetition_penalty = repetition_penalty,
-        #     stop_words = ['<|im_end|>'],
-        #     serve_cfg = dict(
-        #         backend_config = backend_config
-        #     )
-        # )
-        llm = LMDeployPipeline(
+        llm = LMDeployServer(
             path = model_path,
             model_name = model_name,
             meta_template = INTERNLM2_META,
@@ -237,7 +223,21 @@ class WebSearcher:
             max_new_tokens = max_new_tokens,
             repetition_penalty = repetition_penalty,
             stop_words = ['<|im_end|>'],
+            serve_cfg = dict(
+                backend_config = backend_config
+            )
         )
+        # llm = LMDeployPipeline(
+        #     path = model_path,
+        #     model_name = model_name,
+        #     meta_template = INTERNLM2_META,
+        #     top_p = top_p,
+        #     top_k = top_k,
+        #     temperature = temperature,
+        #     max_new_tokens = max_new_tokens,
+        #     repetition_penalty = repetition_penalty,
+        #     stop_words = ['<|im_end|>'],
+        # )
         self.agent = MindSearchAgent(
             llm = llm,
             protocol = MindSearchProtocol(