updated

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from PIL import Image
 import torch
-from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, AutoModelForImageToText
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, pipeline
 from colpali_engine.models import ColPali, ColPaliProcessor
 import os
 
@@ -9,22 +9,21 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 hf_token = os.getenv('HF_TOKEN')
 try:
-    processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-448", use_auth_token=hf_token)
-    model = AutoModelForImageToText.from_pretrained("google/paligemma-3b-mix-448", use_auth_token=hf_token)
+    model = pipeline("image-to-text", model="google/paligemma-3b-mix-448", use_auth_token=hf_token)
 except Exception as e:
     st.error(f"Error loading image-to-text model: {e}")
     st.stop()
 
 try:
-    model_colpali = ColPali.from_pretrained("vidore/colpali-v1.2", torch_dtype=torch.bfloat16, use_auth_token=hf_token).to(device)
-    processor_colpali = ColPaliProcessor.from_pretrained("google/paligemma-3b-mix-448", use_auth_token=hf_token)
+    model_colpali = ColPali.from_pretrained("vidore/colpali-v1.2", torch_dtype=torch.bfloat16).to(device)
+    processor_colpali = ColPaliProcessor.from_pretrained("google/paligemma-3b-mix-448")
 except Exception as e:
     st.error(f"Error loading ColPali model or processor: {e}")
     st.stop()
 
 try:
-    model_qwen = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", use_auth_token=hf_token).to(device)
-    processor_qwen = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", use_auth_token=hf_token)
+    model_qwen = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct").to(device)
+    processor_qwen = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
 except Exception as e:
     st.error(f"Error loading Qwen model or processor: {e}")
     st.stop()