update model_worker.py

model_worker.py

@@ -325,7 +325,8 @@ class ModelWorker:
             "queue_length": self.get_queue_length(),
         }
 
-    @torch.inference_mode()
+    # @torch.inference_mode()
+    @spaces.GPU(duration=120)
     def generate_stream(self, params):
         system_message = params["prompt"][0]["content"]
         send_messages = params["prompt"][1:]
@@ -428,18 +429,19 @@ class ModelWorker:
         )
         logger.info(f"Generation config: {generation_config}")
 
-        thread = Thread(
-            target=self.model.chat,
-            kwargs=dict(
-                tokenizer=self.tokenizer,
-                pixel_values=pixel_values,
-                question=question,
-                history=history,
-                return_history=False,
-                generation_config=generation_config,
-            ),
-        )
-        thread.start()
+        with torch.no_grad():
+            thread = Thread(
+                target=self.model.chat,
+                kwargs=dict(
+                    tokenizer=self.tokenizer,
+                    pixel_values=pixel_values,
+                    question=question,
+                    history=history,
+                    return_history=False,
+                    generation_config=generation_config,
+                ),
+            )
+            thread.start()
 
         generated_text = ""
         for new_text in streamer:
@@ -453,7 +455,6 @@ class ModelWorker:
         )
         self.model.system_message = old_system_message
 
-    @spaces.GPU(duration=120)
     def generate_stream_gate(self, params):
         try:
             for x in self.generate_stream(params):