Update app.py

app.py

@@ -17,7 +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")
 
 def generate_caption(image):
     # Generate the caption for the uploaded image
@@ -26,11 +28,15 @@ def generate_caption(image):
 
 def generate_story(caption):
     # Generate the story based on the caption using the GPT-2 model
-    prompt = f"Imagine a delightful children's fairy tale inspired by the image described as '{caption}'. Itnshould be interesting, easy to understand, and use age-appropriate language suitable for children aged 3-10. Let the magical story unfold:\n\n"
+    prompt = f"Imagine a delightful children's fairy tale inspired by the image described as '{caption}'. The story should be enchanting, easy to understand, and use age-appropriate language suitable for children aged 3-10. Let the magical story unfold:\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 = story.split("\n\n")[1].strip()
+    story_parts = story.split("\n\n")
+    if len(story_parts) > 1:
+        story = "\n\n".join(story_parts[1:]).strip()
+    else:
+        story = story_parts[0].strip()
     
     # Post-process the story (example: remove inappropriate words)
     inappropriate_words = ["violence", "horror", "scary"]