Update app.py

app.py

@@ -19,6 +19,13 @@ def load_pipelines():
 
 captioner, storyer, tts = load_pipelines()
 
+# Function to extract keywords from caption
+def extract_keywords(caption):
+    stop_words = {'the', 'a', 'an', 'in', 'on', 'at', 'of', 'to', 'is', 'are', 'with', 'and'}
+    words = caption.lower().split()
+    keywords = [word.strip(".,!?\"'") for word in words if word not in stop_words and len(word) > 2]
+    return keywords
+
 # Function part
 # Function to generate content from an image
 def generate_content(image):
@@ -29,11 +36,18 @@ def generate_content(image):
     st.write("**🌟 What's in the picture: 🌟**")
     st.write(caption)
 
-    # Create prompt for story
+    # Extract keywords from the caption
+    keywords = extract_keywords(caption)
+    keywords_str = ", ".join(keywords)
+
+    # Create prompt for story with narrative structure
     prompt = (
         f"Write a funny, warm children's story for ages 3-10, 50–100 words, "
-        f"in third-person narrative, that describes this scene exactly: {caption} "
-        f"mention the exact place, location or venue within {caption}"
+        f"in third-person narrative. The story must describe this scene exactly: {caption}. "
+        f"Include these keywords from the caption: {keywords_str}. "
+        f"Feature a named protagonist (a child), a simple problem or goal, and a happy resolution. "
+        f"Mention the exact place, location, or venue from the scene (e.g., park, pool, gym). "
+        f"Keep the tone playful and engaging."
     )
     
     # Generate raw story