Update app.py

app.py

@@ -75,7 +75,7 @@ def get_caption(image_bytes):
 def get_story(caption):
     """
     Generates a humorous and engaging children's story based on the caption.
-    Uses a prompt to instruct the model and limits token generation to 80 tokens.
+    Uses a prompt to instruct the model and limits token generation.
     """
     prompt = (
         f"Write a funny, warm, and imaginative children's story for ages 3-10, 50-100 words, "
@@ -83,22 +83,28 @@ def get_story(caption):
     )
     result = st.session_state.storyer(
         prompt,
-        max_new_tokens=80,
+        max_new_tokens=120,    # Increased from 80 to 120 for more continuation space
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
         return_full_text=False
     )
     
-    # Log the raw result for debugging (this is viewable in the server logs)
+    # Log the raw result for debugging (viewable in server logs)
     print("Story generation raw result:", result)
     
     raw_story = result[0].get("generated_text", "").strip()
     
-    # Remove the prompt from the output if it is included.
+    # If the generated text starts with the prompt, remove it only if there is substantial extra content.
     if raw_story.startswith(prompt):
-        raw_story = raw_story[len(prompt):].strip()
-    
+        # Compute the extra part after the prompt.
+        extra_text = raw_story[len(prompt):].strip()
+        # Only use the extra text if it is longer than a threshold (e.g. 20 characters).
+        if len(extra_text) > 20:
+            raw_story = extra_text
+        else:
+            # If not, use the full raw_story instead.
+            raw_story = raw_story
     words = raw_story.split()
     story = " ".join(words[:100])
     return story
@@ -107,8 +113,8 @@ def get_story(caption):
 def get_audio(story):
     """
     Converts the generated story text into audio.
-    Splits the text into 300-character chunks, processes each via the TTS pipeline,
-    and concatenates the resulting audio arrays. If no audio is generated, 1 second of silence is used.
+    Splits the text into 300-character chunks to reduce repeated TTS calls.
+    Checks each chunk; if no valid audio is produced, creates 1 second of silence.
     """
     chunks = textwrap.wrap(story, width=300)
     audio_chunks = []
@@ -153,7 +159,6 @@ if uploaded_file is not None:
             with st.spinner("Generating story..."):
                 story = get_story(caption)
             st.markdown("<h3 style='text-align: center;'>Your Story:</h3>", unsafe_allow_html=True)
-            # If the story is empty (or consists only of whitespace), display a default message.
             if not story.strip():
                 st.write("No story was generated. Please try again.")
             else: