Update app.py

app.py

@@ -6,6 +6,7 @@ import gradio as gr
 from transformers import AutoProcessor, AutoModelForCausalLM
 import torch
 import numpy as np
+import spaces
 import subprocess
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
@@ -33,6 +34,7 @@ def preprocess_image(image_path, max_size=(800, 800)):
     return image_np, image.size
 
 # Function to run Florence-2-large model
+@spaces.GPU
 def run_florence_model(image_np, image_size, task_prompt, text_input=None):
     if text_input is None:
         prompt = task_prompt
@@ -62,11 +64,9 @@ def run_florence_model(image_np, image_size, task_prompt, text_input=None):
 
 # Function to plot image with bounding boxes
 def plot_image_with_bboxes(image_np, bboxes, labels=None):
-    # Normalize the image array to [0.0, 1.0]
-    if image_np.dtype == np.uint8:
-        image_np = image_np / 255.0
-
     fig, ax = plt.subplots(1)
+    # Normalize the image data to the range [0, 1]
+    image_np = image_np / 255.0
     ax.imshow(image_np)
     colors = ['red', 'blue', 'green', 'yellow', 'purple', 'cyan']
     for i, bbox in enumerate(bboxes):
@@ -80,9 +80,10 @@ def plot_image_with_bboxes(image_np, bboxes, labels=None):
     return fig
 
 # Gradio function to process uploaded images
+@spaces.GPU
 def process_image(image_path):
     image_np, image_size = preprocess_image(image_path)
-    
+
     # Convert image_np to float32
     image_np = image_np.astype(np.float32)