Update app.py

app.py

@@ -1,28 +1,52 @@
 import streamlit as st
-from transformers import pipeline
-from gtts import gTTS
-import io
+import importlib.util
+try:
+    # 检查 accelerate 库是否安装
+    spec = importlib.util.find_spec("accelerate")
+    if spec is None:
+        st.error("缺少 'accelerate' 库，请安装该库以加载 FP8 量化模型。可以使用 'pip install accelerate' 进行安装。")
+        st.stop()
+    from transformers import pipeline
+    from gtts import gTTS
+    import io
+    import tempfile
+    import os
+except ImportError as e:
+    st.error(f"导入库时出错: {e}")
+    st.stop()
 
 # function part
 # img2text
 def img2text(url):
-    image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    text = image_to_text_model(url)[0]["generated_text"]
-    return text
+    try:
+        image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+        text = image_to_text_model(url)[0]["generated_text"]
+        return text
+    except Exception as e:
+        st.error(f"图像描述生成出错: {e}")
+        return None
 
 # text2story
 def text2story(text):
-    story_generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)
-    story = story_generator(text, max_length=200, num_return_sequences=1)[0]['generated_text']
-    return story
+    try:
+        story_generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)
+        story = story_generator(text, max_length=200, num_return_sequences=1)[0]['generated_text']
+        return story
+    except Exception as e:
+        st.error(f"故事生成出错: {e}")
+        return None
 
 # text2audio
 def text2audio(story_text):
-    tts = gTTS(text=story_text, lang='en')
-    audio_file = io.BytesIO()
-    tts.write_to_fp(audio_file)
-    audio_file.seek(0)
-    return audio_file
+    try:
+        tts = gTTS(text=story_text, lang='en')
+        audio_file = io.BytesIO()
+        tts.write_to_fp(audio_file)
+        audio_file.seek(0)
+        return audio_file
+    except Exception as e:
+        st.error(f"文本转语音出错: {e}")
+        return None
 
 st.set_page_config(page_title="Your Image to Audio Story",
                    page_icon="🦜")
@@ -31,29 +55,35 @@ uploaded_file = st.file_uploader("Select an Image...")
 
 if uploaded_file is not None:
     print(uploaded_file)
-    bytes_data = uploaded_file.getvalue()
-    with open(uploaded_file.name, "wb") as file:
-        file.write(bytes_data)
+    # 使用临时文件处理上传的图像
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
+        temp_file.write(uploaded_file.getvalue())
+        temp_file_path = temp_file.name
+
     st.image(uploaded_file, caption="Uploaded Image",
-             use_column_width=True)
+             use_container_width=True)  # 修改为 use_container_width
 
     #Stage 1: Image to Text
     st.text('Processing img2text...')
-    scenario = img2text(uploaded_file.name)
-    st.write(scenario)
-
-    #Stage 2: Text to Story
-    st.text('Generating a story...')
-    story = text2story(scenario)
-    st.write(story)
-
-    #Stage 3: Story to Audio data
-    st.text('Generating audio data...')
-    audio_data = text2audio(story)
-
-    # Play button
-    if st.button("Play Audio"):
-        st.audio(audio_data,
-                 format="audio/mpeg",
-                 start_time=0)
- 
+    scenario = img2text(temp_file_path)
+    if scenario:
+        st.write(scenario)
+
+        #Stage 2: Text to Story
+        st.text('Generating a story...')
+        story = text2story(scenario)
+        if story:
+            st.write(story)
+
+            #Stage 3: Story to Audio data
+            st.text('Generating audio data...')
+            audio_data = text2audio(story)
+            if audio_data:
+                # Play button
+                if st.button("Play Audio"):
+                    st.audio(audio_data,
+                             format="audio/mpeg",
+                             start_time=0)
+
+    # 删除临时文件
+    os.remove(temp_file_path)