Update demo/agent.py

demo/agent.py

@@ -11,10 +11,17 @@ from lmdeploy.vl import load_image
 class ConversationalAgent:
     def __init__(self, 
                  model_path,
-                 outputs_dir) -> None:
-        self.pipe = pipeline(model_path,
-                        chat_template_config=ChatTemplateConfig(model_name='internvl2-internlm2'),
-                        backend_config=TurbomindEngineConfig(session_len=8192))
+                 outputs_dir,
+                 device='cpu') -> None:
+        # 传入 device 参数，并设置 pipeline 时指定设备
+        self.device = device
+        self.pipe = pipeline(
+            model_path,
+            chat_template_config=ChatTemplateConfig(model_name='internvl2-internlm2'),
+            backend_config=TurbomindEngineConfig(session_len=8192),
+            device=self.device  # 设置设备为 'cpu' 或 'cuda'
+        )
+        
         self.uploaded_images_storage = os.path.join(outputs_dir, "uploaded")
         self.uploaded_images_storage = os.path.abspath(self.uploaded_images_storage)
         os.makedirs(self.uploaded_images_storage, exist_ok=True)