Update app.py

app.py

@@ -3,7 +3,8 @@ import spaces
 
 import gradio as gr
 import torch
-from colpali_engine.models.paligemma_colbert_architecture import ColPali
+from .modeling_colflor import ColFlor
+from .processing_colflor import ColFlorProcessor
 from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator
 from colpali_engine.utils.colpali_processing_utils import (
     process_images,
@@ -16,16 +17,15 @@ from tqdm import tqdm
 from transformers import AutoProcessor
 
 # Load model
-model_name = "vidore/colpali-v1.2"
+model_name = "ahmed-masry/ColFlor"
 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 = ColFlor.from_pretrained(
+    model_name, 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)
+processor = ColFlorProcessor.from_pretrained(model_name, token = token)
 
-mock_image = Image.new("RGB", (448, 448), (255, 255, 255))
+mock_image = Image.new("RGB", (768, 768), (255, 255, 255))
 
 
 @spaces.GPU