Update app.py

app.py

@@ -1,14 +1,23 @@
+import os
 import streamlit as st
 from PIL import Image
 import torch
 from transformers import AutoProcessor, AutoModelForVision2Seq
 
+# Get Hugging Face API key from Hugging Face Spaces secrets
+HF_TOKEN = os.getenv("HF_KEY")
+
+# Ensure API key is available
+if not HF_TOKEN:
+    st.error("❌ Hugging Face API key not found! Set it as 'HF_KEY' in Spaces secrets.")
+    st.stop()
+
 # Load the PaliGemma model and processor
 @st.cache_resource
 def load_model():
     model_name = "google/paligemma2-3b-mix-224"
-    processor = AutoProcessor.from_pretrained(model_name)
-    model = AutoModelForVision2Seq.from_pretrained(model_name)
+    processor = AutoProcessor.from_pretrained(model_name, token=HF_TOKEN)
+    model = AutoModelForVision2Seq.from_pretrained(model_name, token=HF_TOKEN)
     return processor, model
 
 processor, model = load_model()
@@ -30,4 +39,4 @@ if uploaded_file:
             output = model.generate(**inputs)
         
         answer = processor.batch_decode(output, skip_special_tokens=True)[0]
-        st.success(f"Answer: {answer}")
+        st.success(f"✅ Answer: {answer}")