Update app.py

app.py

@@ -1,101 +1,79 @@
+
 import streamlit as st
-from PIL import Image
 from transformers import pipeline
+from PIL import Image
+from gtts import gTTS
+from io import BytesIO
 
-# ----------------------------
-# 生成图像描述函数
-# ----------------------------
-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
+# Set up the image captioning pipeline
+@st.cache_resource
+def get_image_captioner():
+    return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# ----------------------------
-# 基于图片描述生成完整故事的函数
-# ----------------------------
-def generate_story(caption):
-    """
-    基于图片描述生成完整故事，确保生成的故事至少包含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:
-        additional = story_generator(prompt, max_length=350, num_return_sequences=1)[0]['generated_text']
-        story += " " + additional
-    return story
+# Set up the story generation pipeline
+@st.cache_resource
+def get_story_generator():
+    return pipeline("text-generation", model="Qwen/Qwen2.5-1.5B-Instruct", padding=True)
 
-# ----------------------------
-# 文字转语音 (TTS) 函数
-# ----------------------------
-def text_to_speech(text, output_file="output.mp3"):
-    """
-    将文本转换为语音并保存为 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
+# Convert text to speech using gTTS
+def text_to_speech(text):
+    tts = gTTS(text=text, lang='en')
+    audio_bytes = BytesIO()
+    tts.write_to_fp(audio_bytes)
+    audio_bytes.seek(0)
+    return audio_bytes
 
-# ----------------------------
-# 主函数：构建 Streamlit 界面
-# ----------------------------
-def main():
-    st.title("儿童故事生成应用")
-    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_column_width=True)
-        
-        # 生成图片描述
-        with st.spinner("正在生成图片描述..."):
-            caption = generate_caption(uploaded_file)
-        st.write("图片描述：", caption)
-        
-        # 根据图片描述生成完整故事
-        with st.spinner("正在生成故事..."):
-            story = generate_story(caption)
-        st.write("生成的故事：")
-        st.write(story)
+# Streamlit UI
+st.title("📖 Kids' Storytelling App")
+st.write("Upload an image and let the magic create a story!")
+
+uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
+
+if uploaded_image is not None:
+    try:
+        # Open and preprocess the image
+        image = Image.open(uploaded_image).convert("RGB")  # Ensure image is in RGB format
+        st.image(image, caption="Your Image", use_column_width=True)
         
-        # 文本转语音
-        with st.spinner("正在转换成语音..."):
-            audio_file = text_to_speech(story)
-        st.audio(audio_file, format="audio/mp3")
+        if st.button("Generate Story!"):
+            with st.spinner("Creating your story..."):
+                # Generate image caption
+                captioner = get_image_captioner()
+                caption_result = captioner(image)
+                caption = caption_result[0]['generated_text']
+                st.subheader("Image Caption")
+                st.write(caption)
+                
+                # Generate story from caption
+                story_gen = get_story_generator()
+                prompt = f"Create a fun, children's story based on this: {caption}. The story must be at least 100 words, imaginative, and suitable for kids aged 3-10. Story:"
+                story_result = story_gen(
+                    prompt,
+                    max_length=100,
+                    num_return_sequences=1,
+                    temperature=0.9,
+                    repetition_penalty=1.2
+                )
+   
 
-if __name__ == "__main__":
-    main()
+story = story_result[0]['generated_text']
+                
+                # Ensure story meets word count
+                if len(story.split()) < 100:
+                    story_result = story_gen(
+                        prompt,
+                        max_length=100,
+                        num_return_sequences=1,
+                        temperature=0.9
+                    )
+                    story = story_result[0]['generated_text']
+                
+                st.subheader("Your Story")
+                st.write(story)
+                
+                # Convert story to audio
+                audio_bytes = text_to_speech(story)
+                st.subheader("Listen to the Story")
+                st.audio(audio_bytes, format="audio/mp3")
+    except Exception as e:
+        st.error(f"An error occurred: {e}")