Update app.py

app.py

@@ -1,8 +1,16 @@
+import os
 import torch
 from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
 from PIL import Image
 import gradio as gr
 
+# -----------------------------------------------------------------------------
+# Load HF token from environment
+# -----------------------------------------------------------------------------
+HF_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HUGGINGFACEHUB_API_TOKEN environment variable not set")
+
 # -----------------------------------------------------------------------------
 # 1) GPU inference function
 # -----------------------------------------------------------------------------
@@ -10,21 +18,20 @@ def run_inference_on_gpu(
     model_id: str,
     image: Image.Image,
     prompt: str = "caption",
-    max_new_tokens: int = 100,
-    use_auth_token: bool = True
+    max_new_tokens: int = 100
 ) -> str:
     # ensure CUDA is available
     assert torch.cuda.is_available(), "CUDA not available—check your PyTorch installation!"
     device = torch.device("cuda")
     dtype = torch.float16
 
-    # load tokenizer + model onto GPU
-    processor = AutoProcessor.from_pretrained(model_id, use_auth_token=use_auth_token)
+    # load tokenizer + model onto GPU with explicit token
+    processor = AutoProcessor.from_pretrained(model_id, use_auth_token=HF_TOKEN)
     model = PaliGemmaForConditionalGeneration.from_pretrained(
         model_id,
         torch_dtype=dtype,
         device_map=None,
-        use_auth_token=use_auth_token
+        use_auth_token=HF_TOKEN
     ).to(device).eval()
 
     # build multimodal prompt
@@ -53,7 +60,6 @@ def run_inference_on_gpu(
     # decode
     return processor.decode(outputs[0].cpu(), skip_special_tokens=True)
 
-
 # -----------------------------------------------------------------------------
 # 2) Gradio UI
 # -----------------------------------------------------------------------------