Update app.py

app.py

@@ -1,41 +1,56 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForImageCaptioning
-import requests
 from PIL import Image
-import numpy as np
+import textwrap
+import google.generativeai as genai
 
-# Initialize the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("microsoft/beit-base-patch16-224-in21k")
-model = AutoModelForImageCaptioning.from_pretrained("microsoft/beit-base-patch16-224-in21k")
+# Function to display formatted Markdown text
+def to_markdown(text):
+    text = text.replace('•', '  *')
+    return textwrap.indent(text, '> ', predicate=lambda _: True)
 
-def generate_caption(image_url):
-  # Get the image from the URL
-  image = Image.open(requests.get(image_url, stream=True).raw)
-  
-  # Preprocess the image
-  input_array = np.array(image) / 255.0
-  input_array = np.transpose(input_array, (2, 0, 1))
-  input_ids = tokenizer(image_url, return_tensors="pt").input_ids
-
-  # Generate the caption
-  output = model.generate(input_ids, max_length=20)
-  caption = tokenizer.batch_decode(output, skip_special_tokens=True)
+# Function to generate content using Gemini API
+def generate_gemini_content(prompt, model_name='gemini-pro', image=None):
+    model = genai.GenerativeModel(model_name)
+    if image:
+        response = model.generate_content([prompt, image])
+    else:
+        response = model.generate_content(prompt)
 
-  return caption[0]
+    return response
 
+# Streamlit app
 def main():
-  # Create a sidebar for the user to input the image URL
-  st.sidebar.header("Image Caption Generator")
-  image_url = st.sidebar.text_input("Enter the URL of an image:")
-
-  # Generate the caption if the user clicks the button
-  if st.sidebar.button("Generate Caption"):
-    if image_url != "":
-      caption = generate_caption(image_url)
-      st.success(f"Caption: {caption}")
-    else:
-      st.error("Please enter a valid image URL.")
+    st.title("Gemini API Demo with Streamlit")
+
+    # Get Gemini API key from user input
+    api_key = st.text_input("Enter your Gemini API key:")
+    genai.configure(api_key=api_key)
+
+    # Choose a model
+    model_name = st.selectbox("Select a Gemini model", ["gemini-pro", "gemini-pro-vision"])
+
+    # Get user input prompt
+    prompt = st.text_area("Enter your prompt:")
+
+    # Get optional image input
+    image_file = st.file_uploader("Upload an image (if applicable):", type=["jpg", "jpeg", "png"])
+
+    # Display image if provided
+    if image_file:
+        st.image(image_file, caption="Uploaded Image", use_column_width=True)
+
+    # Generate content on button click
+    if st.button("Generate Content"):
+        st.markdown("### Generated Content:")
+        if image_file:
+            # If an image is provided, use gemini-pro-vision model
+            image = Image.open(image_file)
+            response = generate_gemini_content(prompt, model_name='gemini-pro-vision', image=image)
+        else:
+            response = generate_gemini_content(prompt, model_name=model_name)
+
+        # Display the generated content in Markdown format
+        st.markdown(to_markdown(response.text))
 
-# Run the main function
 if __name__ == "__main__":
     main()