Update server.py

server.py

@@ -9,13 +9,23 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from huggingface_hub import snapshot_download
 
-# Setup Logging
+# 1. SETUP LOGGING
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# 2. DEFINE PATHS
+DATASET_REPO = "aniketkumar1106/orbit-data" 
+IMAGE_DIR = "Productimages"
+DB_TARGET = "orbiitt.db" 
+
+# 3. CRITICAL FIX: Create directory BEFORE FastAPI initialization
+if not os.path.exists(IMAGE_DIR):
+    os.makedirs(IMAGE_DIR, exist_ok=True)
+    logger.info(f"Created initial directory: {IMAGE_DIR}")
+
+# 4. INITIALIZE APP
 app = FastAPI()
 
-# Enable CORS for Hugging Face environment
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -24,19 +34,16 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-DATASET_REPO = "aniketkumar1106/orbit-data" 
-IMAGE_DIR = "Productimages"
-DB_TARGET = "orbiitt.db" 
+# 5. MOUNT STORAGE (Will no longer crash)
+app.mount("/Productimages", StaticFiles(directory=IMAGE_DIR), name="Productimages")
 
-# Global state
+# 6. GLOBAL ENGINE STATE
 engine = None
-loading_status = "System Initializing..."
+loading_status = "System Booting..."
 
 def normalize_filename(name):
-    """Handles smart quotes, Unicode characters, and URL-unsafe symbols."""
-    # Convert curly quotes/dashes to standard ASCII
+    """Handles smart quotes and complex symbols to prevent 404s."""
     name = name.replace('’', "'").replace('‘', "'").replace('“', '"').replace('”', '"').replace('–', '-')
-    # Decompose Unicode characters and strip non-ASCII
     name = unicodedata.normalize('NFKD', name).encode('ascii', 'ignore').decode('ascii')
     return name.strip()
 
@@ -45,11 +52,14 @@ def setup_environment():
     global engine, loading_status
     token = os.environ.get("HF_TOKEN")
     
-    if os.path.exists(IMAGE_DIR): shutil.rmtree(IMAGE_DIR)
-    os.makedirs(IMAGE_DIR)
+    # We clean and re-create inside startup to ensure a fresh sync
+    if os.path.exists(IMAGE_DIR): 
+        shutil.rmtree(IMAGE_DIR)
+    os.makedirs(IMAGE_DIR, exist_ok=True)
 
     try:
         loading_status = "Syncing Orbit Assets..."
+        logger.info(loading_status)
         snapshot_download(repo_id=DATASET_REPO, repo_type="dataset", token=token, local_dir=".")
 
         if os.path.exists("orbiitt_db.zip"):
@@ -62,12 +72,13 @@ def setup_environment():
                     if f.lower() in ["orbiitt_db", "orbiitt.db"] or (f.endswith(".db") and not f.startswith(".")):
                         shutil.copy(src, f"./{DB_TARGET}")
                     elif f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
-                        # Clean name to prevent 404 mismatches
+                        # Normalize filename on disk
                         clean_name = normalize_filename(urllib.parse.unquote(f))
                         shutil.copy(src, os.path.join(IMAGE_DIR, clean_name))
             shutil.rmtree("temp_extract")
 
-        loading_status = "Loading SigLIP 2 AI Model..."
+        loading_status = "Waking up AI Engine..."
+        logger.info(loading_status)
         try:
             from orbiitt_engine import OrbiittEngine
             engine = OrbiittEngine()
@@ -79,8 +90,7 @@ def setup_environment():
 
     except Exception as e:
         loading_status = f"Setup Failed: {e}"
-
-app.mount("/Productimages", StaticFiles(directory=IMAGE_DIR), name="Productimages")
+        logger.error(loading_status)
 
 @app.get("/health")
 def health():
@@ -95,11 +105,11 @@ async def search(text: str = Form(None), weight: float = Form(0.5), file: Upload
     try:
         if file:
             t_path = f"tmp_{file.filename}"
-            with open(t_path, "wb") as b: b.write(await file.read())
+            with open(t_path, "wb") as b: 
+                b.write(await file.read())
         
         results = engine.search(text_query=text, image_file=t_path, text_weight=weight)
         
-        # Fuzzy match results against disk to ensure valid URLs
         all_files = os.listdir(IMAGE_DIR)
         valid_results = []
 
@@ -107,12 +117,13 @@ async def search(text: str = Form(None), weight: float = Form(0.5), file: Upload
             raw_path = r.get('url') or r.get('path') or ""
             fname = normalize_filename(os.path.basename(raw_path))
             
+            # Fuzzy match check
             match = None
             if fname in all_files:
                 match = fname
             else:
                 for disk_file in all_files:
-                    if fname[:20] in disk_file: # Substring fallback
+                    if fname[:20] in disk_file: 
                         match = disk_file
                         break
             
@@ -122,9 +133,11 @@ async def search(text: str = Form(None), weight: float = Form(0.5), file: Upload
                 
         return {"results": valid_results}
     except Exception as e:
+        logger.error(f"Search error: {e}")
         return {"results": [], "error": str(e)}
     finally:
-        if t_path and os.path.exists(t_path): os.remove(t_path)
+        if t_path and os.path.exists(t_path): 
+            os.remove(t_path)
 
 if __name__ == "__main__":
     import uvicorn