Update app.py

app.py

@@ -1,9 +1,13 @@
 import streamlit as st
 from PIL import Image
 from transformers import pipeline
+from gtts import gTTS
+from diffusers import StableDiffusionPipeline
+import torch
+import os
 
 # ----------------------------
-# 生成图像描述函数
+# 1. 图像描述生成函数
 # ----------------------------
 def generate_caption(image_file):
     """
@@ -13,67 +17,85 @@ def generate_caption(image_file):
     返回:
         caption: 生成的图片描述文本
     """
-    # 打开图片（如果上传的是文件流，可以直接传给 pipeline）
     image = Image.open(image_file)
-    # 利用 image-to-text pipeline 加载 Salesforce/blip-image-captioning-base 模型
     caption_generator = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    # 直接将图片传入 pipeline，返回结果是一个列表，每个元素是一个字典
     caption_results = caption_generator(image)
-    caption = caption_results[0]['generated_text']  # 取第一个结果
+    caption = caption_results[0]['generated_text']  # 取返回结果的第一个描述
     return caption
 
 # ----------------------------
-# 基于图片描述生成完整故事的函数
+# 2. 故事生成函数
 # ----------------------------
-def generate_story(caption):
+def generate_story(prompt):
     """
-    基于图片描述生成完整故事，确保生成的故事至少包含100个单词。
+    基于提示语生成故事段落，要求至少100个单词，如果生成的文本字数不够，则再次补充
     参数:
-        caption: 图片描述文本
+        prompt: 文本生成的提示语
     返回:
-        story: 生成的故事文本
+        story: 生成的故事文本片段
     """
-    # 使用 text-generation pipeline 加载 GPT-2 模型
     story_generator = pipeline("text-generation", model="gpt2")
-    # 构建生成故事的提示语
-    prompt = f"Based on the following image caption: '{caption}', generate a complete fairy tale story for children with at least 100 words. "
-    
-    # 生成故事文本
     result = story_generator(prompt, max_length=300, num_return_sequences=1)
     story = result[0]['generated_text']
     
-    # 简单检查生成的故事单词数是否达到100，否则再生成部分文本补充
     if len(story.split()) < 100:
         additional = story_generator(prompt, max_length=350, num_return_sequences=1)[0]['generated_text']
         story += " " + additional
     return story
 
 # ----------------------------
-# 文字转语音 (TTS) 函数
+# 3. 图像生成（配图）相关函数
+# ----------------------------
+@st.cache_resource
+def load_image_generator():
+    """
+    加载稳定扩散模型，使用 Diffusers 库生成插图
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-v1-5",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    )
+    pipe = pipe.to(device)
+    return pipe
+
+def generate_illustration(prompt):
+    """
+    基于输入的提示语生成一张配图
+    参数:
+        prompt: 用于生成图像的文本提示
+    返回:
+        generated_image: 生成的 PIL Image 图像
+    """
+    pipe = load_image_generator()
+    image_result = pipe(prompt)
+    generated_image = image_result.images[0]
+    return generated_image
+
+# ----------------------------
+# 4. 文字转语音 (TTS) 函数
 # ----------------------------
 def text_to_speech(text, output_file="output.mp3"):
     """
-    将文本转换为语音并保存为 mp3 文件
+    将输入文本转换为语音，并保存为 mp3 文件
     参数:
         text: 要转换的文本
         output_file: 保存的音频文件名
     返回:
-        output_file: 转换后的音频文件路径
+        output_file: 转换后生成的音频文件路径
     """
-    from gtts import gTTS
-    # 这里语言参数设为英语 "en"，
-    # 如需中文可修改 lang="zh-cn"，但对应文本生成模型也需生成中文
-    tts = gTTS(text=text, lang="en")
+    tts = gTTS(text=text, lang="en")  # 如需中文，lang 可设置为 "zh-cn"
     tts.save(output_file)
     return output_file
 
 # ----------------------------
-# 主函数：构建 Streamlit 界面
+# 5. 主函数：构建 Streamlit 交互式应用
 # ----------------------------
 def main():
-    st.title("儿童故事生成应用")
-    st.write("上传一张图片，我们将根据图片生成有趣的故事，并转换成语音播放！")
+    st.title("互动式故事生成与配图应用")
+    st.write("上传一张图片，我们会基于该图片生成描述，并自动生成一个儿童故事。你可以选择继续扩展改故事，也可以结束互动。每个生成的故事段落都会搭配 AI 配图。")
     
+    # 图片上传
     uploaded_file = st.file_uploader("选择一张图片", type=["png", "jpg", "jpeg"])
     
     if uploaded_file is not None:
@@ -84,18 +106,51 @@ def main():
         # 生成图片描述
         with st.spinner("正在生成图片描述..."):
             caption = generate_caption(uploaded_file)
-        st.write("图片描述：", caption)
+        st.write("图片描述:", caption)
         
-        # 根据图片描述生成完整故事
-        with st.spinner("正在生成故事..."):
-            story = generate_story(caption)
-        st.write("生成的故事：")
-        st.write(story)
+        # 使用 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]
+        
+        st.write("### 生成的故事：")
+        st.write(st.session_state.story)
+        
+        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")
         
-        # 文本转语音
-        with st.spinner("正在转换成语音..."):
-            audio_file = text_to_speech(story)
-        st.audio(audio_file, format="audio/mp3")
-
 if __name__ == "__main__":
     main()