Update app.py

app.py

@@ -1,17 +1,18 @@
-import streamlit as st
-from transformers import pipeline
 from diffusers import DiffusionPipeline
+from transformers import pipeline
+import streamlit as st
 from PIL import Image
 import requests
 from io import BytesIO
 
+text_to_image = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
 
-# Load the image-to-text pipeline
+# Initialize the pipeline
 image_to_text = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
 
-# Load the text-to-image model
-text_to_image = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
+st.title('Image Captioning Application')
 
+# Function to load images from URL
 def load_image_from_url(url):
     try:
         response = requests.get(url)
@@ -21,30 +22,39 @@ def load_image_from_url(url):
         st.error(f"Error loading image from URL: {e}")
         return None
 
-def main():
-    st.title("Image to Story to Image Converter")
-
-    # User input for image URL
-    image_url = st.text_input("Enter the image URL:")
-
-    # Process the image and display the generated image
-    if st.button("Generate Story and Image"):
-        if image_url:
-            try:
-                image = load_image_from_url(image_url)
-                # Generate caption from image
-                image_text = image_to_text(image)
-                rephrased_text = "I want to buy " + image_text + " and [MASK] for my children"
+# User option to select input type: Upload or URL
+input_type = st.radio("Select input type:", ("Upload Image", "Image URL"))
+
+if input_type == "Upload Image":
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption='Uploaded Image', use_column_width=True)
+elif input_type == "Image URL":
+    image_url = st.text_input("Enter the image URL here:", "")
+    if image_url:
+        image = load_image_from_url(image_url)
+        if image:
+            st.image(image, caption='Image from URL', use_column_width=True)
+
+# Generate caption button
+if st.button('Generate Caption'):
+    if not image:
+        st.warning("Please upload an image or enter an image URL.")
+    else:
+        with st.spinner("Generating caption..."):
+            # Process the image and generate caption
+            if input_type == "Upload Image":
+                # Save the uploaded image to a temporary file to pass its path to the model
+                with open("temp_image.jpg", "wb") as f:
+                    f.write(uploaded_file.getbuffer())
+                result = image_to_text("temp_image.jpg")
+            elif input_type == "Image URL" and image_url:
+                result = image_to_text(image_url)
                 
-                # Generate image from rephrased text
-                generated_image = text_to_image(rephrased_text)
-
-                # Display the generated image
-                st.image(generated_image, caption='Generated Image')
-            except Exception as e:
-                st.error(f"Error processing image: {e}")
-        else:
-            st.warning("Please enter an image URL.")
-
-if __name__ == "__main__":
-    main()
+            if result:
+                generated_text = result[0]['generated_text']
+                st.success(f'Generated Caption: {generated_text}')
+            else:
+                st.error("Failed to generate caption.")
+