Update app.py

app.py

@@ -5,7 +5,7 @@ from gtts import gTTS
 
 st.set_page_config(page_title="Image to Audio Story", page_icon="🦜")
 caption_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-story_pipeline = pipeline("text-generation", model="Qwen/Qwen2.5-0.5B-Instruct")
+story_pipeline = pipeline("text-generation", model="Qwen/Qwen2-1.5B")
 
 
 def extract_image_caption(image_data):
@@ -15,22 +15,27 @@ def extract_image_caption(image_data):
 
 
 def compose_story_from_caption(caption_detail):
-    prompt_text = (
-        "You are a talented and imaginative storyteller for children aged 3 to 10. "
-        "Using the details derived from the image below, craft a captivating tale that includes three main characters, "
-        "an adventurous journey, and delightful surprises. "
-        "Your story should be vivid, original, and between 90 and 100 words in length.\n\n"
-        f"Image Details: {caption_detail}\n\nStory:"
-    )
-    
-    # 生成故事，设置最大长度为150，以确保有足够空间生成90-100词的内容
-    story_results = story_pipeline(prompt_text, num_return_sequences=1, max_length=150)
-    story_text = story_results[0]['generated_text']
-    
-    # 提取故事文本
-    if "Story:" in story_text:
-        return story_text.split("Story:", 1)[1].strip()
-    return story_text.strip()
+    while True:
+        prompt_text = (
+            "You are a talented and imaginative storyteller for children aged 3 to 10. "
+            "Using the details derived from the image below, craft a complete and captivating tale that includes three main characters, "
+            "an adventurous journey, and delightful surprises. "
+            "Your story should have a clear beginning, middle, and end, and be between 80 and 100 words in length.\n\n"
+            f"Image Details: {caption_detail}\n\nStory:"
+        )
+        
+        story_results = story_pipeline(prompt_text, num_return_sequences=1, max_length=150)
+        story_text = story_results[0]['generated_text']
+        
+        # 提取故事
+        if "Story:" in story_text:
+            story = story_text.split("Story:", 1)[1].strip()
+        else:
+            story = story_text.strip()
+        
+        # 检查长度
+        if 80 <= len(story.split()) <= 100:
+            return story
     
 def convert_text_to_audio(text_content, audio_path="output.mp3"):
     """