Update app.py

app.py

@@ -4,6 +4,7 @@ import safetensors
 import timm
 from transformers import AutoProcessor
 import gradio as gr
+import spaces
 import torch
 import time
 from florence2_implementation.modeling_florence2 import Florence2ForConditionalGeneration
@@ -239,6 +240,7 @@ for idx, tag in enumerate(allowed_tags):
 
 pruner = Pruner("tags-2024-05-05.csv")
 
+@spaces.GPU
 def generate_prompt(image, expected_caption_length):
     global THRESHOLD, tree, tokenizer, model, tagger_model, tagger_transform
     tagger_input = tagger_transform(image.convert('RGBA')).unsqueeze(0)
@@ -254,7 +256,7 @@ def generate_prompt(image, expected_caption_length):
     task_prompt = pruner.prompt_construction_pipeline_florence2(final_tags, expected_caption_length)
     return task_prompt
 
-
+@spaces.GPU
 def inference_caption(image, expected_caption_length, seq_len=512,):
     start_time = time.time()
     prompt_input = generate_prompt(image, expected_caption_length)