vdefrv

app.py

@@ -133,6 +133,8 @@ def search(query: str, ds, images, k):
     results = []
     for idx in top_k_indices:
         results.append((images[idx])) #, f"Page {idx}"
+    del model
+    del processor
     print("done")
     return results
 
@@ -158,7 +160,17 @@ def convert_files(files):
 @spaces.GPU
 def index_gpu(images, ds):
     """Example script to run inference with ColPali"""
-    
+        # Load colpali model
+    model_name = "vidore/colpali-v1.2"
+    token = os.environ.get("HF_TOKEN")
+    model = ColPali.from_pretrained(
+        "vidore/colpaligemma-3b-pt-448-base", torch_dtype=torch.bfloat16, device_map="cuda", token = token).eval()
+
+    model.load_adapter(model_name)
+    model = model.eval()
+    processor = AutoProcessor.from_pretrained(model_name, token = token)
+
+    mock_image = Image.new("RGB", (448, 448), (255, 255, 255))
     # run inference - docs
     dataloader = DataLoader(
         images,
@@ -178,6 +190,9 @@ def index_gpu(images, ds):
             batch_doc = {k: v.to(device) for k, v in batch_doc.items()}
             embeddings_doc = model(**batch_doc)
         ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
+    del model
+    del processor
+    print("done")
     return f"Uploaded and converted {len(images)} pages", ds, images
 
 @spaces.GPU