Update app.py

app.py

@@ -25,14 +25,19 @@ def generate_caption(image):
     caption = caption_model(image)[0]["generated_text"]
     return caption
 
-def generate_story(caption):
+#def generate_story(caption):
     # Generate the story based on the caption using the GPT-2 model
-    prompt = f"Starting with 'Once upon a time', based on the image described as '{caption}', here is a short and interesting story for children aged 3-10. The story is positive and happy in tone, with added imagination:\n\n"
-    story = story_generator(prompt, max_length=500, num_return_sequences=1)[0]["generated_text"]
+    #prompt = f"Starting with 'Once upon a time', based on the image described as '{caption}', here is a short and interesting story for children aged 3-10. The story is positive and happy in tone, with added imagination:\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 = story.split("\n\n")[1].strip()
     
+    #return story
+
+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"]
     return story
 
 def convert_to_audio(story):