Update app.py

app.py

@@ -26,14 +26,23 @@ def generate_caption(image):
     caption = image_to_text(image)[0]['generated_text']
     return caption
 
-# Function to generate story from caption
+# Function to generate story from caption (less than 100 words)
 def generate_story(caption):
     prompt = f"Once upon a time, {caption} "
-    story = story_generator(prompt, max_length=200, do_sample=True)[0]['generated_text']
-    # Ensure the story is at least 100 words
-    while len(story.split()) < 100:
-        additional_text = story_generator(story, max_length=100, do_sample=True)[0]['generated_text']
-        story += additional_text
+    
+    # Set max_length to control story length (approximately 100 words)
+    # Typical English word is ~5 characters, so ~500 characters ≈ 100 words
+    story = story_generator(prompt, max_length=100, do_sample=True)[0]['generated_text']
+    
+    # Ensure story doesn't exceed 100 words
+    words = story.split()
+    if len(words) > 100:
+        words = words[:100]
+        story = " ".join(words)
+        # Add period to the end if needed
+        if not story.endswith(('.', '!', '?')):
+            story += '.'
+    
     return story
 
 # Function to convert text to speech
@@ -60,7 +69,8 @@ if uploaded_file is not None:
             
             # Generate story
             story = generate_story(caption)
-            st.write("### Your Story")
+            word_count = len(story.split())
+            st.write(f"### Your Story ({word_count} words)")
             st.write(story)
             
             # Generate audio