# Import necessary libraries
import streamlit as st  # Streamlit for web application
from transformers import pipeline  # Hugging Face transformer pipeline
from PIL import Image  # Python Imaging Library for image handling

# Set the title of the Streamlit app
st.set_page_config(page_title="Storytelling Friend",
                   page_icon="haha")
 # Title of the application

# Create a file uploader for the image
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])  # User uploads an image

# Load the image captioning model
caption_model = pipeline("image-captioning", model="facebook/blip-image-captioning-base")  # Load pre-trained model

# Load the text generation model
story_model = pipeline("text-generation", model="gpt2")  # Load a text generation model

# Function to generate a story from the caption
def generate_story(caption):  # Define a function to generate a story
    story_input = f"Once upon a time, {caption}"  # Create a story prompt
    story = story_model(story_input, max_length=150, num_return_sequences=1)[0]['generated_text']  # Generate the story
    return story  # Return the generated story

# Process the uploaded image and generate story
if uploaded_file is not None:  # Check if a file is uploaded
    image = Image.open(uploaded_file)  # Open the uploaded image
    st.image(image, caption="Uploaded Image", use_column_width=True)  # Display the uploaded image

    caption = caption_model(image)[0]['caption']  # Generate caption from the image
    st.subheader("Image Caption:")  # Subtitle for the caption
    st.write(caption)  # Display the caption

    # Generate story based on the caption
    story = generate_story(caption)  # Call the story generation function
    st.subheader("Generated Story:")  # Subtitle for the generated story
    st.write(story)  # Display the generated story