Update app.py

app.py

@@ -60,16 +60,19 @@ def generate_caption(image_bytes):
 # Function to generate a story from a caption
 def generate_story(caption):
     try:
-        prompt = [f"Based on this text:'{caption}', ellaborate a short story with no more than 100 words for young children."]
+        device = "cpu"
+        prompt = [{"role": "user", 
+                   "content": f"Based on this text:'{caption}', ellaborate a short story with no more than 100 words for young children."}]
+        input_text = tokenizer.apply_chat_template(prompt, tokenize=False)
+        inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
         story_output = story_generator(
-            prompt,
+            inputs,
             max_length=150,
             do_sample=True,
             temperature=0.2,
-            top_p=0.9,
-            num_return_sequences=1
+            top_p=0.9
         )
-        story = story_output[0]["generated_text"]
+        story = tokenizer.decode(story_output[0])  #story_output[0]["generated_text"]
         story = truncate_to_100_words(story)
         return story
     except Exception as e: