Update app.py

app.py

@@ -1,59 +1,56 @@
-from transformers import pipeline
 import streamlit as st
+from transformers import pipeline
+from diffusers import DiffusionPipeline
 from PIL import Image
 import requests
-from io import BytesIO
-from diffusers import DiffusionPipeline
+import io
 
+# Load the image-to-text 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")
-# Initialize the pipeline
-image_to_text = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
 
-st.title('Image Captioning Application')
-
-# Function to load images from URL
-def load_image_from_url(url):
-    try:
-        response = requests.get(url)
-        img = Image.open(BytesIO(response.content))
-        return img
-    except Exception as e:
-        st.error(f"Error loading image from URL: {e}")
-        return None
-
-# 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)
-                
-            if result:
-                generated_text = result[0]['generated_text']
-                st.success(f'Generated Caption: {generated_text}')
-            else:
-                st.error("Failed to generate caption.")
+def main():
+    st.title("Image to Story to Image Converter")
+
+    # User input for text or URL
+    input_option = st.radio("Select input option:", ("Text", "URL"))
+
+    # Input text
+    if input_option == "Text":
+        text_input = st.text_input("Enter the text:")
+        if st.button("Generate Story and Image") and text_input:
+            generate_image(text_input)
+
+    # Input URL
+    elif input_option == "URL":
+        uploaded_file = st.file_uploader("Upload an image file:", 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)
+            if st.button("Generate Story and Image"):
+                image_text = image_to_text_from_file(uploaded_file)
+                generate_image(image_text)
+        else:
+            image_url = st.text_input("Enter the image URL:")
+            if st.button("Generate Story and Image") and image_url:
+                image_text = image_to_text_from_url(image_url)
+                generate_image(image_text)
+
+def image_to_text_from_file(uploaded_file):
+    image_bytes = io.BytesIO(uploaded_file.read())
+    return image_to_text(image_bytes)[0]['generated_text']
+
+def image_to_text_from_url(image_url):
+    response = requests.get(image_url)
+    image_bytes = io.BytesIO(response.content)
+    return image_to_text(image_bytes)[0]['generated_text']
+
+def generate_image(text):
+    rephrased_text = "I want to buy " + text + " and [MASK] for my children"
+    generated_image = text_to_image(rephrased_text)
+    st.image(generated_image, caption="Generated Image", use_column_width=True)
+
+if __name__ == "__main__":
+    main()