Update app.py

app.py

@@ -1,3 +1,5 @@
+import os
+
 import gradio as gr
 from pdf2image import convert_from_path
 
@@ -6,6 +8,7 @@ from PIL import Image
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 from transformers import AutoProcessor
+imoort os
 
 from custom_colbert.models.paligemma_colbert_architecture import ColPali
 from custom_colbert.trainer.retrieval_evaluator import CustomEvaluator
@@ -87,9 +90,10 @@ def index(file):
 COLORS = ['#4285f4', '#db4437', '#f4b400', '#0f9d58', '#e48ef1']
 # Load model
 model_name = "coldoc/colpali-3b-mix-448"
-model = ColPali.from_pretrained("google/paligemma-3b-mix-448", torch_dtype=torch.bfloat16, device_map="cuda").eval()
+token = os.environ.get("HF_TOKEN")
+model = ColPali.from_pretrained("google/paligemma-3b-mix-448", torch_dtype=torch.bfloat16, device_map="cuda", token=token).eval()
 model.load_adapter(model_name)
-processor = AutoProcessor.from_pretrained(model_name)
+processor = AutoProcessor.from_pretrained(model_name, token=token)
 device = model.device
 mock_image = Image.new("RGB", (448, 448), (255, 255, 255))