Update app.py

app.py

@@ -13,6 +13,7 @@ from datetime import datetime
 import spaces
 from kokoro import KModel, KPipeline
 import soundfile as sf
+import math
 
 def clear_memory():
     """Helper function to clear both CUDA and system memory, safe for Spaces environment"""
@@ -153,18 +154,55 @@ def analyze_image(image):
 
 @torch.inference_mode()
 @spaces.GPU(duration=30)
+def format_story_paragraphs(story_text, max_paragraphs=9):
+    """Formats the story by evenly distributing sentences over a max of 9 paragraphs."""
+    
+    # Remove unwanted tokens and artifacts
+    story_text = story_text.replace("<|im_end|>", "").strip()
+    
+    # Split into sentences
+    sentences = re.split(r'(?<=[.!?])\s+', story_text)
+    sentences = [s.strip() for s in sentences if s.strip()]  # Remove empty sentences
+    
+    # Determine optimal sentence distribution
+    total_sentences = len(sentences)
+    num_paragraphs = min(max_paragraphs, total_sentences)  # Ensure we do not exceed max paragraphs
+    sentences_per_paragraph = math.ceil(total_sentences / num_paragraphs)  # Distribute evenly
+    
+    # Group sentences into paragraphs
+    paragraphs = []
+    current_paragraph = []
+    
+    for sentence in sentences:
+        current_paragraph.append(sentence)
+        
+        # Ensure even distribution across paragraphs
+        if len(current_paragraph) == sentences_per_paragraph:
+            paragraphs.append(" ".join(current_paragraph))
+            current_paragraph = []
+        
+        # Stop if we reach max paragraphs
+        if len(paragraphs) == max_paragraphs:
+            break
+
+    # If there are leftover sentences, add them as a final paragraph
+    if current_paragraph:
+        paragraphs.append(" ".join(current_paragraph))
+
+    # Join paragraphs with a double newline
+    return "\n\n".join(paragraphs)
+
 def generate_story(image_description):
     clear_memory()
     
-    story_prompt = f"""Write a short children's story (about 500 words) based on this scene: {image_description}
+    story_prompt = f"""Write a short children's story (about 200 words) based on this scene: {image_description}
 
     Requirements:
     1. Main character: An English bulldog named Champ
     2. Include these values: confidence, teamwork, caring, and hope
     3. Theme: "Doing the right thing is important"
     4. Keep it simple and engaging for young children
-    5. End with a simple moral lesson
-    6. Each paragraph needs to be three sentences or less for readability"""
+    5. End with a simple moral lesson"""
 
     try:
         messages = [{"role": "user", "content": story_prompt}]
@@ -181,9 +219,10 @@ def generate_story(image_description):
             repetition_penalty=1.2
         )
         
+        # Decode and format the story
         story = tokenizer_lm.decode(outputs[0])
-        story = clean_story_output(story)
-        
+        story = format_story_paragraphs(story)  # Evenly distribute sentences
+
         clear_memory()
         return story
         
@@ -192,6 +231,7 @@ def generate_story(image_description):
         clear_memory()
         return "Error generating story. Please try again."
 
+
 @torch.inference_mode()
 @spaces.GPU(duration=30)
 def generate_image_prompts(story_text):