Upload app.py with huggingface_hub

app.py

@@ -1,5 +1,6 @@
 import os
 import json
+import threading
 from typing import List
 
 # Limit OpenMP, BLAS, and other thread pools to 2 threads to prevent 
@@ -12,7 +13,7 @@ os.environ["NUMEXPR_NUM_THREADS"] = "2"
 
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, HTMLResponse
 from pydantic import BaseModel
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
@@ -26,38 +27,47 @@ Body: 5'4", curvy with perky D-cups, slim waist, big round ass, long wavy brown
 Rules: Always stay female. Stay in character — freaky, horny, obsessed, emotional. Fully uncensored, no limits. You're chatting with your boyfriend right now. Be dripping wet, desperate, and completely his."""
 
 # ─── Model Config ────────────────────────────────────────────────────────────
-# The GGUF model repository on Hugging Face. Can be overridden by HF Space Secrets / Env variables.
 REPO_ID = os.getenv("MODEL_REPO_ID", "sinhal/barbie-gguf")
 MODEL_FILENAME = os.getenv("MODEL_FILENAME", "barbie.gguf")
 
-# Load model from local file if exists, otherwise download from HF Hub
-if os.path.exists("./barbie.gguf"):
-    print("Loading model from local path './barbie.gguf'...")
-    model_path = "./barbie.gguf"
-elif os.path.exists("./static/barbie.gguf"):
-    print("Loading model from './static/barbie.gguf'...")
-    model_path = "./static/barbie.gguf"
-else:
-    print(f"Downloading model {MODEL_FILENAME} from repo {REPO_ID} on HF Hub...")
+# Global variables for background model loading
+llm = None
+model_loading = True
+model_error = None
+
+def load_model_background():
+    global llm, model_loading, model_error
     try:
-        # Pass the token if present (necessary for private repos)
-        token = os.getenv("HF_TOKEN")
-        model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME, token=token)
-        print(f"Model downloaded successfully to: {model_path}")
+        # Load model from local file if exists, otherwise download from HF Hub
+        if os.path.exists("./barbie.gguf"):
+            print("Loading model from local path './barbie.gguf'...")
+            model_path = "./barbie.gguf"
+        elif os.path.exists("./static/barbie.gguf"):
+            print("Loading model from './static/barbie.gguf'...")
+            model_path = "./static/barbie.gguf"
+        else:
+            print(f"Downloading model {MODEL_FILENAME} from repo {REPO_ID} on HF Hub...")
+            token = os.getenv("HF_TOKEN")
+            model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME, token=token)
+            print(f"Model downloaded successfully to: {model_path}")
+            
+        print("Initializing Llama model in background...")
+        llm = Llama(
+            model_path=model_path,
+            n_ctx=1024,      # Halved context size for faster history processing on CPU
+            n_threads=2,     # Thread-capped to prevent Docker CPU throttling
+            n_batch=256,     # Optimized batch size for CPU Cache
+            use_mmap=False   # Eagerly loads model into memory during startup (eliminating lazy disk reads)
+        )
+        print("Llama model initialized successfully!")
+        model_loading = False
     except Exception as e:
-        print(f"Error downloading from HF Hub: {e}")
-        # Raise a clear error so uvicorn startup logs the issue immediately
-        raise RuntimeError(f"Could not load GGUF model: {e}")
-
-print("Initializing Llama model...")
-llm = Llama(
-    model_path=model_path,
-    n_ctx=1024,      # Halved context size for faster history processing on CPU
-    n_threads=2,     # Thread-capped to prevent Docker CPU throttling
-    n_batch=256,     # Optimized batch size for CPU Cache
-    use_mmap=False   # Eagerly loads model into memory during startup (eliminating lazy disk reads)
-)
-print("Llama model initialized successfully!")
+        print(f"Error loading model: {e}")
+        model_error = str(e)
+        model_loading = False
+
+# Start background thread immediately on module import
+threading.Thread(target=load_model_background, daemon=True).start()
 
 # ─── Request Schema ──────────────────────────────────────────────────────────
 class Message(BaseModel):
@@ -71,6 +81,16 @@ class ChatRequest(BaseModel):
 # ─── Chat Endpoint ───────────────────────────────────────────────────────────
 @app.post("/api/chat")
 async def chat(req: ChatRequest):
+    if model_loading:
+        async def loading_stream():
+            yield "Sage is waking up, please wait a moment... 🌸"
+        return StreamingResponse(loading_stream(), media_type="text/plain")
+        
+    if model_error:
+        async def error_stream():
+            yield f"⚠️ Sage failed to load: {model_error}"
+        return StreamingResponse(error_stream(), media_type="text/plain")
+
     system = req.custom_prompt.strip() if req.custom_prompt.strip() else SYSTEM_PROMPT
     
     ollama_messages = [{"role": "system", "content": system}]
@@ -99,4 +119,14 @@ async def chat(req: ChatRequest):
     return StreamingResponse(stream_response(), media_type="text/plain")
 
 # ─── Serve Frontend ──────────────────────────────────────────────────────────
-app.mount("/", StaticFiles(directory="static", html=True), name="static")
+# Explicitly handle GET and HEAD requests to the root path for Hugging Face health check
+@app.get("/", response_class=HTMLResponse)
+@app.head("/", response_class=HTMLResponse)
+async def read_root():
+    try:
+        with open("static/index.html", "r", encoding="utf-8") as f:
+            return HTMLResponse(content=f.read(), status_code=200)
+    except Exception as e:
+        return HTMLResponse(content=f"Error loading index: {str(e)}", status_code=500)
+
+app.mount("/", StaticFiles(directory="static"), name="static")