Update app.py

app.py

@@ -2,7 +2,6 @@ import streamlit as st
 from PIL import Image
 from transformers import pipeline
 from gtts import gTTS
-from diffusers import DiffusionPipeline
 import torch
 import os
 
@@ -50,9 +49,19 @@ def generate_story(prompt):
 def load_image_generator():
     """
     加载稳定扩散模型，使用 Diffusers 库生成插图
+    使用 StableDiffusionPipeline 替代 DiffusionPipeline
     """
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
+    
+    # 导入 StableDiffusionPipeline
+    from diffusers import StableDiffusionPipeline
+
+    # 对于 GPU，采用 fp16 精度以加速推理；否则使用默认精度
+    torch_dtype = torch.float16 if device == "cuda" else torch.float32
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-v1-5",
+        torch_dtype=torch_dtype
+    )
     pipe = pipe.to(device)
     return pipe
 
@@ -89,10 +98,9 @@ def text_to_speech(text, output_file="output.mp3"):
 # 5. 主函数：构建 Streamlit 交互式应用
 # ----------------------------
 def main():
-    st.title("互动式故事生成与配图应用")
-    st.write("上传一张图片，我们会基于该图片生成描述，并自动生成一个儿童故事。你可以选择继续扩展改故事，也可以结束互动。每个生成的故事段落都会搭配 AI 配图。")
+    st.title("儿童故事生成应用")
+    st.write("上传一张图片，我们将根据图片生成有趣的故事，并转换成语音播放！")
     
-    # 图片上传
     uploaded_file = st.file_uploader("选择一张图片", type=["png", "jpg", "jpeg"])
     
     if uploaded_file is not None:
@@ -103,51 +111,26 @@ def main():
         # 生成图片描述
         with st.spinner("正在生成图片描述..."):
             caption = generate_caption(uploaded_file)
-        st.write("图片描述:", caption)
+        st.write("图片描述：", caption)
         
-        # 使用 session_state 保存生成的故事和插图历史
-        if "story" not in st.session_state:
-            # 生成初始故事段落（至少100个单词）
-            with st.spinner("正在生成初始故事..."):
-                initial_prompt = f"Based on the image caption: '{caption}', generate a complete fairy tale story for children with at least 100 words."
-                story_segment = generate_story(initial_prompt)
-                st.session_state.story = story_segment
-            # 生成初始配图，使用初始故事的前200个字符作为提示
-            with st.spinner("正在生成初始配图..."):
-                illustration = generate_illustration(st.session_state.story[:200])
-                st.session_state.illustrations = [illustration]
+        # 根据图片描述生成完整故事
+        with st.spinner("正在生成故事..."):
+            story = generate_story(caption)
+        st.write("生成的故事：")
+        st.write(story)
         
-        st.write("### 生成的故事：")
-        st.write(st.session_state.story)
+        # 生成配图
+        # 这里使用故事内容的前200个字符作为提示生成配图，实际中可以根据需要调整策略
+        with st.spinner("正在生成插图..."):
+            illustration = generate_illustration(story[:200])
         
         st.write("### 故事配图：")
-        for idx, illus in enumerate(st.session_state.illustrations):
-            st.image(illus, caption=f"配图段落 {idx+1}", use_column_width=True)
-        
-        st.write("---")
-        st.write("是否继续生成故事？如果不再扩展，请点击“结束互动”。")
-        
-        # 接收用户输入的额外情节提示（可选）
-        user_input = st.text_input("请输入你希望添加的故事情节（可选）：", value="")
-        
-        col1, col2 = st.columns(2)
-        if col1.button("继续生成故事"):
-            # 使用现有故事作为上下文，并附加用户输入的提示语生成新段落
-            additional_prompt = st.session_state.story + " " + (user_input if user_input.strip() != "" else "")
-            with st.spinner("正在生成新的故事段落..."):
-                new_segment = generate_story(additional_prompt)
-            st.session_state.story += " " + new_segment
-            # 为新段落生成配图，取新段落前200个字符作为提示
-            with st.spinner("正在生成新的配图..."):
-                new_illustration = generate_illustration(new_segment[:200])
-            st.session_state.illustrations.append(new_illustration)
-            st.experimental_rerun()
-        
-        if col2.button("结束互动"):
-            with st.spinner("正在生成故事音频..."):
-                audio_file = text_to_speech(st.session_state.story)
-            st.write("故事生成完毕！请点击下方按钮播放故事音频。")
-            st.audio(audio_file, format="audio/mp3")
+        st.image(illustration, caption="配图", use_column_width=True)
         
+        # 文本转语音
+        with st.spinner("正在转换成语音..."):
+            audio_file = text_to_speech(story)
+        st.audio(audio_file, format="audio/mp3")
+
 if __name__ == "__main__":
     main()