Update app.py

app.py

@@ -27,11 +27,11 @@ def generate_caption(image):
 
 def generate_story(caption):
     # Generate the story based on the caption using the GPT-2 model
-    prompt = f"Could you please suggest a story or fairytale in 100 words based on the {caption}\n\n"
-    story = story_generator(prompt, max_length=500, num_return_sequences=1)[0]["generated_text"]
+    prompt = f"Write a short, simple children's story in 100 words based on the following image description:\n\n{caption}\n\nStory:"
+    story = story_generator(prompt, max_length=200, num_return_sequences=1)[0]["generated_text"]
     
     # Extract the story text from the generated output
-    story = story.strip()
+    story = story.split("Story:")[1].strip()
     
     # Post-process the story (example: remove inappropriate words)
     inappropriate_words = ["violence", "horror", "scary", "adult", "death", "gun", "shoot"]