Enhance model loading in indexing functions for BiGemma3 and ColGemma3

app.py

@@ -143,6 +143,10 @@ def unload_models():
 @spaces.GPU
 def index_bigemma_images(images: List[Image.Image]) -> torch.Tensor:
     """Index images with BiGemma3 model."""
+    # Ensure model is loaded
+    if doc_index.bigemma_model is None:
+        load_bigemma_model()
+
     model, processor = doc_index.bigemma_model, doc_index.bigemma_processor
     batch_images = processor.process_images(images).to(device)
     embeddings = model(**batch_images, embedding_dim=768)
@@ -151,6 +155,10 @@ def index_bigemma_images(images: List[Image.Image]) -> torch.Tensor:
 @spaces.GPU
 def index_colgemma_images(images: List[Image.Image]) -> torch.Tensor:
     """Index images with ColGemma3 model."""
+    # Ensure model is loaded
+    if doc_index.colgemma_model is None:
+        load_colgemma_model()
+
     model, processor = doc_index.colgemma_model, doc_index.colgemma_processor
     batch_images = processor.process_images(images).to(device)
     embeddings = model(**batch_images)
@@ -175,27 +183,13 @@ def index_document(pdf_file, model_choice: str):
 
         # Index with BiGemma3
         if model_choice in ["NetraEmbed (BiGemma3)", "Both"]:
-            if doc_index.bigemma_model is None:
-                status.append("⏳ Loading BiGemma3 model...")
-                load_bigemma_model()
-                status.append("✓ BiGemma3 loaded")
-            else:
-                status.append("✓ Using cached BiGemma3 model")
-
-            status.append("⏳ Encoding images with BiGemma3...")
+            status.append("⏳ Loading & encoding with BiGemma3...")
             doc_index.bigemma_embeddings = index_bigemma_images(doc_index.images)
             status.append(f"✓ Indexed with BiGemma3 (shape: {doc_index.bigemma_embeddings.shape})")
 
         # Index with ColGemma3
         if model_choice in ["ColNetraEmbed (ColGemma3)", "Both"]:
-            if doc_index.colgemma_model is None:
-                status.append("⏳ Loading ColGemma3 model...")
-                load_colgemma_model()
-                status.append("✓ ColGemma3 loaded")
-            else:
-                status.append("✓ Using cached ColGemma3 model")
-
-            status.append("⏳ Encoding images with ColGemma3...")
+            status.append("⏳ Loading & encoding with ColGemma3...")
             doc_index.colgemma_embeddings = index_colgemma_images(doc_index.images)
             status.append(f"✓ Indexed with ColGemma3 (shape: {doc_index.colgemma_embeddings.shape})")
 
@@ -293,6 +287,10 @@ def generate_colgemma_heatmap(
 @spaces.GPU
 def query_bigemma(query: str, top_k: int) -> Tuple[str, List]:
     """Query indexed documents with BiGemma3."""
+    # Ensure model is loaded
+    if doc_index.bigemma_model is None:
+        load_bigemma_model()
+
     model, processor = doc_index.bigemma_model, doc_index.bigemma_processor
 
     # Encode query
@@ -322,6 +320,10 @@ def query_bigemma(query: str, top_k: int) -> Tuple[str, List]:
 @spaces.GPU
 def query_colgemma(query: str, top_k: int, show_heatmap: bool = False) -> Tuple[str, List]:
     """Query indexed documents with ColGemma3."""
+    # Ensure model is loaded
+    if doc_index.colgemma_model is None:
+        load_colgemma_model()
+
     model, processor = doc_index.colgemma_model, doc_index.colgemma_processor
 
     # Encode query