Update app.py

app.py

@@ -16,39 +16,10 @@ def img2text(url):
 
 # text2story
 def text2story(text):
-    prompt_template = f"Write a children's story for ages 3-10 based on: {text}"
-        # 生成完整提示词（避免重复插入）
-    full_prompt = prompt_template.format(text=text)
-    
-    # 初始化生成管道
-    pipe = pipeline(
-        "text-generation", 
-        model="pranavpsv/genre-story-generator-v2",
-        max_new_tokens=180,
-        min_new_tokens=130,
-        temperature=0.7,
-        pad_token_id=50256  # 添加模型专用参数
-    )
-    
-    # 生成原始文本
-    raw_output = pipe(full_prompt, return_full_text=False)[0]['generated_text']
-    
-    # 增强版清洗逻辑
-    def clean_output(generated_text):
-        # 使用正则表达式匹配提示词块
-        prompt_pattern = re.compile(
-            r'\[PROMPT_START\].*?\[PROMPT_END\]', 
-            re.DOTALL  # 匹配多行内容
-        )
-        # 移除整个提示词块
-        cleaned = re.sub(prompt_pattern, '', generated_text)
-        
-        # 二次清理残留内容
-        cleaned = re.sub(r'^Write a children.*?\n', '', cleaned)  # 处理可能的开头残留
-        cleaned = re.sub(r'Requirements:.*?\n', '', cleaned)      # 移除要求残留
-        return cleaned.strip()
-
-    return clean_output(raw_output)
+    prompt = f"Write a children's story for ages 3-10 based on: {text}"
+    pipe = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2", max_new_tokens=160, min_new_tokens=130, num_return_sequences=1)
+    story_text = pipe(prompt, return_full_text=False)[0]['generated_text']
+    return story_text
 
 
 # text2audio