Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
+import spaces
 
 def load_model(model_name):
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -21,6 +22,7 @@ def load_model(model_name):
     )
     return generator
 
+@spaces.GPU
 def generate_text(prompt, model_name):
     generator = load_model(model_name)
     messages = [{"role": "user", "content": prompt}]
@@ -48,3 +50,4 @@ demo = gr.Interface(
 )
 
 demo.launch()
+