Fix startup timeout with lazy model loading

backend/model_handler.py

@@ -13,14 +13,24 @@ class ModelHandler:
             self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
             print(f"Using device: {self.device}", flush=True)
             self.transform = build_transform()
-            self.load_models()
+            
+            # Initialize model placeholders
+            self.model_int = None
+            self.tokenizer_int = None
+            self.reader = None
+            self.model_clip = None
+            self.processor_clip = None
+            
         except Exception as e:
             print(f"CRITICAL ERROR in ModelHandler.__init__: {e}", flush=True)
             import traceback
             traceback.print_exc()
 
-    def load_models(self):
-        # MODEL 1: InternVL
+    def load_internvl(self):
+        if self.model_int is not None and self.tokenizer_int is not None:
+            return
+
+        print("Loading InternVL model...", flush=True)
         try:
             # Check if local path exists, otherwise use HF Hub ID
             local_path = os.path.join("Models", "InternVL2_5-1B-MPO")
@@ -51,7 +61,11 @@ class ModelHandler:
             self.model_int = None
             self.tokenizer_int = None
 
-        # MODEL 2: EasyOCR
+    def load_easyocr(self):
+        if self.reader is not None:
+            return
+
+        print("Loading EasyOCR model...", flush=True)
         try:
             # EasyOCR automatically handles downloading if not present
             self.reader = easyocr.Reader(['en', 'hi'], gpu=False)
@@ -60,7 +74,11 @@ class ModelHandler:
             print(f"\nError initializing EasyOCR reader: {e}")
             self.reader = None
 
-        # MODEL 3: CLIP
+    def load_clip(self):
+        if self.model_clip is not None and self.processor_clip is not None:
+            return
+
+        print("Loading CLIP model...", flush=True)
         try:
             local_path = os.path.join("Models", "clip-vit-base-patch32")
             if os.path.exists(local_path):
@@ -79,6 +97,7 @@ class ModelHandler:
             self.processor_clip = None
 
     def easyocr_ocr(self, image):
+        self.load_easyocr()
         if not self.reader:
             return ""
         image_np = np.array(image)
@@ -95,6 +114,7 @@ class ModelHandler:
         return ordered_text
 
     def intern(self, image, prompt, max_tokens):
+        self.load_internvl()
         if not self.model_int or not self.tokenizer_int:
             return ""
             
@@ -123,6 +143,7 @@ class ModelHandler:
         return response
 
     def clip(self, image, labels):
+        self.load_clip()
         if not self.model_clip or not self.processor_clip:
             return None
 
@@ -138,6 +159,8 @@ class ModelHandler:
         return processed
 
     def get_clip_probs(self, image, labels):
+        # clip() calls load_clip(), so we don't strictly need it here, but good for safety if clip() implementation changes
+        self.load_clip() 
         inputs = self.clip(image, labels)
         if inputs is None:
             return None

verify_startup.py

@@ -0,0 +1,46 @@
+import sys
+import os
+
+# Add the current directory to sys.path so we can import backend
+sys.path.append(os.getcwd())
+
+print("Importing backend.model_handler...", flush=True)
+from backend.model_handler import model_handler
+
+print("\nChecking model states...", flush=True)
+
+errors = []
+
+if model_handler.model_int is not None:
+    errors.append("InternVL model should be None on startup")
+else:
+    print("PASS: InternVL model is None")
+
+if model_handler.tokenizer_int is not None:
+    errors.append("InternVL tokenizer should be None on startup")
+else:
+    print("PASS: InternVL tokenizer is None")
+
+if model_handler.reader is not None:
+    errors.append("EasyOCR reader should be None on startup")
+else:
+    print("PASS: EasyOCR reader is None")
+
+if model_handler.model_clip is not None:
+    errors.append("CLIP model should be None on startup")
+else:
+    print("PASS: CLIP model is None")
+
+if model_handler.processor_clip is not None:
+    errors.append("CLIP processor should be None on startup")
+else:
+    print("PASS: CLIP processor is None")
+
+if errors:
+    print("\nFAILED:", flush=True)
+    for error in errors:
+        print(f"- {error}", flush=True)
+    sys.exit(1)
+else:
+    print("\nSUCCESS: All models are lazily loaded.", flush=True)
+    sys.exit(0)