Update app.py

app.py

@@ -1,13 +1,16 @@
-import streamlit as st
-from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration
-from PIL import Image
-import torch
+import streamlit as st  # Don't forget to include `streamlit` in your `requirements.txt` file to ensure the app runs properly on Hugging Face Spaces.
+from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration  # Make sure that the Hugging Face `transformers` library version supports the `PaliGemma2` model. You may need to specify the version in `requirements.txt`.
+from PIL import Image  # Ensure the `pillow` library is included in your `requirements.txt`.
+import torch  # Since PyTorch is required for this app, specify the appropriate version of `torch` in `requirements.txt` based on compatibility with the model.
 import os
 
 def load_model():
-    """Load PaliGemma2 model and processor."""
-    processor = PaliGemmaProcessor.from_pretrained("google/paligemma2")
-    model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2")
+    """Load PaliGemma2 model and processor with Hugging Face token."""
+    token = os.getenv("HUGGINGFACEHUB_API_TOKEN")  # Retrieve token from environment variable
+    if not token:
+        raise ValueError("Hugging Face API token not found. Please set it in the environment variables.")
+    processor = PaliGemmaProcessor.from_pretrained("google/paligemma2", token=token)
+    model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2", token=token)
     return processor, model
 
 def process_image(image, processor, model):
@@ -29,8 +32,12 @@ def main():
 
     # Load model and processor
     with st.spinner("Loading PaliGemma2 model... This may take a few moments."):
-        processor, model = load_model()
-    st.success("Model loaded successfully!")
+        try:
+            processor, model = load_model()
+            st.success("Model loaded successfully!")
+        except ValueError as e:
+            st.error(str(e))
+            st.stop()
 
     # User input: upload image
     uploaded_image = st.file_uploader("Upload an image containing text", type=["png", "jpg", "jpeg"])