Update app.py

app.py

@@ -4,9 +4,6 @@ import torch
 from transformers import pipeline
 
 
-@spaces.GPU
-def warm_up(): 
-    pass
 
 # Set up the model
 model_id = "meta-llama/Llama-3.2-1B"
@@ -14,14 +11,13 @@ pipe = pipeline(
     "text-generation", 
     model=model_id, 
     torch_dtype=torch.bfloat16, 
-    device_map="cpu"
+    device_map="cuda"
 )
 
 
 
-# Text generation function
+@spaces.GPU
 def generate_text(prompt):
-    pipe.model.to('cuda')
     response = pipe(prompt, max_length=100, num_return_sequences=1)
     return response[0]['generated_text']
 
@@ -30,7 +26,7 @@ css = """
 body {background-color: #f0f8ff; font-family: 'Arial', sans-serif;}
 .gradio-container {max-width: 700px; margin: 0 auto; border-radius: 15px; padding: 20px; background-color: white; box-shadow: 0 4px 8px rgba(0,0,0,0.1);}
 textarea {border-radius: 10px; padding: 10px; font-size: 16px; border: 2px solid #ddd;}
-button {background-color: #4caf50; color: white; border-radius: 5px; padding: 10px 20px; font-size: 16px; border: none; cursor: pointer;}
+button {background-color: #4caf50; color: #000000; border-radius: 5px; padding: 10px 20px; font-size: 16px; border: none; cursor: pointer;}
 button:hover {background-color: #45a049;}
 h1 {color: #333; font-size: 32px; text-align: center;}
 footer {text-align: center; padding: 10px;}