Refactor model loading and inferring

app.py

@@ -25,20 +25,14 @@ else:
 
     class spaces:
         @staticmethod
-        def GPU(func, duration: int = 60):
+        def GPU(func):
             def wrapper(*args, **kwargs):
                 return func(*args, **kwargs)
 
             return wrapper
 
 
-@spaces.GPU
-def dummy():
-    return
-
-
 # Define constants
-GPU_DURATION = 300
 EXAMPLES_DIR = Path(__file__).parent / "examples"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 MODEL_IDS = [
@@ -199,20 +193,20 @@ with gr.Blocks() as demo:
             ).eval()
             current_processor = AutoProcessor.from_pretrained(model_id)
             current_model_id = model_id
-
         return current_model, current_processor
 
-    def run(
+    @spaces.GPU
+    def generate(
+        model,
+        processor,
         image,
         model_id: str,
         system_prompt: str,
         user_prompt: str,
-        max_new_tokens: int = 1024,
-        image_resize: str = "Yes",
-        image_target_size: int | None = None,
+        max_new_tokens: int,
+        image_resize: str,
+        image_target_size: int | None,
     ):
-        model, processor = load_model(model_id)
-
         base64_image = image_to_base64(
             resize_image(image, image_target_size)
             if image_resize == "Yes" and image_target_size
@@ -247,8 +241,7 @@ with gr.Blocks() as demo:
         )
         inputs = inputs.to(DEVICE)
 
-        generate = spaces.GPU(model.generate, duration=GPU_DURATION)
-        generated_ids = generate(**inputs, max_new_tokens=max_new_tokens)
+        generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)
         generated_ids_trimmed = [
             out_ids[len(in_ids) :]
             for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
@@ -287,6 +280,29 @@ with gr.Blocks() as demo:
 
         return [(image, bboxes), str(json.dumps(output_json))]
 
+    def run(
+        image,
+        model_id: str,
+        system_prompt: str,
+        user_prompt: str,
+        max_new_tokens: int = 1024,
+        image_resize: str = "Yes",
+        image_target_size: int | None = None,
+    ):
+        model, processor = load_model(model_id)
+
+        return generate(
+            model,
+            processor,
+            image,
+            model_id,
+            system_prompt,
+            user_prompt,
+            max_new_tokens,
+            image_resize,
+            image_target_size,
+        )
+
     with gr.Row():
         with gr.Column():
             gr.Markdown("## Examples")