Update app.py

app.py

@@ -3,32 +3,49 @@ import json
 import uuid
 import httpx
 import gradio as gr
+import torch
 from fastapi import FastAPI, HTTPException, Request
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, set_default_dtype
 import uvicorn
 import asyncio
 
-# ✅ Load Model Configuration
+# ✅ Use float16 to reduce memory usage (for Hugging Face Spaces)
+set_default_dtype(torch.float16)
+
+# ✅ Hugging Face API Token
 HF_TOKEN = os.getenv("HF_TOKEN")
 MODEL_NAME = "hpyapali/tinyllama-workout"
-event_store = {}  
+event_store = {}  # Store AI responses for polling fallback
 
 app = FastAPI()
 
-# ✅ Load AI Model
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, token=HF_TOKEN)
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    print("✅ AI Model Loaded Successfully!")
-except Exception as e:
-    print(f"❌ Error loading model: {e}")
-    pipe = None
+# ✅ Lazy Load AI Model (to prevent Space timeout)
+pipe = None
+
+def get_pipeline():
+    global pipe
+    if pipe is None:
+        try:
+            print("🔄 Loading AI Model...")
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                token=HF_TOKEN,
+                torch_dtype=torch.float16,   # Lower memory usage
+                device_map="auto"            # Load on available device (CPU/GPU)
+            )
+            pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+            print("✅ AI Model Loaded Successfully!")
+        except Exception as e:
+            print(f"❌ Error loading model: {e}")
+            pipe = None
+    return pipe
 
 
 # ✅ AI Function - Processes and ranks workouts
 def analyze_workouts(last_workouts: str):
     """Generates AI-based workout rankings based on heart rate recovery."""
+    pipe = get_pipeline()
     if pipe is None:
         return "❌ AI model is not loaded."
 
@@ -101,6 +118,25 @@ async def root():
     return {"message": "Workout Analysis & Ranking AI is running!"}
 
 
-# ✅ Start FastAPI
-if __name__ == "__main__":
+# ✅ Gradio UI for Testing
+iface = gr.Interface(
+    fn=analyze_workouts,
+    inputs="text",
+    outputs="text",
+    title="Workout Analysis & Ranking AI",
+    description="Enter workout data to analyze effectiveness, rank workouts, and receive improvement recommendations."
+)
+
+
+# ✅ Start Both FastAPI & Gradio
+def start_gradio():
+    iface.launch(server_name="0.0.0.0", server_port=7860, share=True)
+
+def start_fastapi():
     uvicorn.run(app, host="0.0.0.0", port=7861)
+
+# ✅ Run both servers in parallel
+if __name__ == "__main__":
+    import threading
+    threading.Thread(target=start_gradio).start()
+    threading.Thread(target=start_fastapi).start()