import streamlit as st
import requests
from PIL import Image, ImageDraw, ImageFont
import os
from transformers import BlipProcessor, BlipForConditionalGeneration

# Load the BLIP model for creative caption generation
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# Function to generate a creative caption
def generate_caption(image):
    # Prepare image for the model
    inputs = processor(images=image, return_tensors="pt")
    
    # Generate the caption
    out = model.generate(**inputs)
    
    # Decode the output to get a readable caption
    caption = processor.decode(out[0], skip_special_tokens=True)
    
    # Return the creative caption
    return caption

# Streamlit app
def main():
    st.title("Creative Image Caption Generator")
    st.write("Upload an image, and let the AI generate a creative and descriptive caption for it.")
    
    # Upload image
    uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"])
    if uploaded_file is not None:
        # Open the uploaded image
        image = Image.open(uploaded_file)
        
        # Display the uploaded image
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Generate caption button
        if st.button("Generate Creative Caption"):
            caption = generate_caption(image)
            
            # Display the generated caption
            st.write("### Generated Creative Caption: ")
            st.write(caption)
            
            # Save image with caption overlay
            if st.button("Save Image with Caption"):
                # Draw the caption on the image
                draw = ImageDraw.Draw(image)
                font = ImageFont.load_default()
                draw.text((10, 10), caption, fill="white", font=font)
                
                # Save the image with the caption
                save_path = os.path.join("saved_images", "image_with_caption.jpg")
                image.save(save_path)

                # Provide download button for the captioned image
                with open(save_path, "rb") as f:
                    st.download_button("Download Image with Caption", f, file_name="image_with_caption.jpg")

if __name__ == "__main__":
    # Create a directory to save the image
    if not os.path.exists("saved_images"):
        os.makedirs("saved_images")
        
    main()