Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import spaces
 
 # Load the fine-tuned model
 def load_model():
@@ -12,6 +13,7 @@ def load_model():
 model, tokenizer = load_model()
 
 # Function to generate an image
+@spaces.GPU
 def generate_image(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(inputs["input_ids"], max_length=50)