import streamlit as st
from transformers import pipeline
from PIL import Image
from gtts import gTTS
import tempfile

# Page configuration
st.set_page_config(page_title="🧸 Story Generator (CPU Friendly)", page_icon="📚")
st.title("🖼️ Image to Story Generator (CPU Version)")
st.write("Upload an image and enjoy a short story with audio narration — all without a GPU!")

# Image upload
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_container_width=True)

    if st.button("Generate Story"):
        with st.spinner("📷 Generating caption..."):
            # Image captioning model (CPU-friendly)
            captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
            caption = captioner(image)[0]['generated_text'].strip()

        with st.spinner("✍️ Generating story..."):
            # Using Falcon-rw-1b text generation model (CPU-adapted)
            story_prompt = f"A short and fun story for children about: {caption}"
            generator = pipeline("text-generation", model="tiiuae/falcon-rw-1b")
            story = generator(
                story_prompt,
                max_length=150,
                do_sample=True,
                temperature=0.9,
                top_p=0.95
            )[0]['generated_text'].strip()

            # Limit maximum word count to 100
            story = story.replace("\n", " ")
            words = story.split()
            if len(words) > 100:
                story = " ".join(words[:100]) + "..."

        with st.spinner("🔊 Converting to speech..."):
            # Text-to-speech
            tts = gTTS(text=story, lang='en')
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
            tts.save(temp_file.name)

            # Display results
            st.subheader("📖 Generated Story")
            st.write(story)

            st.subheader("🔊 Listen to the Story")
            st.audio(temp_file.name, format="audio/mp3")