Update app.py

app.py

@@ -1,4 +1,4 @@
-# storygen_tts_final.py
+# storygen_tts_final.py (修正版)
 import streamlit as st
 from transformers import (
     BlipForConditionalGeneration,
@@ -69,8 +69,8 @@ Story:"""
     # 使用GPT-2生成故事
     generated = story_generator(
         prompt,
-        max_length=100,
-        min_length=50,
+        max_length=300,
+        min_length=150,
         num_return_sequences=1,
         temperature=0.85,
         repetition_penalty=2.0
@@ -79,15 +79,17 @@ Story:"""
     # 提取生成文本并清理
     full_text = generated[0]['generated_text']
     story = full_text.split("Story:")[-1].strip()
-    return story[:600].replace(caption, "").strip()
+    return story[:500].replace(caption, "").strip()  # 缩短最大长度
 
 def text_to_speech(text, processor, model, vocoder, embeddings_dataset):
     """文本转语音"""
     try:
+        # 添加文本截断和长度限制
         inputs = processor(
             text=text, 
             return_tensors="pt",
-            voice_preset=None
+            truncation=True,
+            max_length=600  # 确保不超过模型限制
         )
         input_ids = inputs["input_ids"].to(torch.int64)
         
@@ -164,6 +166,7 @@ def main():
             
         except Exception as e:
             st.error(f"Magic failed: {str(e)}")
+            st.session_state.generated = False  # 确保失败时重置状态
     
     # 显示结果
     if st.session_state.generated:
@@ -176,6 +179,15 @@ def main():
         st.subheader("Listen 🎧")
         audio_data, sr = st.session_state.audio
         st.audio(audio_data, sample_rate=sr)
+        
+        st.markdown("---")
+        if st.button("Create New Story", use_container_width=True):
+            # 彻底清除所有相关状态
+            keys_to_clear = ['generated', 'story', 'audio']
+            for key in keys_to_clear:
+                if key in st.session_state:
+                    del st.session_state[key]
+            st.rerun()
 
 if __name__ == "__main__":
     main()