import streamlit as st
from PIL import Image
import time
from transformers import pipeline
import tempfile
import os


# Generate caption from image
def image_to_caption(image_path):
    """Generates a caption for the given image using a pre-trained model."""
    imgtocaption = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    caption = imgtocaption(image_path)
    return caption[0]['generated_text']

# Generate story from caption
def caption_to_story(text):
    captiontostory = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
    story = captiontostory(text, max_length=150, min_length=50)
    return story[0]['generated_text']

# Convert story to speech
def story_to_audio(text):
    tts_pipe = pipeline("text-to-audio", model="facebook/mms-tts-eng")
    audio_output = tts_pipe(text[:1000]) 
    return audio_output['audio'], audio_output['sampling_rate']
        
# Child-Friendly Interface Design
st.set_page_config(page_title="Magic Story House", page_icon="🧚")
st.title("🧚 Magic Image Story Generator")
st.markdown("Upload an image and generate your exclusive fairy tale!")

# File Upload
uploaded_image = st.file_uploader("Choose a picture", type=["jpg", "jpeg", "png"], key="image_uploader")

# Main application
def main():
    if uploaded_image is not None:
        try:
            # Process image
            with st.spinner("Processing image..."):
                image = Image.open(uploaded_image)
                st.image(image, caption="Uploaded Image", use_column_width=True)
                
                # Save temporary file
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
                    image.save(temp_file.name)
                    image_path = temp_file.name

            # Generate caption
            with st.spinner("Generating caption..."):
                caption = image_to_caption(image_path)
                st.subheader("Generated Caption")
                st.write(caption)

            # Generate story
            with st.spinner("Generating story..."):
                story = caption_to_story(caption)
                st.subheader("Generated Story")
                st.write(story)

            # Generate speech
            with st.spinner("Generating audio..."):
                audio_array, sample_rate = story_to_audio(story)
                if audio_array is not None:
                    st.subheader("Audio Narration")
                    st.audio(audio_array, sample_rate=sample_rate)

        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
        finally:
            # Clean up temporary file
            if 'image_path' in locals() and os.path.exists(image_path):
                os.remove(image_path)

if __name__ == "__main__":
    main()