Update app.py

app.py

@@ -1,58 +1,100 @@
-from transformers import BlipProcessor, BlipForConditionalGeneration
+import streamlit as st
 from PIL import Image
-# Load model directly
-from transformers import AutoProcessor, AutoModelForImageTextToText
+from transformers import pipeline
 
-def generate_caption(image_path):
-    image = Image.open(image_path)
-    processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-    model = AutoModelForImageTextToText.from_pretrained("Salesforce/blip-image-captioning-base")
-    inputs = processor(image, return_tensors="pt")
-    output = model.generate(**inputs)
-    caption = processor.decode(output[0], skip_special_tokens=True)
+# ----------------------------
+# 生成图像描述函数
+# ----------------------------
+def generate_caption(image_file):
+    """
+    使用 Hugging Face pipeline 的 image-to-text 模型生成图片描述
+    参数:
+        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']  # 取第一个结果
     return caption
-from transformers import pipeline
 
+# ----------------------------
+# 基于图片描述生成完整故事的函数
+# ----------------------------
 def generate_story(caption):
-    # 使用文本生成 pipeline
-    generator = pipeline("text-generation", model="openai-community/gpt2")
-    # Use a pipeline as a high-level helper
-    prompt = f"由以下图片得到的描述: '{caption}'，请根据这个描述生成一个完整的童话故事，故事至少100个单词。"
-    result = generator(prompt, max_length=300, num_return_sequences=1)
+    """
+    基于图片描述生成完整故事，确保生成的故事至少包含100个单词。
+    参数:
+        caption: 图片描述文本
+    返回:
+        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:
-        # 可以进行递归调用或其他逻辑扩充文本
-        story += " " + generate_story(caption)
+        additional = story_generator(prompt, max_length=350, num_return_sequences=1)[0]['generated_text']
+        story += " " + additional
     return story
-from gtts import gTTS
 
+# ----------------------------
+# 文字转语音 (TTS) 函数
+# ----------------------------
 def text_to_speech(text, output_file="output.mp3"):
-    tts = gTTS(text=text, lang='en')  # 注意可根据需要选择语言或使用中文
+    """
+    将文本转换为语音并保存为 mp3 文件
+    参数:
+        text: 要转换的文本
+        output_file: 保存的音频文件名
+    返回:
+        output_file: 转换后的音频文件路径
+    """
+    from gtts import gTTS
+    # 这里语言参数设为英语 "en"，
+    # 如需中文可修改 lang="zh-cn"，但对应文本生成模型也需生成中文
+    tts = gTTS(text=text, lang="en")
     tts.save(output_file)
     return output_file
-import streamlit as st
-from PIL import Image
 
+# ----------------------------
+# 主函数：构建 Streamlit 界面
+# ----------------------------
 def main():
     st.title("儿童故事生成应用")
-    st.write("上传一张图片，我们将根据图片生成有趣的故事，并转换成语音播放给你听！!!")
+    st.write("上传一张图片，我们将根据图片生成有趣的故事，并转换成语音播放！")
     
     uploaded_file = st.file_uploader("选择一张图片", type=["png", "jpg", "jpeg"])
+    
     if uploaded_file is not None:
+        # 显示上传的图片
         image = Image.open(uploaded_file)
-        st.image(image, caption="上传的图片", use_container_width=True)
+        st.image(image, caption="上传的图片", use_column_width=True)
         
-        # 调用图像描述函数
-        caption = generate_caption(uploaded_file)
-        st.write("图片描述:", caption)
+        # 生成图片描述
+        with st.spinner("正在生成图片描述..."):
+            caption = generate_caption(uploaded_file)
+        st.write("图片描述：", caption)
         
-        # 生成故事
-        story = generate_story(caption)
-        st.write("生成的故事:", story)
+        # 根据图片描述生成完整故事
+        with st.spinner("正在生成故事..."):
+            story = generate_story(caption)
+        st.write("生成的故事：")
+        st.write(story)
         
-        # 文字转语音
-        audio_file = text_to_speech(story)
+        # 文本转语音
+        with st.spinner("正在转换成语音..."):
+            audio_file = text_to_speech(story)
         st.audio(audio_file, format="audio/mp3")
 
 if __name__ == "__main__":