Update app.py

app.py

@@ -1,28 +1,49 @@
 import streamlit as st
-from transformers import pipeline
+from PIL import Image
+import requests
+from transformers import BlipProcessor, BlipForConditionalGeneration
 
-# Initialize the image captioning pipeline
-captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+# Load the model and processor outside the main function to avoid reloading on every run
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-# Streamlit app title
-st.title("Image to Text Captioning")
+def generate_caption(img_url):
+    raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
 
-# Input for image URL
-image_url = st.text_input("Enter the URL of the image:")
+    # Conditional image captioning
+    text = "a photography of"
+    inputs = processor(raw_image, text, return_tensors="pt")
+    out = model.generate(**inputs)
+    conditional_caption = processor.decode(out[0], skip_special_tokens=True)
 
-# If an image URL is provided
-if image_url:
-    try:
-        # Display the image
-        st.image(image_url, caption="Provided Image", use_column_width=True)
-        
-        # Generate the caption
-        caption = captioner(image_url)
-        
-        # Display the caption
-        st.write("**Generated Caption:**")
-        st.write(caption[0]['generated_text'])
-    except Exception as e:
-        st.error(f"An error occurred: {e}")
+    # Unconditional image captioning
+    inputs = processor(raw_image, return_tensors="pt")
+    out = model.generate(**inputs)
+    unconditional_caption = processor.decode(out[0], skip_special_tokens=True)
 
-# To run this app, save this code to a file (e.g., `app.py`) and run `streamlit run app.py` in your terminal.
+    return conditional_caption, unconditional_caption
+
+def main():
+    st.title("Image Captioning App")
+    
+    img_url = st.text_input("Enter the image URL:")
+    if img_url:
+        try:
+            # Display the image
+            image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
+            st.image(image, caption='Input Image', use_column_width=True)
+            
+            # Generate captions
+            conditional_caption, unconditional_caption = generate_caption(img_url)
+            
+            # Display captions
+            st.subheader("Conditional Image Caption")
+            st.write(conditional_caption)
+            
+            st.subheader("Unconditional Image Caption")
+            st.write(unconditional_caption)
+        except Exception as e:
+            st.error(f"Error processing the image: {e}")
+
+if __name__ == "__main__":
+    main()