Update app.py

app.py

@@ -1,14 +1,19 @@
-from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 import gradio as gr
+import spaces
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
+import torch
+
+device = 'cuda' 
 
 model_name = "ailm/pegsus-text-summarization"
-model = PegasusForConditionalGeneration.from_pretrained(model_name)
+model = PegasusForConditionalGeneration.from_pretrained(model_name).to(device)
 tokenizer = PegasusTokenizer.from_pretrained(model_name)
 
+@spaces.GPU
 def summarize(text):
-    tokens = tokenizer(text, truncation=True, padding="longest", return_tensors="pt")
+    tokens = tokenizer(text, truncation=True, padding="longest", return_tensors="pt").to(device)
     summary = model.generate(**tokens)
     return tokenizer.decode(summary[0], skip_special_tokens=True)
 
 iface = gr.Interface(fn=summarize, inputs="text", outputs="text")
-iface.launch()
+iface.launch()