# Import necessary libraries
import streamlit as st
from transformers import pipeline
import pyttsx3
import os

# Function to convert image to text
def img2text(url):
    image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    text = image_to_text_model(url)[0]["generated_text"]
    return text

# Function to generate a story from text using GPT-2
def text2story(text):
    text_generator = pipeline("text-generation", model="gpt2")
    story = text_generator(text, max_length=95, num_return_sequences=1)[0]["generated_text"]
    return story

# Function to convert text to audio using pyttsx3
def text2audio(story_text):
    engine = pyttsx3.init()
    audio_file = "story_audio.wav"
    engine.save_to_file(story_text, audio_file)
    engine.runAndWait()
    return audio_file

# Main application
st.set_page_config(page_title="Image to Story", page_icon="📖")
st.header("📖 Image to Story")
st.markdown("### Turn your image into a fun story!")

uploaded_file = st.file_uploader("Select an Image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Save the uploaded file
    bytes_data = uploaded_file.getvalue()
    with open(uploaded_file.name, "wb") as file:
        file.write(bytes_data)

    st.image(uploaded_file, caption="Your Uploaded Image", use_column_width=True)

    # Stage 1: Image to Text
    st.text('🖼️ Processing image...')
    scenario = img2text(uploaded_file.name)
    st.write("**What I see:**", scenario)

    # Stage 2: Text to Story
    st.text('📝 Creating a story...')
    story = text2story(scenario)
    st.write("**Your Story:**", story)

    # Stage 3: Story to Audio
    st.text('🎙️ Turning your story into audio...')
    audio_file = text2audio(story)

    # Play button for audio
    if st.button("🎧 Listen to the Story"):
        st.audio(audio_file, format="audio/wav")

    # Clean up the generated audio file
    os.remove(audio_file)

# Add some fun prompts for kids
st.markdown("### 🎨 Tips for a Great Story!")
st.write("1. Upload a picture of your favorite animal, place, or toy!")
st.write("2. Imagine what's happening in the picture and let the story begin!")
st.write("3. Listen to your story and share it with your friends!")