Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -23,6 +24,7 @@ def load_model(name):
         cached[name] = (tok, mod)
     return cached[name]
 
+@spaces.GPU()
 def chat_fn(message, history, selected_model):
     tokenizer, model = load_model(selected_model)
     # Build chat template (single‑turn for simplicity)
@@ -49,6 +51,7 @@ with gr.Blocks(title="Multi‑Model Chat") as demo:
             label="Choose Model"
         )
     chat = gr.ChatInterface(
+        fn=chat_gpu,
         fn=lambda msg, hist: chat_fn(msg, hist, model_select.value),
         additional_inputs=[model_select],
     )