Update app.py

app.py

@@ -1,81 +1,89 @@
+# app.py
+"""
+app.py
+
+Streamlit application for Image-to-Story demo.
+Allows users to upload an image, generates a caption, creates a child-friendly story,
+and plays it back as audio.
+"""
 import streamlit as st
-from transformers import pipeline
 from PIL import Image
-import numpy as np
 import warnings
+from modules import (
+    load_captioner,
+    load_story_gen,
+    load_tts,
+    generate_caption,
+    generate_story,
+    generate_audio
+)
 
+# Suppress deprecation warnings for cleaner UI
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
-st.title("🖼️ → 📖 Image-to-Story App (with Bark TTS)")
-st.write("Upload an image, get a caption, spin it into a story, and play it aloud!")
 
-# 1) Define your cached loaders
-@st.cache_resource
-def load_captioner():
-    return pipeline("image-to-text", model="unography/blip-large-long-cap")
+def main():
+    st.title("🖼️ → 📖 Image-to-Story App for Kids")
+    st.write(
+        "Upload an image and get an engaging story suitable for 3–10 year-olds, "
+        "with audio playback powered by Hugging Face pipelines!"
+    )
 
-@st.cache_resource
-def load_story_gen():
-    return pipeline("text-generation", model="gpt2", tokenizer="gpt2")
+    # Load pipelines (cached) with a friendly spinner
+    if "models_loaded" not in st.session_state:
+        with st.spinner("Loading AI models, please wait…"):
+            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
 
-@st.cache_resource
-def load_bark_tts():
-    # Zero-config Bark text-to-speech
-    return pipeline("text-to-speech", model="suno/bark")
+    captioner = st.session_state.captioner
+    story_gen = st.session_state.story_gen
+    tts_pipe  = st.session_state.tts_pipe
 
-# 2) Load all models under one spinner
-if "models_loaded" not in st.session_state:
-    with st.spinner("Getting everything ready…"):
-        st.session_state.captioner = load_captioner()
-        st.session_state.story_gen  = load_story_gen()
-        st.session_state.tts_pipe   = load_bark_tts()
-    st.session_state.models_loaded = True
+    # File uploader for images
+    uploaded = st.file_uploader(
+        "Upload an image (PNG, JPG, JPEG)", type=["png", "jpg", "jpeg"]
+    )
+    if not uploaded:
+        return
 
-captioner = st.session_state.captioner
-story_gen = st.session_state.story_gen
-tts_pipe  = st.session_state.tts_pipe
+    # Display uploaded image
+    img = Image.open(uploaded)
+    st.image(img, use_container_width=True)
 
-# 3) Upload image
-uploaded = st.file_uploader("Upload an image", type=["png","jpg","jpeg"])
-if not uploaded:
-    st.stop()
+    # Generate caption once
+    if "caption" not in st.session_state:
+        with st.spinner("Generating image caption…"):
+            st.session_state.caption = generate_caption(captioner, img)
+    st.markdown(f"**Caption:** {st.session_state.caption}")
 
-img = Image.open(uploaded)
-st.image(img, use_container_width=True)
+    # Generate story once
+    if "story" not in st.session_state:
+        with st.spinner("Creating story for kids…"):
+            st.session_state.story = generate_story(
+                story_gen,
+                st.session_state.caption,
+                min_words=50,
+                max_words=100
+            )
+    st.markdown(f"**Story:** {st.session_state.story}")
 
-# 4) Generate caption once
-if "caption" not in st.session_state:
-    with st.spinner("Describing your image…"):
-        caps = captioner(img)
-        first = caps[0]
-        st.session_state.caption = first.get("generated_text", "") if isinstance(first, dict) else str(first)
-st.write("**Caption:**", st.session_state.caption)
+    # Generate TTS audio once
+    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
 
-# 5) Generate story once
-if "story" not in st.session_state:
-    with st.spinner("Weaving a story…"):
-        out = story_gen(
-            st.session_state.caption,
-            max_length=200,
-            do_sample=True,
-            top_p=0.9,
-            num_return_sequences=1
+    # 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.session_state.story = out[0]["generated_text"]
-st.write("**Story:**", st.session_state.story)
 
-# 6) Generate Bark audio once
-if "audio_data" not in st.session_state:
-    with st.spinner("Converting to speech…"):
-        speech = tts_pipe(st.session_state.story)
-        # Bark returns a dict: {"audio": np.ndarray, "sampling_rate": int}
-        st.session_state.audio_data = speech["audio"]
-        st.session_state.audio_sr   = speech["sampling_rate"]
 
-# 7) Play 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
-    )
+if __name__ == "__main__":
+    main()