Update app.py

app.py

@@ -59,22 +59,59 @@ def img2text(image):
     text = image_to_text(image)[0]["generated_text"]
     return text
 
-# Improved text-to-story function with natural ending
+# Helper function to count words
+def count_words(text):
+    return len(text.split())
+
+# Improved text-to-story function without "Once upon a time" constraint
 def text2story(text):
     generator = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-    prompt = f"Write a short children's story based on this: {text}. The story should have a clear beginning, middle, and end. Keep it under 150 words. Once upon a time, "
+    
+    # Ask for a story without specifying how to start
+    prompt = f"""Write a children's story based on this: {text}. 
+    The story should have a clear beginning, middle, and end.
+    Make the story approximately 150-200 words long with descriptive language.
+    """
     
     # Generate a longer text to ensure we get a complete story
     story_result = generator(
         prompt,
-        max_length=300,
+        max_length=500,
         num_return_sequences=1,
         temperature=0.7,
         do_sample=True
     )
    
-    story_text = story_result[0]['generated_text']
-    story_text = story_text.replace(prompt, "Once upon a time, ")
+    full_text = story_result[0]['generated_text']
+    
+    # Try to extract just the story part (after the prompt)
+    # Look for paragraph breaks or clear story beginnings
+    potential_starts = [
+        "\n\n",
+        "\n",
+        ". ",
+        "! ",
+        "? "
+    ]
+    
+    # Find where the prompt ends and the actual story begins
+    story_text = full_text
+    
+    # First remove the exact prompt if it appears verbatim
+    if prompt in story_text:
+        story_text = story_text.replace(prompt, "")
+    else:
+        # Look for paragraph breaks or sentence endings that might indicate
+        # where the prompt instructions end and the story begins
+        for start_marker in potential_starts:
+            if start_marker in story_text:
+                parts = story_text.split(start_marker, 1)
+                if len(parts[0]) < len(story_text) * 0.5:  # If the first part is reasonably short
+                    story_text = parts[1]
+                    break
+    
+    # Clean up any leading/trailing whitespace
+    story_text = story_text.strip()
     
     # Find natural ending points (end of sentences)
     periods = [i for i, char in enumerate(story_text) if char == '.']
@@ -84,20 +121,39 @@ def text2story(text):
     # Combine all ending punctuation and sort
     all_endings = sorted(periods + question_marks + exclamation_marks)
     
+    # Target approximately 100 words
+    target_word_count = 100
+    min_acceptable_words = 80
+    
     # If we have any sentence endings
     if all_endings:
-        # Get the index where the story should reasonably end (after at least 100 characters)
-        min_story_length = 100
-        suitable_endings = [i for i in all_endings if i >= min_story_length]
+        # Find the sentence ending that gets us closest to 100 words
+        closest_ending = None
+        closest_word_diff = float('inf')
         
-        if suitable_endings:
-            # Find an ending that completes a thought (not just the first sentence)
-            if len(suitable_endings) > 2:
-                # Use the third sentence ending or later for a more complete story
-                return story_text[:suitable_endings[2]+1]
-            else:
-                # If we don't have many sentences, use the last one we found
-                return story_text[:suitable_endings[-1]+1]
+        for ending_idx in all_endings:
+            candidate_text = story_text[:ending_idx+1]
+            candidate_word_count = count_words(candidate_text)
+            
+            # Only consider endings that give us at least min_acceptable_words
+            if candidate_word_count >= min_acceptable_words:
+                word_diff = abs(candidate_word_count - target_word_count)
+                
+                if word_diff < closest_word_diff:
+                    closest_ending = ending_idx
+                    closest_word_diff = word_diff
+        
+        # If we found a suitable ending, use it
+        if closest_ending is not None:
+            return story_text[:closest_ending+1]
+    
+    # If we couldn't find a good ending near 100 words, but we have some sentence endings,
+    # use the last one that results in a story with at least min_acceptable_words words
+    if all_endings:
+        for ending_idx in reversed(all_endings):
+            candidate_text = story_text[:ending_idx+1]
+            if count_words(candidate_text) >= min_acceptable_words:
+                return candidate_text
     
     # If no good ending is found, return as is
     return story_text
@@ -121,7 +177,10 @@ if uploaded_file is not None:
     # Text to Story
     with st.spinner("Creating story..."):
         story = text2story(caption)
-    st.write(f"Story: {story}")
+        # Display word count for transparency
+        word_count = len(story.split())
+        st.write(f"Story ({word_count} words):")
+        st.write(story)
     
     # Text to Audio
     with st.spinner("Generating audio..."):