Update app.py

app.py

@@ -18,7 +18,8 @@ except ImportError:
 caption_model = pipeline("image-to-text", model="unography/blip-large-long-cap")
 
 # Load the GPT-2 model for story generation
-story_generator = pipeline("text-generation", model="gpt2")
+#story_generator = pipeline("text-generation", model="gpt2")
+story_generator = pipeline("text-generation", model="distilbert/distilgpt2")
 
 def generate_caption(image):
     # Generate the caption for the uploaded image
@@ -35,9 +36,24 @@ def generate_caption(image):
     
     #return story
 
-def generate_story(caption):
+#def generate_story(caption):
     # Generate the story based on the caption
-    story = story_generator(caption, max_length=200, num_return_sequences=1)[0]["generated_text"]
+    #story = story_generator(caption, max_length=200, num_return_sequences=1)[0]["generated_text"]
+    #return story
+
+def generate_story(caption):
+    # Generate the story based on the caption using the GPT-2 model
+    prompt = f"Once upon a time, based on the image described as '{caption}', here is a short, simple, and engaging story for children aged 3-10. The story should be easy to understand, use age-appropriate language, and convey a positive message. Focus on the main elements in the image and create a story that sparks their imagination:\n\n"
+    story = story_generator(prompt, max_length=100, num_return_sequences=1)[0]["generated_text"]
+    
+    # Extract the story text from the generated output
+    story = story.split("\n\n")[1].strip()
+    
+    # Post-process the story (example: remove inappropriate words)
+    inappropriate_words = ["violence", "horror", "scary"]
+    for word in inappropriate_words:
+        story = story.replace(word, "")
+    
     return story
 
 def convert_to_audio(story):