Add @spaces.GPU decorator for ZeroGPU support

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
@@ -18,40 +19,40 @@ MODELS = {
     }
 }
 
-# Global cache for loaded models
-loaded_models = {}
-
-def load_model(model_name):
-    """Load model and tokenizer, cache if already loaded"""
-    if model_name not in loaded_models:
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map="auto"
-        )
-        loaded_models[model_name] = (model, tokenizer)
-
-    return loaded_models[model_name]
+# Global storage for models (will be loaded on GPU)
+current_model = None
+current_tokenizer = None
+current_model_name = None
 
+@spaces.GPU
 def generate_text(model_name, prompt, max_tokens, temperature, top_p):
-    """Generate text using selected model"""
+    """Generate text using selected model with ZeroGPU"""
+    global current_model, current_tokenizer, current_model_name
+
     try:
-        model, tokenizer = load_model(model_name)
+        # Load model if not loaded or different model selected
+        if current_model is None or current_model_name != model_name:
+            current_tokenizer = AutoTokenizer.from_pretrained(model_name)
+            current_model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            current_model_name = model_name
 
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
 
         with torch.no_grad():
-            outputs = model.generate(
+            outputs = current_model.generate(
                 **inputs,
                 max_new_tokens=max_tokens,
                 temperature=temperature,
                 top_p=top_p,
                 do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
+                pad_token_id=current_tokenizer.eos_token_id
             )
 
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = current_tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response
 
     except Exception as e: