Update app.py

app.py

@@ -54,11 +54,8 @@ def load_models():
                 model="aspis/gpt2-genre-story-generation"
             )
         except Exception as e:
-            logging.warning(f"Failed to load aspis/gpt2-genre-story-generation: {e}. Falling back to gpt2.")
-            st.session_state.storyer = pipeline(
-                "text-generation",
-                model="gpt2"
-            )
+            logging.error(f"Failed to load aspis/gpt2-genre-story-generation: {e}")
+            raise Exception("Text generation model could not be loaded.")
     if "tts" not in st.session_state:
         st.session_state.tts = pipeline(
             "text-to-speech",
@@ -80,48 +77,55 @@ def get_caption(image_bytes):
 @st.cache_data(show_spinner=False)
 def get_story(caption):
     """
-    Generates a humorous and engaging children's story based on the caption.
-    Uses a simplified prompt and robust output parsing.
+    Generates a humorous children's story (50-100 words) based on the caption.
+    Optimized for faster generation with aspis/gpt2-genre-story-generation.
     """
     prompt = (
-        f"Create a funny, warm children's story (50-100 words) for ages 3-10 based on: {caption}. "
-        f"Use third-person narrative, as if playfully describing the scene."
+        f"Write a funny, warm children's story (50-100 words) for ages 3-10 based on: {caption}. "
+        f"Third-person narrative, playful tone."
     )
     try:
         result = st.session_state.storyer(
             prompt,
-            max_new_tokens=150,    # Increased to allow more room for story
+            max_new_tokens=80,     # Reduced for faster generation
             do_sample=True,
-            temperature=0.8,       # Slightly higher for creativity
-            top_p=0.9,
+            temperature=0.7,       # Balanced creativity
+            top_k=40,              # Faster sampling with smaller k
+            top_p=0.85,            # Tighter sampling for coherence
             return_full_text=False
         )
         logging.info(f"Story generation raw result: {result}")
         
-        # Extract generated text
+        # Extract and clean generated text
         raw_story = result[0].get("generated_text", "").strip()
         
-        # If no meaningful output, generate a fallback story
-        if not raw_story or len(raw_story.split()) < 10:
+        # Ensure story is 50-100 words
+        words = raw_story.split()
+        if len(words) < 50 or not raw_story:
             logging.warning("Generated story too short or empty. Using fallback.")
             raw_story = (
-                f"Once upon a time, in a land of {caption}, a silly squirrel named Sammy "
-                f"found a shiny treasure! He danced with joy, but oh no! It was a magic acorn! "
-                f"It grew into a giant tree, and Sammy climbed to the top, giggling all the way. "
-                f"The tree sang funny songs, and all the animals joined in for a big party!"
+                f"In a land of {caption}, a silly bunny named Bouncy found a shiny star! "
+                f"It sparkled, making Bouncy giggle and hop high. The star said, 'Dance!' "
+                f"So Bouncy twirled with squirrels and birds. They threw a forest party, "
+                f"singing silly songs under the twinkling sky, laughing all night."
             )
+            words = raw_story.split()
         
-        # Truncate to 100 words
-        words = raw_story.split()
+        # Truncate to 100 words if too long
         story = " ".join(words[:100])
+        # Pad if too short
+        if len(words) < 50:
+            story += " And they all lived happily ever after!"
+        
         return story
     except Exception as e:
         logging.error(f"Story generation failed: {e}")
-        # Fallback story in case of errors
+        # Fallback story
         return (
-            f"Once upon a time, in a land of {caption}, a silly squirrel named Sammy "
-            f"found a shiny treasure! He danced with joy, but oh no! It was a magic acorn! "
-            f"It grew into a giant tree, and Sammy climbed to the top, giggling all the way."
+            f"In a land of {caption}, a silly bunny named Bouncy found a shiny star! "
+            f"It sparkled, making Bouncy giggle and hop high. The star said, 'Dance!' "
+            f"So Bouncy twirled with squirrels and birds. They threw a forest party, "
+            f"singing silly songs under the twinkling sky, laughing all night."
         )
 
 @st.cache_data(show_spinner=False)