Update app.py

app.py

@@ -1,6 +1,7 @@
 import base64
 import gradio as gr
 import torch
+import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 MODEL_ID = "ApolloAIs/phoebus-3.1-1.5b-merged"
@@ -9,17 +10,17 @@ print("Loading Phoebus 3.1...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float32,
-    device_map="cpu"
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
-print("Model loaded!")
+print("Model ready!")
 
 def format_emem_cell64(claim: str) -> str:
     encoded = base64.b64encode(claim.encode("utf-8")).decode("utf-8")[:12]
     return f"emem:fact:cell64_{encoded}"
 
+@spaces.GPU
 def phoebus_respond(message, history, enable_tinyfish, enable_emem):
-    # Base system persona
     system_prompt = (
         "You are Phoebus 3.1, a reasoning model specialized in deep space telemetry, "
         "Earth geography, and coding, developed by Clovr (atclovr.dev). "
@@ -34,7 +35,6 @@ def phoebus_respond(message, history, enable_tinyfish, enable_emem):
         receipt = format_emem_cell64(message)
         context += f"\n\n[System Protocol - Vortx AI emem]: Attach receipt `{receipt}` for physical spatial validation."
 
-    # Format messages for the model
     formatted_messages = [{"role": "system", "content": system_prompt}]
     for user_turn, bot_turn in history:
         formatted_messages.append({"role": "user", "content": user_turn})
@@ -48,7 +48,7 @@ def phoebus_respond(message, history, enable_tinyfish, enable_emem):
         add_generation_prompt=True
     )
 
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -61,7 +61,6 @@ def phoebus_respond(message, history, enable_tinyfish, enable_emem):
     reply = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return reply
 
-# Create the ChatGPT-like UI
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple")) as demo:
     gr.Markdown("# 🛰️ Phoebus 3.1 Console\n*Reasoning model by Clovr (atclovr.dev)*")