Update app.py

app.py

@@ -19,14 +19,14 @@ logger = logging.getLogger(__name__)
 # ======================
 @st.cache_resource
 def load_image_model():
-    """Load official Hugging Face image captioning model"""
+    """Load official 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("❌ 图像模型加载失败，请检查网络连接")
+        st.error("❌ Image model failed to load")
         raise
 
 def stage1_generate_caption(uploaded_file):
@@ -34,12 +34,12 @@ def stage1_generate_caption(uploaded_file):
     processor, model = load_image_model()
     try:
         img = Image.open(uploaded_file).convert("RGB")
-        img.thumbnail((512, 512))  # Resize for speed
+        img.thumbnail((512, 512))
         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)}")
+        st.error(f"Image processing failed: {str(e)}")
         return "children playing"
 
 # ======================
@@ -47,27 +47,27 @@ def stage1_generate_caption(uploaded_file):
 # ======================
 @st.cache_resource
 def load_story_model():
-    """Load Microsoft DialoGPT model"""
+    """Load story generation model"""
     try:
-        tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-        model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+        tokenizer = AutoTokenizer.from_pretrained("pranavpsv/gpt-genre-story-generator")
+        model = AutoModelForCausalLM.from_pretrained("pranavpsv/gpt-genre-story-generator", use_auth_token=True)
         logger.info("Stage 2 model loaded")
         return tokenizer, model
     except Exception as e:
-        st.error("❌ 故事模型加载失败，请检查模型名称")
+        st.error("❌ Story model failed to load")
         raise
 
 def stage2_generate_story(keyword):
-    """Generate children's story"""
+    """Generate structured story"""
     tokenizer, model = load_story_model()
     
     # Optimized prompt template
-    prompt = f"""写一个儿童故事，包含以下要素:
-    - 主题: {keyword}
-    - 角色: 小动物
-    - 字数: 100字左右
+    prompt = f"""Generate a children's story with:
+    - Theme: {keyword}
+    - Characters: Animals
+    - Word count: 100 words
     
-    故事开头: 有一天，小熊嘟嘟在公园里发现"""
+    Story: Once upon a time, a little bear named Honey discovered"""
     
     try:
         inputs = tokenizer(prompt, return_tensors="pt", max_length=100, truncation=True)
@@ -76,14 +76,13 @@ def stage2_generate_story(keyword):
             max_length=300,
             temperature=0.9,
             top_k=50,
-            repetition_penalty=1.2,
-            pad_token_id=tokenizer.eos_token_id
+            repetition_penalty=1.2
         )
         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 "小熊和朋友们玩得很开心！"
+        st.error(f"Story generation failed: {str(e)}")
+        return "The animals had a wonderful adventure!"
 
 # ======================
 # Stage 3: Text-to-Speech
@@ -91,43 +90,43 @@ def stage2_generate_story(keyword):
 def stage3_generate_audio(text):
     """Convert text to audio"""
     try:
-        tts = gTTS(text=text[:300], lang='zh-CN')  # Chinese support
+        tts = gTTS(text=text[:300], lang='en')
         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)}")
+        st.error(f"Audio generation failed: {str(e)}")
         return None
 
 # ======================
 # Main Application
 # ======================
 def main():
-    st.title("📚 智能故事生成器")
+    st.title("📚 Smart Story Generator")
     
-    uploaded_file = st.file_uploader("上传儿童照片", type=["jpg", "png", "jpeg"])
+    uploaded_file = st.file_uploader("Upload Photo (JPG/PNG)", type=["jpg", "png", "jpeg"])
     
     if uploaded_file:
         # Stage 1
         st.image(uploaded_file, use_container_width=True)
-        with st.spinner("正在分析图片..."):
+        with st.spinner("Analyzing image..."):
             caption = stage1_generate_caption(uploaded_file)
-            st.write(f"✨ 识别主题: **{caption}**")
+            st.write(f"✨ Detected Theme: **{caption}**")
         
         # Stage 2
-        with st.spinner("正在生成故事..."):
+        with st.spinner("Generating story..."):
             story = stage2_generate_story(caption)
-            st.subheader("生成故事")
+            st.subheader("Generated Story")
             st.write(story)
         
         # Stage 3
-        if len(story) > 10:  # Minimum length check
-            with st.spinner("正在生成语音..."):
+        if len(story) > 20:
+            with st.spinner("Creating audio..."):
                 audio = stage3_generate_audio(story)
                 if audio:
                     st.audio(audio, format="audio/mp3")
-                    st.download_button("下载语音", audio.getvalue(), "story.mp3")
+                    st.download_button("Download Audio", audio.getvalue(), "story.mp3")
 
 if __name__ == "__main__":
     main()