import streamlit as st
from transformers import pipeline
from PIL import Image
import requests
import os
from io import BytesIO
import tempfile

# Load Hugging Face pipelines
@st.cache_resource
def load_pipelines():
    image_captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    story_generator = pipeline("text-generation", model="Tevatron/tiny-stories-generator", max_length=200)
    text_to_speech = pipeline("text-to-speech", model="espnet/kan-bayashi_ljspeech_vits", framework="pt")
    return image_captioner, story_generator, text_to_speech

# Define main interface
def main():
    st.set_page_config(page_title="Kids Story Maker 🧒📖", page_icon="📸", layout="centered")

    # Child-friendly header
    st.markdown("<h1 style='color:#FF69B4; font-family:Comic Sans MS;'>🎨 Welcome to Kids Story Maker!</h1>", unsafe_allow_html=True)
    st.markdown("<h3 style='color:#4CAF50;'>Upload a picture, and we'll turn it into a magical story and voice! 🐻✨</h3>", unsafe_allow_html=True)

    image_captioner, story_generator, text_to_speech = load_pipelines()

    uploaded_file = st.file_uploader("🖼️ Upload an image:", type=["png", "jpg", "jpeg"])

    if uploaded_file:
        image = Image.open(uploaded_file)
        st.image(image, caption="Your Image", use_column_width=True)

        with st.spinner("🔍 Generating a description..."):
            caption = image_captioner(image)[0]['generated_text']
            st.success(f"📝 Description: {caption}")

        # Prompt template for storytelling
        story_prompt = f"Write a short story for children aged 3 to 10 based on this description: {caption}. The story should be creative, friendly, and use simple words."

        with st.spinner("✍️ Creating your story..."):
            story = story_generator(story_prompt)[0]['generated_text']
            st.success("📖 Here's your story:")
            st.write(story)

        with st.spinner("🔊 Turning your story into voice..."):
            speech = text_to_speech(story)[0]['audio']
            st.audio(speech, format="audio/wav")

# Run the app
if __name__ == "__main__":
    main()