lazy loading of model

api_backend.py

@@ -19,24 +19,61 @@ import json
 import base64
 import os
 
-# 🔧 Import local processing modules
-from yolo_module import run_yolo
-from ocr_module import extract_text, count_elements, validate_structure
-from graph_module import map_arrows, build_flowchart_json
-from summarizer_module import summarize_flowchart
-
 # 🔥 Initialize FastAPI app
 app = FastAPI()
 
-# 🔓 Enable CORS to allow frontend (e.g., Streamlit on localhost) to connect
+# 🔓 Enable CORS with more specific configuration for Hugging Face Spaces
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # In production, replace with allowed frontend domain
+    allow_origins=["*", "https://venkatviswa-flowchart-to-text.hf.space"],  # Include your specific domain
     allow_credentials=True,
-    allow_methods=["*"],
+    allow_methods=["GET", "POST", "OPTIONS"],  # Explicitly allow methods
     allow_headers=["*"],
+    expose_headers=["*"],
 )
 
+# Add a health check endpoint
+@app.get("/")
+async def health_check():
+    """Health check endpoint to verify API is running."""
+    return {"status": "ok", "message": "API is running"}
+
+
+@app.options("/process-image")
+async def options_process_image():
+    """Handle OPTIONS requests for the process-image endpoint."""
+    return {}
+
+
+# Lazy-loading for ML modules to avoid startup issues
+SKIP_MODEL_LOADING = os.getenv("SKIP_MODEL_LOADING", "0") == "1"
+yolo_module = None
+ocr_module = None
+graph_module = None
+summarizer_module = None
+
+def load_modules():
+    global yolo_module, ocr_module, graph_module, summarizer_module
+    
+    if yolo_module is None:
+        # Only import these when needed, not during startup
+        from yolo_module import run_yolo as yolo_run
+        from ocr_module import extract_text as ocr_extract, count_elements, validate_structure
+        from graph_module import map_arrows, build_flowchart_json
+        from summarizer_module import summarize_flowchart
+        
+        yolo_module = {"run_yolo": yolo_run}
+        ocr_module = {
+            "extract_text": ocr_extract,
+            "count_elements": count_elements,
+            "validate_structure": validate_structure
+        }
+        graph_module = {
+            "map_arrows": map_arrows,
+            "build_flowchart_json": build_flowchart_json
+        }
+        summarizer_module = {"summarize_flowchart": summarize_flowchart}
+
 
 @app.post("/process-image")
 async def process_image(
@@ -54,6 +91,9 @@ async def process_image(
     Returns:
         JSONResponse: Contains flowchart structure, summary, debug output, and optional YOLO overlay.
     """
+    # Lazy load modules when first request comes in
+    load_modules()
+    
     debug_mode = debug.lower() == "true"
     debug_log = []
 
@@ -69,25 +109,24 @@ async def process_image(
     print("✅ Image converted to RGB")
 
     # 📦 YOLO Detection for boxes and arrows
-    boxes, arrows, vis_debug = run_yolo(image)
+    boxes, arrows, vis_debug = yolo_module["run_yolo"](image)
     if debug_mode:
         debug_log.append(f"📦 Detected {len(boxes)} boxes, {len(arrows)} arrows")
 
     # 🔍 Run OCR on each detected box
     for box in boxes:
-        box["text"] = extract_text(image, box["bbox"], debug=debug_mode)
+        box["text"] = ocr_module["extract_text"](image, box["bbox"], debug=debug_mode)
         print(f"🔍 OCR for {box['id']}: {box['text']}")
         if debug_mode:
             debug_log.append(f"🔍 {box['id']}: {box['text']}")
 
-
     # 🧠 Build structured JSON from nodes and edges
-    flowchart_json = build_flowchart_json(boxes, arrows)
+    flowchart_json = graph_module["build_flowchart_json"](boxes, arrows)
     print("🧠 Flowchart JSON:", json.dumps(flowchart_json, indent=2))
 
     # ✅ Validate structure
-    structure_info = count_elements(boxes, arrows, debug=debug_mode)
-    validation = validate_structure(
+    structure_info = ocr_module["count_elements"](boxes, arrows, debug=debug_mode)
+    validation = ocr_module["validate_structure"](
         flowchart_json,
         expected_boxes=structure_info["box_count"],
         expected_arrows=len(arrows),
@@ -97,7 +136,7 @@ async def process_image(
         debug_log.append(f"🧾 Validation: {validation}")
 
     # ✍️ Generate plain-English summary
-    summary = summarize_flowchart(flowchart_json)
+    summary = summarizer_module["summarize_flowchart"](flowchart_json)
     print("📝 Summary:", summary)
 
     # 🖼️ Encode YOLO debug image (if debug enabled)

combined_app.py

@@ -1,10 +1,7 @@
 """
 combined_app.py - A unified approach that runs both FastAPI and Streamlit in a single process.
 This is specifically designed for Hugging Face Spaces deployments.
-
-Usage:
-1. Make this your entry point in Hugging Face Spaces
-2. It will start the FastAPI backend and then launch Streamlit in the same process
+Uses lazy loading to avoid loading models during startup.
 """
 
 import streamlit as st
@@ -20,7 +17,11 @@ import subprocess
 # Use environment variable to determine if we're in Spaces
 IS_SPACE = "SPACE_ID" in os.environ
 
-# Import FastAPI app from your backend module
+# Import FastAPI app but avoid importing modules that load models
+# This is critical to prevent infinite loops during deployment
+os.environ["SKIP_MODEL_LOADING"] = "1"
+
+# Now import the FastAPI app
 from api_backend import app as fastapi_app
 
 # Add a route for checking if the API is running