support ZeroGPU, fixed

app.py

@@ -1,4 +1,3 @@
-+import spaces
 import sys
 import os
 import argparse
@@ -25,7 +24,6 @@ def start_controller():
     print(controller_command)
     return subprocess.Popen(controller_command)
 
-[email protected]
 def start_worker(model_path: str, bits=16):
     print(f"Starting the model worker for the model {model_path}")
     model_name = model_path.strip("/").split("/")[-1]

llava/serve/model_worker.py

@@ -1,6 +1,7 @@
 """
 A model worker executes the model.
 """
+import spaces
 import argparse
 import asyncio
 import json
@@ -119,6 +120,7 @@ class ModelWorker:
             "queue_length": self.get_queue_length(),
         }
 
+    @spaces.GPU
     @torch.inference_mode()
     def generate_stream(self, params):
         tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor