Update app.py

app.py

@@ -1,3 +1,4 @@
+import torchimport spaces
 import os
 import uuid
 import time
@@ -14,7 +15,12 @@ from huggingface_hub import HfApi, login
 
 # Model configuration
 checkpoint = "WillHeld/soft-raccoon"
-device = "cuda"
+# Set device based on availability
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+    print("CUDA not available, using CPU instead. This will be much slower.")
 
 # Dataset configuration
 DATASET_NAME = "your-username/soft-raccoon-conversations"  # Change to your username
@@ -71,6 +77,7 @@ def save_to_dataset():
     return dataset, status_msg
 
 
+@spaces.GPU(duration=120)
 def predict(message, chat_history, temperature, top_p, conversation_id=None):
     """Generate a response using the model and save the conversation"""
     # Create/retrieve conversation ID for tracking
@@ -203,8 +210,8 @@ with gr.Blocks(theme=theme, title="Stanford Soft Raccoon Chat with Dataset Colle
     with gr.Row():
         with gr.Column(scale=3):
             chatbot = gr.Chatbot(
-                label="Soft Raccoon Chat",
-                avatar_images=(None, "🦝"),
+                label="Stanford Soft Raccoon Chat",
+                avatar_images=(None, "🌲"),  # Stanford tree emoji
                 height=600
             )
             
@@ -303,5 +310,4 @@ with gr.Blocks(theme=theme, title="Stanford Soft Raccoon Chat with Dataset Colle
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch(share=True)
-    
+    demo.launch(share=True)