Update app.py

app.py

@@ -33,14 +33,14 @@ st.markdown(
     unsafe_allow_html=True
 )
 
-# ------------------ Model Loading ------------------
+# ------------------ Lazy Model Loading ------------------
 def load_models():
     """
-    Lazy-load the pipelines and store them in session state.
+    Lazy-load the required pipelines and store them in session state.
     Pipelines:
       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 spoken audio.
+      3. TTS: Converts text into audio.
     """
     if "captioner" not in st.session_state:
         st.session_state.captioner = pipeline(
@@ -62,11 +62,10 @@ def load_models():
 @st.cache_data(show_spinner=False)
 def get_caption(image_bytes):
     """
-    Convert the image bytes into a lower resolution image to speed up captioning,
-    then generate and return the caption.
+    Convert the image bytes into a lower resolution image and generate the caption.
     """
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    # Resize the image to speed up processing (keeping aspect ratio)
+    # Resize the image (preserving aspect ratio) to speed up processing.
     image.thumbnail((384, 384))
     caption = st.session_state.captioner(image)[0]["generated_text"]
     return caption
@@ -75,21 +74,33 @@ def get_caption(image_bytes):
 def get_story(caption):
     """
     Generate a humorous and engaging children's story using the caption.
+    The prompt instructs the model to write a playful story for ages 3-10 (50-100 words) that includes:
+      - The venue or location (e.g. park, school, or home),
+      - Specific characters (like a little girl named Lily or a boy named Jack),
+      - Humorous actions.
+    The generated text should not include the prompt.
     """
     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 (e.g. a park, school, or home), describe specific characters "
+        "Explicitly mention the exact venue or location (such as 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."
     )
+    # Set return_full_text=False so the prompt isn't included in the output.
     raw_story = st.session_state.storyer(
         prompt,
         max_new_tokens=100,
         do_sample=True,
         temperature=0.7,
-        top_p=0.9
+        top_p=0.9,
+        return_full_text=False
     )[0]["generated_text"].strip()
+    
+    # Fallback check: if the generated text starts with the prompt, remove it.
+    if raw_story.startswith(prompt):
+        raw_story = raw_story[len(prompt):].strip()
+    
     words = raw_story.split()
     return " ".join(words[:100])
 
@@ -97,7 +108,8 @@ def get_story(caption):
 def get_audio(story):
     """
     Convert the generated story text into audio.
-    The text is split into 300-character chunks to reduce repeated TTS calls.
+    The text is split into 300-character chunks to reduce repeated TTS calls,
+    the resulting audio arrays are concatenated, and the output is stored in an in‑memory WAV buffer.
     """
     chunks = textwrap.wrap(story, width=300)
     audio_chunks = [st.session_state.tts(chunk)["audio"].squeeze() for chunk in chunks]
@@ -111,12 +123,13 @@ def get_audio(story):
 uploaded_file = st.file_uploader("Choose a Picture...", type=["jpg", "jpeg", "png"])
 if uploaded_file is not None:
     try:
-        load_models()  # Ensure models are loaded once
+        load_models()  # Ensure models are loaded
         image_bytes = uploaded_file.getvalue()
-        # Display image
+        # Display the 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("Generating caption..."):
                 caption = get_caption(image_bytes)