Update app.py

app.py

@@ -1,63 +1,59 @@
-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
-
+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():
-    # Set page configuration
-    st.set_page_config(page_title="Text Reading with PaliGemma2", layout="centered")
-    st.title("Text Reading from Images using PaliGemma2")
-
-    # Load model and processor
-    with st.spinner("Loading PaliGemma2 model... This may take a few moments."):
-        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"])
-
-    if uploaded_image is not None:
-        # Display uploaded image
-        image = Image.open(uploaded_image)
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-
-        # Extract text button
-        if st.button("Extract Text"):
-            with st.spinner("Processing image..."):
-                extracted_text = process_image(image, processor, model)
-                st.success("Text extraction complete!")
-                st.subheader("Extracted Text")
-                st.write(extracted_text)
-
-    # Footer
-    st.markdown("---")
-    st.markdown("**Built with [PaliGemma2](https://huggingface.co/google/paligemma2) and Streamlit**")
+    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()
+    main()