Update app.py

app.py

@@ -1,3 +1,7 @@
+# ======================================
+# Package Import
+# ======================================
+
 import streamlit as st
 from PIL import Image
 import time
@@ -7,6 +11,10 @@ from datasets import load_dataset
 import soundfile as sf
 import torch
 
+# ======================================
+# Basic Initialization
+# ======================================
+
 # Initialize image captioning pipeline with pretrained model
 # Model source: Hugging Face Model Hub
 _image_caption_pipeline = pipeline(
@@ -31,6 +39,10 @@ _SPEECH_PIPELINE = pipeline("text-to-speech", model="microsoft/speecht5_tts")
 _EMBEDDINGS_DATASET = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
 _DEFAULT_SPEAKER_EMBEDDING = torch.tensor(_EMBEDDINGS_DATASET[7306]["xvector"]).unsqueeze(0)
 
+# ======================================
+# Function settings
+# ======================================
+
 def generate_image_caption(input_image):
     """
     Generate a textual description for an input image using a pretrained model.
@@ -196,35 +208,138 @@ def generate_audio_from_story(story_text: str, output_path: str = "output.wav")
         raise RuntimeError(f"Audio synthesis failed: {str(error)}") from error
 
 
-# App title
-st.title("Best Story Teller")
+# ======================================
+# Page Configuration & Custom Styling
+# ======================================
+st.set_page_config(
+    page_title="Magic Story Generator",
+    page_icon="🧚",
+    layout="wide",
+    initial_sidebar_state="collapsed"
+)
 
-# Write some text
-st.write("Upload a picture and start your journey of creativeness and imagination")
+# Custom CSS styling for child-friendly interface
+st.markdown("""
+<style>
+    /* Primary title styling */
+    .main-title {
+        color: #E91E63;
+        font-size: 2.8rem;
+        text-align: center;
+        padding: 20px;
+        text-shadow: 2px 2px #FFC107;
+    }
+    
+    /* Prompt buttons styling */
+    .prompt-btn {
+        background: #4CAF50 !important;
+        border-radius: 15px !important;
+        padding: 15px 30px !important;
+        font-size: 1.1rem !important;
+        margin: 10px;
+    }
+    
+    /* Story container styling */
+    .story-container {
+        background: #FFF3E0;
+        border-radius: 20px;
+        padding: 25px;
+        margin: 20px 0;
+        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+    }
+    
+    /* Progress spinner styling */
+    .stSpinner > div {
+        font-size: 1.2rem !important;
+        color: #9C27B0 !important;
+    }
+</style>
+""", unsafe_allow_html=True)
 
-# File uploader for image and audio
-uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+# ======================================
+# Main Application Interface
+# ======================================
+st.markdown('<p class="main-title">🧚 Welcome to Magic Story Maker!</p>', unsafe_allow_html=True)
 
-# Display image with spinner
+# File upload section
+with st.container():
+    st.subheader("Step 1: Upload Your Picture")
+    uploaded_image = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"],label_visibility="collapsed")
+
+# Main processing flow
 if uploaded_image is not None:
-    with st.spinner("Loading image..."):
+    # Display uploaded image
+    with st.spinner("✨ Magical image processing..."):
         image = Image.open(uploaded_image)
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-    with st.spinner("Captioning image..."):
-        caption_from_file = generate_image_caption(image)
-    with st.spinner("Adding some magics and imagination..."):
-        system_prompt = """"Create magical children's stories (ages 4-8) from image captions. Include:
-            Friendly animal/human characters
-            Simple plot with happy resolution
-            Sensory details ("glowing firefly wings") and sounds ("Splash!")
-            Gentle lessons about friendship/courage/nature
-            Choices ("Should the fox knock or sing?")
-            Key image elements woven into the story
-            Diverse characters, no stereotypes
-            End with happy twist + rhyming moral ("Brave small paws, when they try/Can touch the stars in the sky"). Keep language warm and playful."""
-        user_prompt = caption_from_file
-        story = generate_story_content(system_prompt, user_prompt)
-        st.write(story)
-    with st.spinner("Finding the best voice actor"):
-        generated_audio = generate_audio_from_story(story,"childrens_story.wav")
-        st.audio(generated_audio)
+        st.image(image, caption="Your Magical Image", use_column_width=True)
+
+    # Prompt selection section
+    with st.container():
+        st.subheader("Step 2: Choose Story Style")
+        
+        # Create three columns for prompt buttons
+        col1, col2, col3 = st.columns(3)
+        with col1:
+            if st.button("📚 Learning Story", 
+                        help="Generate educational story with life lessons",
+                        key="edu_btn"):
+                st.session_state.selected_prompt = "educational"
+        with col2:
+            if st.button("🌠 Fantasy Adventure", 
+                        help="Create magical adventure story",
+                        key="fantasy_btn"):
+                st.session_state.selected_prompt = "adventure"
+        with col3:
+            if st.button("🐻 Animal Friends", 
+                        help="Make story about friendly animals",
+                        key="animal_btn"):
+                st.session_state.selected_prompt = "animal"
+
+    # Define prompt templates
+    PROMPT_TEMPLATES = {
+        "educational": {
+            "system": "You are a children's educator. Create a simple 150-word story that teaches basic life skills or moral lessons.",
+            "icon": "📚"
+        },
+        "adventure": {
+            "system": "You are a fantasy writer. Create a 150-word magical adventure story suitable for children.",
+            "icon": "🌠"
+        },
+        "animal": {
+            "system": "You are an animal expert. Create a 150-word story about friendly animals learning together.",
+            "icon": "🐻"
+        }
+    }
+
+    # Story generation section
+    with st.spinner(f"{PROMPT_TEMPLATES[st.session_state.selected_prompt]['icon']} Creating your story..."):
+        # Generate image caption
+        image_caption = generate_image_caption(image)
+        
+        # Generate story content
+        selected_template = PROMPT_TEMPLATES[st.session_state.selected_prompt]
+        story_text = generate_story_content(
+            system_prompt=selected_template["system"],
+            user_prompt=image_caption
+        )
+        
+        # Display formatted story
+        st.subheader("Step 3: Your Magical Story")
+        st.markdown(f'<div class="story-container">{story_text}</div>', 
+                   unsafe_allow_html=True)
+
+    # Audio generation section
+    with st.spinner("🔮 Preparing story narration..."):
+        audio_file = generate_audio_from_story(story_text, "story_audio.wav")
+        st.subheader("🎧 Listen to Your Story")
+        st.audio(audio_file)
+
+# Help section
+st.markdown("---")
+st.subheader("🌟 How to Use:")
+st.info("""
+1. Upload any picture (animals, nature, or people work best!)
+2. Choose your favorite story style
+3. Wait for magic to happen!
+4. Listen to your personalized story
+""")