Update vision_llm.py

vision_llm.py

@@ -2,12 +2,15 @@
 from transformers import AutoProcessor, AutoModelForVision2Seq
 import torch
 from PIL import Image
+import os
 
 class VisionLLM:
-    def __init__(self, device="cuda", model_id="google/paligemma2-3b-pt-224"):  # Corrected model ID
+    def __init__(self, device="cuda", model_id="google/paligemma2-3b-pt-224", use_auth_token=None):  # Corrected __init__
         self.device = device
-        self.processor = AutoProcessor.from_pretrained(model_id)
-        self.model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16).to(self.device)
+        if use_auth_token is None:
+            use_auth_token=os.environ.get("HF_TOKEN", None)
+        self.processor = AutoProcessor.from_pretrained(model_id, use_auth_token=use_auth_token)  #Use the auth token
+        self.model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16, use_auth_token=use_auth_token).to(self.device) #Use the auth token
 
     def describe_images(self, images, prompt="", max_length=128):
         if isinstance(images, list):