update model_worker.py

model_worker.py

@@ -228,7 +228,7 @@ class ModelWorker:
         )
         self.heart_beat_thread.start()
 
-    @spaces.GPU
+    @spaces.GPU(duration=120)
     def import_flash_attn(self):
         try:
             import flash_attn
@@ -325,10 +325,8 @@ class ModelWorker:
             "queue_length": self.get_queue_length(),
         }
 
-    @spaces.GPU
     @torch.inference_mode()
     def generate_stream(self, params):
-
         system_message = params["prompt"][0]["content"]
         send_messages = params["prompt"][1:]
         max_input_tiles = params["max_input_tiles"]
@@ -455,6 +453,7 @@ 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):