Update app.py

app.py

@@ -33,16 +33,14 @@ st.markdown(
     unsafe_allow_html=True
 )
 
-# ------------------ Lazy Model Loading ------------------
+# ------------------ Model Loading ------------------
 def load_models():
     """
-    Lazy-load the required pipelines and store them in session state.
-    This avoids heavy model initialization during container build.
-    
+    Lazy-load the pipelines and store them in session state.
     Pipelines:
-      1. Captioner: Turns an image into descriptive text.
+      1. Captioner: Generates descriptive text from an image.
       2. Storyer: Generates a humorous children's story using aspis/gpt2-genre-story-generation.
-      3. TTS: Converts text into audio.
+      3. TTS: Converts text into spoken audio.
     """
     if "captioner" not in st.session_state:
         st.session_state.captioner = pipeline(
@@ -60,32 +58,31 @@ def load_models():
             model="facebook/mms-tts-eng"
         )
 
-# ------------------ Helper Functions ------------------
-def generate_caption(image):
+# ------------------ Caching Functions ------------------
+@st.cache_data(show_spinner=False)
+def get_caption(image_bytes):
     """
-    Generate a caption from the image using the captioner model.
+    Convert the image bytes into a lower resolution image to speed up captioning,
+    then generate and return the caption.
     """
-    caption_results = st.session_state.captioner(image)
-    return caption_results[0]["generated_text"]
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    # Resize the image to speed up processing (keeping aspect ratio)
+    image.thumbnail((384, 384))
+    caption = st.session_state.captioner(image)[0]["generated_text"]
+    return caption
 
-def generate_story(caption):
+@st.cache_data(show_spinner=False)
+def get_story(caption):
     """
-    Generate a humorous and engaging children's story based on the caption.
-    
-    The prompt instructs the model to write a funny, warm, and imaginative story for ages 3-10 
-    (50-100 words) in a playful authorial tone. It explicitly requests details such as:
-      - A clear venue/location (e.g., a park, school, or home),
-      - Specific characters (e.g., a little girl named Lily or a boy named Jack),
-      - Humorous actions.
+    Generate a humorous and engaging children's story using the caption.
     """
     prompt = (
         f"Write a funny, warm, and imaginative children's story for ages 3-10, 50-100 words, "
         f"in third-person narrative, as if the author is playfully describing the scene in the image: {caption}. "
-        "Explicitly mention the exact venue or location (such as a park, school, or home), describe specific characters "
+        "Explicitly mention the exact venue or location (e.g. a park, school, or home), describe specific characters "
         "(for example, a little girl named Lily or a boy named Jack), and detail the humorous actions they perform. "
         "Ensure the story is playful, engaging, and ends with a complete sentence."
     )
-    # Using sampling parameters that work well with the smaller GPT-2 genre model.
     raw_story = st.session_state.storyer(
         prompt,
         max_new_tokens=100,
@@ -96,11 +93,11 @@ def generate_story(caption):
     words = raw_story.split()
     return " ".join(words[:100])
 
-def generate_audio(story):
+@st.cache_data(show_spinner=False)
+def get_audio(story):
     """
-    Convert the story text to audio.
-    The text is split into 300-character chunks (to reduce the number of TTS calls),
-    the audio chunks are concatenated, and then saved into an in‑memory WAV buffer.
+    Convert the generated story text into audio.
+    The text is split into 300-character chunks to reduce repeated TTS calls.
     """
     chunks = textwrap.wrap(story, width=300)
     audio_chunks = [st.session_state.tts(chunk)["audio"].squeeze() for chunk in chunks]
@@ -114,31 +111,30 @@ def generate_audio(story):
 uploaded_file = st.file_uploader("Choose a Picture...", type=["jpg", "jpeg", "png"])
 if uploaded_file is not None:
     try:
-        # Open and display the uploaded image
-        image = Image.open(uploaded_file).convert("RGB")
+        load_models()  # Ensure models are loaded once
+        image_bytes = uploaded_file.getvalue()
+        # Display image
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         st.image(image, caption="Your Amazing Picture!", use_column_width=True)
         st.markdown("<h3 style='text-align: center;'>Ready for your story?</h3>", unsafe_allow_html=True)
-        
         if st.button("Story, Please!"):
-            with st.spinner("Loading models and processing your picture..."):
-                load_models()
-                caption = generate_caption(image)
+            with st.spinner("Generating caption..."):
+                caption = get_caption(image_bytes)
             st.markdown("<h3 style='text-align: center;'>Caption:</h3>", unsafe_allow_html=True)
             st.write(caption)
             
-            with st.spinner("Creating a fun story just for you..."):
-                story = generate_story(caption)
+            with st.spinner("Generating story..."):
+                story = get_story(caption)
             st.markdown("<h3 style='text-align: center;'>Your Story:</h3>", unsafe_allow_html=True)
             st.write(story)
             
-            with st.spinner("Turning your story into a magical song..."):
-                audio_buffer = generate_audio(story)
+            with st.spinner("Generating audio..."):
+                audio_buffer = get_audio(story)
             st.audio(audio_buffer, format="audio/wav", start_time=0)
             st.markdown(
                 "<p style='text-align: center; font-weight: bold;'>Enjoy your magical story! 🎶</p>",
                 unsafe_allow_html=True
             )
-    
     except Exception as e:
         st.error("Oops! Something went wrong. Please try a different picture or check the file format!")
         st.error(f"Error details: {e}")