Update app.py

app.py

@@ -7,21 +7,19 @@ import os
 
 def img2text(img_path):
     captioner = pipeline(
-        "image-to-text", 
-        model="nlpconnect/vit-gpt2-image-captioning",
-        device_map="auto"
+        'image-to-text',
+        model='nlpconnect/vit-gpt2-image-captioning',
+        device_map='auto'
     )
     result = captioner(img_path)
     return result[0]["generated_text"]
 
 def text2story(scenario):
     generator = pipeline(
-        "text-generation",
-        model="gpt2",
-        trust_remote_code=True,
-        device_map="auto",    # 必须添加
-        torch_dtype=torch.float16, 
-        max_length=100,
+        'text-generation',
+        model='gpt2',
+        device_map='auto',
+        max_length=200,
         num_return_sequences=1
     )
     prompt = f"Create a children's story based on: {scenario}"
@@ -29,17 +27,22 @@ def text2story(scenario):
     return story
 
 def text2audio(story_text):
-    tts = gTTS(text=story_text, lang="en")
-    audio_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
+    tts = gTTS(text=story_text, lang='en')
+    audio_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp3')
     tts.save(audio_file.name)
     return audio_file.name
 
-    
 def main():
-    st.set_page_config(page_title="Image to Story")
+    st.set_page_config(
+        page_title='Image to Story',
+        page_icon='📖'
+    )
     st.header("Upload Your Image")
     
-    uploaded_file = st.file_uploader("Choose image", type=["jpg", "png", "jpeg"])
+    uploaded_file = st.file_uploader(
+        "Choose Image",
+        type=['jpg', 'png', 'jpeg']
+    )
     
     if uploaded_file:
         temp_img = os.path.join(tempfile.gettempdir(), uploaded_file.name)
@@ -48,38 +51,28 @@ def main():
         
         st.image(uploaded_file)
         
-        # 图片转文字处理状态
-        with st.status("🖼️ Processing img2text...", expanded=True) as status:
+        # 图片分析
+        with st.status('🖼️ Processing image...'):
             scenario = img2text(temp_img)
-            status.update(label="✅ Image analysis completed!", state="complete")
+            st.write("Image Caption:", scenario)
         
-        st.write("Image Caption:", scenario)
-        
-        # 故事生成处理状态
-        with st.status("📝 Generating a story...", expanded=True) as status:
+        # 生成故事
+        with st.status('📖 Generating story...'):
             story = text2story(scenario)
-            status.update(label="📖 Story generation completed!", state="complete")
-        
-        st.subheader("Story")
-        st.write(story)
+            st.subheader('Story')
+            st.write(story)
         
-        # 音频生成处理状态
-        with st.status("🔊 Converting to audio...", expanded=True) as status:
+        # 生成音频
+        with st.status('🔊 Converting audio...'):
             audio_path = text2audio(story)
-            status.update(label="🎧 Audio conversion completed!", state="complete")
         
-        st.audio(audio_path)
+        # 添加播放按钮（新增部分）
+        if st.button('▶️ Play Audio Story'):
+            st.audio(audio_path, format="audio/mp3")
         
+        # 清理文件
         os.unlink(temp_img)
         os.unlink(audio_path)
 
-
-        if st.button("🎵 Play Audio Story"):
-            st.audio(
-            audio_data["audio"],
-            format="audio/mp3",
-            start_time=0,
-            sample_rate=audio_data["sampling_rate"]
-            )
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()