Update app.py

app.py

@@ -1,27 +1,76 @@
 import spaces
 import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load your model
+# Load tokenizer globally (CPU operation)
 model_name = "loocorez/reverse-text-warmup"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-model.to("cuda")
 
+# Load model globally but keep on CPU initially
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16  # Use half precision for memory efficiency
+)
+
+@spaces.GPU(duration=60)  # Reserve GPU for 60 seconds
 def reverse_text(input_text):
-    # Add your model inference logic here
-    inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=100)
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return result
+    # Move model to GPU only when needed
+    model.to("cuda")
+    
+    try:
+        # Tokenize and move to GPU
+        inputs = tokenizer(
+            input_text, 
+            return_tensors="pt",
+            max_length=512,
+            truncation=True
+        ).to("cuda")
+        
+        # Generate
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=100,
+                do_sample=True,
+                temperature=0.7,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode result
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        generated_text = result[len(input_text):].strip()
+        
+        return generated_text
+        
+    except Exception as e:
+        return f"Error: {str(e)}"
+    
+    finally:
+        # Move model back to CPU to free GPU memory
+        model.to("cpu")
+        torch.cuda.empty_cache()
 
-# Create Gradio interface
+# Create interface
 demo = gr.Interface(
     fn=reverse_text,
-    inputs=gr.Textbox(label="Input Text"),
-    outputs=gr.Textbox(label="Reversed Text"),
-    title="Reverse Text Model Demo",
-    description="Test the reverse-text-warmup model"
+    inputs=gr.Textbox(
+        label="Input Text",
+        placeholder="Enter text to process...",
+        lines=3
+    ),
+    outputs=gr.Textbox(
+        label="Generated Text",
+        lines=3
+    ),
+    title="🔄 Reverse Text Model Demo",
+    description="Test your custom reverse-text-warmup model using ZeroGPU",
+    examples=[
+        ["Hello world"],
+        ["The quick brown fox"],
+        ["Machine learning is amazing"]
+    ]
 )
 
 demo.launch()