Update app.py

app.py

@@ -1,17 +1,18 @@
-import gradio as gr, transformers
-
-pipe = transformers.pipeline("text-generation",
-                             model="marcelbinz/Llama-3.1-Minitaur-8B")
-
-def infer(prompt, max_tokens):
-    return pipe(prompt, max_new_tokens=int(max_tokens))[0]["generated_text"]
-
-demo = gr.Interface(
-        fn=infer,
-        inputs=[gr.Textbox(label="Prompt"),
-                gr.Slider(5,128,32,step=1,label="Max tokens")],
-        outputs="text",
-        title="Minitaur",
-        description="Just type and hit *Run*"
-)
+import gradio as gr, transformers
+
+pipe = transformers.pipeline("text-generation",
+                             model="marcelbinz/Llama-3.1-Minitaur-8B")
+pipe.to('cuda')
+
+@spaces.GPU
+def infer(prompt, max_tokens):
+    return pipe(prompt, max_new_tokens=int(max_tokens))[0]["generated_text"]
+
+demo = gr.Interface(
+        fn=infer,
+        inputs=gr.Text(),
+        outputs="text",
+        title="Minitaur",
+        description="Just type and hit *Run*"
+)
 demo.launch()