Update app.py

app.py

@@ -1,102 +1,133 @@
 import streamlit as st
 from PIL import Image
-from transformers import BlipProcessor, BlipForConditionalGeneration
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    BlipProcessor, 
+    BlipForConditionalGeneration,
+    AutoTokenizer,
+    AutoModelForCausalLM
+)
 from gtts import gTTS
 import io
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # ======================
-# Stage1: Image Captioning
+# Stage 1: Image Captioning
 # ======================
 @st.cache_resource
-def load_stage1_model():
-    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-    return processor, model
+def load_image_model():
+    """Load official Hugging Face image captioning model"""
+    try:
+        processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+        model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+        logger.info("Stage 1 model loaded")
+        return processor, model
+    except Exception as e:
+        st.error("❌ 图像模型加载失败，请检查网络连接")
+        raise
 
 def stage1_generate_caption(uploaded_file):
-    processor, model = load_stage1_model()
-    img = Image.open(uploaded_file).convert("RGB")
-    inputs = processor(images=img, return_tensors="pt", padding=True)
-    outputs = model.generate(**inputs)
-    return processor.decode(outputs[0], skip_special_tokens=True)
+    """Generate image caption"""
+    processor, model = load_image_model()
+    try:
+        img = Image.open(uploaded_file).convert("RGB")
+        img.thumbnail((512, 512))  # Resize for speed
+        inputs = processor(images=img, return_tensors="pt", padding=True)
+        outputs = model.generate(**inputs, max_length=30)
+        return processor.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        st.error(f"图像处理失败: {str(e)}")
+        return "children playing"
 
 # ======================
-# Stage2: Story Generation
+# Stage 2: Story Generation
 # ======================
 @st.cache_resource
-def load_stage2_model():
-    tokenizer = AutoTokenizer.from_pretrained("pranavpsv/gpt-genre-story-generator")
-    model = AutoModelForCausalLM.from_pretrained("pranavpsv/gpt-genre-story-generator")
-    return tokenizer, model
+def load_story_model():
+    """Load Microsoft DialoGPT model"""
+    try:
+        tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+        model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+        logger.info("Stage 2 model loaded")
+        return tokenizer, model
+    except Exception as e:
+        st.error("❌ 故事模型加载失败，请检查模型名称")
+        raise
 
 def stage2_generate_story(keyword):
-    tokenizer, model = load_stage2_model()
-    
-    # 专业prompt模板
-    prompt_template = f"""Generate a children's story in English with these elements:
-    - Main theme: {keyword}
-    - Characters: Friendly animals
-    - Plot: Daily adventure
-    - Moral lesson: Sharing is caring
-    - Word count: 50-100 words
+    """Generate children's story"""
+    tokenizer, model = load_story_model()
     
-    Story: Once upon a time, there was a little rabbit named Fluffy who loved"""
+    # Optimized prompt template
+    prompt = f"""写一个儿童故事，包含以下要素:
+    - 主题: {keyword}
+    - 角色: 小动物
+    - 字数: 100字左右
     
-    inputs = tokenizer(prompt_template, return_tensors="pt")
-    outputs = model.generate(
-        inputs.input_ids,
-        max_length=300,
-        temperature=0.85,
-        top_k=50,
-        repetition_penalty=1.2,
-        num_return_sequences=1
-    )
-    full_story = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    故事开头: 有一天，小熊嘟嘟在公园里发现"""
     
-    # 提取生成部分并标准化格式
-    generated_part = full_story.replace(prompt_template, "").strip()
-    return _format_story(generated_part)
-
-def _format_story(raw_text):
-    # 后处理：添加段落结构
-    sentences = raw_text.split(". ")
-    return "\n\n".join([". ".join(sentences[i:i+3]) + "." for i in range(0, len(sentences), 3)])
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt", max_length=100, truncation=True)
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=300,
+            temperature=0.9,
+            top_k=50,
+            repetition_penalty=1.2,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return full_text.replace(prompt, "").strip()
+    except Exception as e:
+        st.error(f"故事生成失败: {str(e)}")
+        return "小熊和朋友们玩得很开心！"
 
 # ======================
-# Stage3: Text-to-Speech
+# Stage 3: Text-to-Speech
 # ======================
-def stage3_generate_audio(story_text):
-    tts = gTTS(text=story_text, lang='en')
-    audio_buffer = io.BytesIO()
-    tts.write_to_fp(audio_buffer)
-    audio_buffer.seek(0)
-    return audio_buffer
+def stage3_generate_audio(text):
+    """Convert text to audio"""
+    try:
+        tts = gTTS(text=text[:300], lang='zh-CN')  # Chinese support
+        audio_buffer = io.BytesIO()
+        tts.write_to_fp(audio_buffer)
+        audio_buffer.seek(0)
+        return audio_buffer
+    except Exception as e:
+        st.error(f"语音生成失败: {str(e)}")
+        return None
 
 # ======================
 # Main Application
 # ======================
 def main():
-    st.title("📚 Smart Story Generator")
+    st.title("📚 智能故事生成器")
     
-    uploaded_file = st.file_uploader("Upload children's photo", type=["jpg", "png"])
+    uploaded_file = st.file_uploader("上传儿童照片", type=["jpg", "png", "jpeg"])
     
     if uploaded_file:
-        # Stage1
+        # Stage 1
         st.image(uploaded_file, use_container_width=True)
-        caption = stage1_generate_caption(uploaded_file)
-        st.write(f"✨ Detected Theme: **{caption}**")
+        with st.spinner("正在分析图片..."):
+            caption = stage1_generate_caption(uploaded_file)
+            st.write(f"✨ 识别主题: **{caption}**")
         
-        # Stage2
-        story = stage2_generate_story(caption)
-        st.subheader("Magic Story")
-        st.write(story)
+        # Stage 2
+        with st.spinner("正在生成故事..."):
+            story = stage2_generate_story(caption)
+            st.subheader("生成故事")
+            st.write(story)
         
-        # Stage3
-        audio = stage3_generate_audio(story[:500])  # Limit for TTS
-        st.audio(audio, format="audio/mp3")
-        st.download_button("Download Story", story, "story.txt")
-        st.download_button("Download Audio", audio.getvalue(), "story.mp3")
+        # Stage 3
+        if len(story) > 10:  # Minimum length check
+            with st.spinner("正在生成语音..."):
+                audio = stage3_generate_audio(story)
+                if audio:
+                    st.audio(audio, format="audio/mp3")
+                    st.download_button("下载语音", audio.getvalue(), "story.mp3")
 
 if __name__ == "__main__":
     main()