Update app.py

app.py

@@ -17,9 +17,9 @@ except ImportError:
 # Load the image captioning model
 caption_model = pipeline("image-to-text", model="unography/blip-large-long-cap")
 
-#story_generator = pipeline("text-generation", model="distilbert/distilgpt2")
+story_generator = pipeline("text-generation", model="distilbert/distilgpt2")
 
-story_generator = pipeline("text-generation", model="isarth/distill_gpt2_story_generator")
+#story_generator = pipeline("text-generation", model="isarth/distill_gpt2_story_generator")
 
 def generate_caption(image):
     # Generate the caption for the uploaded image
@@ -28,18 +28,18 @@ def generate_caption(image):
 
 def generate_story(caption):
     # Generate the story based on the caption using the GPT-2 model
-    prompt = f"Once upon a time, in a world inspired by the image of {caption}, a magical story took place. The fairy tale, suitable for children aged 3-10, unfolds as follows:\n\nIntroduction: Set the scene and introduce the main characters.\n\nConflict: Present a problem or challenge the characters face.\n\nResolution: Describe how the characters solve the problem or overcome the challenge.\n\nMoral: Conclude with a positive message or lesson learned.\n\nThe story should be well-structured, engaging, and easy to understand for young children. Let's begin:\n\n"
+    prompt = f"Once upon a time, in a world inspired by the image of {caption}, a delightful children's story took place. The story, suitable for ages 3-10, goes like this:\n\nIntroduction (1-2 sentences): Introduce the main character(s) and the setting.\n\nBeginning (2-3 sentences): Describe the character's normal life or routine.\n\nMiddle (3-4 sentences): Present a problem or challenge the character faces.\n\nEnd (2-3 sentences): Show how the character solves the problem or learns a lesson.\n\nThe story should be simple, engaging, and convey a positive message. Let's begin the tale:\n\n"
     story = story_generator(prompt, max_length=500, num_return_sequences=1)[0]["generated_text"]
     
     # Extract the story text from the generated output
     story_parts = story.split("\n\n")
-    if len(story_parts) > 5:
-        story = "\n\n".join(story_parts[5:]).strip()
+    if len(story_parts) > 7:
+        story = "\n\n".join(story_parts[7:]).strip()
     else:
         story = "\n\n".join(story_parts).strip()
     
     # Post-process the story (example: remove inappropriate words)
-    inappropriate_words = ["violence", "horror", "scary"]
+    inappropriate_words = ["violence", "horror", "scary", "adult", "death", "gun", "shoot"]
     for word in inappropriate_words:
         story = story.replace(word, "")