Update app.py

app.py

@@ -17,10 +17,12 @@ 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="TheBloke/storytime-13B-GPTQ")
+
 def generate_caption(image):
     # Generate the caption for the uploaded image
     caption = caption_model(image)[0]["generated_text"]
@@ -28,15 +30,11 @@ 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 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"
+    prompt = f"Write a short, simple children's story inspired by the image of {caption}. Here's the story:\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) > 7:
-        story = "\n\n".join(story_parts[7:]).strip()
-    else:
-        story = "\n\n".join(story_parts).strip()
+    story = story.strip()
     
     # Post-process the story (example: remove inappropriate words)
     inappropriate_words = ["violence", "horror", "scary", "adult", "death", "gun", "shoot"]