Update app.py

app.py

@@ -1,59 +1,26 @@
-import streamlit as st # Don't forget to include `streamlit` in your `requirements.txt`
-from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration
-
-# Set up authentication
-if "hf_token" not in st.session_state:
-    st.title("Authentication Required")
-    st.write("Please authenticate with Hugging Face using the following token:")
-    hf_token = st.text_input("Enter your token", type="password")
-
-    if hf_token == "":
-        st.stop()
-else:
-    hf_token = st.session_state.hf_token
-
-# Load Token from Storage
-if "hf_token_local" not in st.session_state:
-    st.title("Load Token from Storage")
-    st.write("Please load your token from storage (e.g., environment variable, file)")
-    hf_token_local = st.text_input("Enter your token", type="password")
-
-    if hf_token_local == "":
-        st.stop()
-else:
-    hf_token_local = st.session_state.hf_token_local
-
-# Load Processor and Model
-if hf_token or hf_token_local:
-    processor = PaliGemmaProcessor.from_pretrained(
-        "google/paligemma2",
-        token=hf_token,
-        local_file_dir="/tmp/",
-    )
-    model = PaliGemmaForConditionalGeneration.from_pretrained(
-        "google/paligemma2",
-        token=hf_token,
-        local_file_dir="/tmp/",
-    )
-
-    # Rest of your code
-else:
-    st.title("No Token Found")
-    st.write("Please authenticate with Hugging Face or load token from storage")
-
-# Use the model
-def main():
-    if "output" not in st.session_state:
-        st.write("Model output")
-    else:
-        st.write(st.session_state.output)
-
-    # Add a button to generate text using the model
-    if st.button("Generate Text"):
-        input_text = st.text_input("Input text")
-        if input_text:
-            output = model.generate(input_text, max_length=50)
-            st.session_state.output = output
-
-if __name__ == "__main__":
-    main()
+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 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):
+    """Extract text from image using PaliGemma2."""
+    # Preprocess the image
+    inputs = processor(images=image, return_tensors="pt")
+    
+    # Generate predictions
+    with torch.no_grad():
+        generated_ids = model.generate(**inputs)
+        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    
+    return text