Update app.py

app.py

@@ -1,12 +1,19 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import gradio as gr
+import spaces
+import torch
+
 
 tokenizer = AutoTokenizer.from_pretrained("merve/chatgpt-prompt-generator-v12")
 model = AutoModelForSeq2SeqLM.from_pretrained("merve/chatgpt-prompt-generator-v12", from_tf=True)
 
-def generate(prompt):
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
+@spaces.GPU
+def generate(prompt):
     batch = tokenizer(prompt, return_tensors="pt")
+    batch.to(device)
     generated_ids = model.generate(batch["input_ids"], max_new_tokens=150)
     output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
     return output[0]