Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from PIL import Image
 import torch
 from transformers import AutoProcessor, LlavaNextForConditionalGeneration
+import spaces
 
 # Load the processor and model
 model_id = "llava-hf/llava-v1.6-mistral-7B-hf"
@@ -12,6 +13,7 @@ model = LlavaNextForConditionalGeneration.from_pretrained(
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+@spaces.GPU()
 def llava_inference(image: Image.Image, prompt: str):
     # Format the input as a conversation
     conversation = [