Update app.py

app.py

@@ -1,115 +1,82 @@
 # app.py
+
 """
 app.py
 
 Streamlit application for Image-to-Story demo.
-Allows users to upload an image or use a demo image, generates a caption, creates a child-friendly story,
-and plays it back as audio.
-Suitable for deployment on Hugging Face Spaces.
+Allows demo/upload image, generates a caption, a trimmed story,
+and plays back as MP3 via gTTS.
 """
 import streamlit as st
 from PIL import Image
 import warnings
 from modules import (
-    load_captioner,
-    load_story_gen,
-    load_tts,
-    generate_caption,
-    generate_story,
-    generate_audio,
-    generate_story_simple
+    load_captioner, load_story_gen,
+    generate_caption, generate_story_simple,
+    generate_audio
 )
 
-# Suppress deprecation warnings for cleaner UI
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
-
 def reset_state():
-    """
-    Clear generated caption, story, and audio when image source changes.
-    """
-    for key in ["caption", "story", "audio_data", "audio_sr"]:
+    for key in ["caption", "story", "audio_bytes", "audio_mime"]:
         if key in st.session_state:
             del st.session_state[key]
 
-
 def main():
     st.title("🖼️ → 📖 Image-to-Story App for Kids")
-    st.write(
-        "Upload an image or use the demo image to get an engaging story suitable for 3–10 year-olds, "
-        "with audio playback powered by Hugging Face pipelines!"
-    )
+    st.write("Upload or demo an image to get a 50–100 word story and audio!")
 
-    # Choose image source with callback to reset state
-    source = st.radio(
-        "Choose image source:",
-        ("Upload my own image", "Use demo image"),
-        on_change=reset_state
-    )
+    source = st.radio("Image source:",
+                      ("Upload my own image", "Use demo image"),
+                      on_change=reset_state)
 
-    # Load pipelines (cached) with a friendly spinner
+    # Load pipelines once
     if "models_loaded" not in st.session_state:
-        with st.spinner("Getting things ready, please wait…"):
+        with st.spinner("Loading models…"):
             st.session_state.captioner = load_captioner()
             st.session_state.story_gen = load_story_gen()
-            st.session_state.tts_pipe  = load_tts()
         st.session_state.models_loaded = True
 
     captioner = st.session_state.captioner
     story_gen = st.session_state.story_gen
-    tts_pipe  = st.session_state.tts_pipe
 
-    # Get image object
+    # Acquire image
     if source == "Use demo image":
-        img = Image.open("test_kids_playing.jpg")
+        img = Image.open("test_kids_playing.jpg").convert("RGB")
     else:
-        uploaded = st.file_uploader(
-            "Upload an image (PNG, JPG, JPEG)",
-            type=["png", "jpg", "jpeg"],
-            key="upload"
-        )
+        uploaded = st.file_uploader("Upload an image",
+                                    type=["png", "jpg", "jpeg"])
         if not uploaded:
             return
-        img = Image.open(uploaded)
+        img = Image.open(uploaded).convert("RGB")
 
-    # Display image
     st.image(img, use_container_width=True)
 
-    # Generate caption
+    # Caption
     if "caption" not in st.session_state:
-        with st.spinner("Generating image caption…"):
+        with st.spinner("Captioning image…"):
             st.session_state.caption = generate_caption(captioner, img)
     st.markdown(f"**Caption:** {st.session_state.caption}")
 
-    # Generate story
+    # Story
     if "story" not in st.session_state:
-        with st.spinner("Creating story for kids…"):
-            st.session_state.story = generate_story_simple(story_gen, st.session_state.caption,
-                                                           min_words=50, max_words=100)
-            
-            # st.session_state.story = generate_story(
-            #    story_gen,
-            #    st.session_state.caption,
-            #    min_words=50,
-            #    max_words=100
-            # )
+        with st.spinner("Creating story…"):
+            st.session_state.story = generate_story_simple(
+                story_gen, st.session_state.caption, 50, 100
+            )
     st.markdown(f"**Story:** {st.session_state.story}")
 
-    # Generate audio
-    if "audio_data" not in st.session_state:
-        with st.spinner("Synthesizing speech…"):
-            audio_array, sr = generate_audio(tts_pipe, st.session_state.story)
-            st.session_state.audio_data = audio_array
-            st.session_state.audio_sr   = sr
+    # Audio
+    if "audio_bytes" not in st.session_state:
+        with st.spinner("Generating audio…"):
+            audio_bytes, mime = generate_audio(st.session_state.story)
+            st.session_state.audio_bytes = audio_bytes
+            st.session_state.audio_mime = mime
 
-    # Audio playback button
     if st.button("🔊 Play Story Audio"):
-        st.audio(
-            data=st.session_state.audio_data,
-            format="audio/wav",
-            sample_rate=st.session_state.audio_sr
-        )
-
+        st.audio(data=st.session_state.audio_bytes,
+                 format=st.session_state.audio_mime)
 
 if __name__ == "__main__":
     main()