Update app.py

app.py

@@ -1,52 +1,102 @@
 import streamlit as st
 from PIL import Image
 from transformers import BlipProcessor, BlipForConditionalGeneration
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from gtts import gTTS
 import io
 
+# ======================
+# Stage1: Image Captioning
+# ======================
 @st.cache_resource
-def load_models():
-    return (
-        BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base"),
-        BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base"),
-        GPT2Tokenizer.from_pretrained("gpt2"),
-        GPT2LMHeadModel.from_pretrained("gpt2")
+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 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)
+
+# ======================
+# Stage2: 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 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
+    
+    Story: Once upon a time, there was a little rabbit named Fluffy who loved"""
+    
+    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)])
 
+# ======================
+# Stage3: 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
+
+# ======================
+# Main Application
+# ======================
 def main():
-    st.title("Stable Story Maker")
-    
-    img_processor, img_model, text_tokenizer, text_model = load_models()
+    st.title("📚 Smart Story Generator")
     
-    uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png"])
+    uploaded_file = st.file_uploader("Upload children's photo", type=["jpg", "png"])
     
     if uploaded_file:
+        # Stage1
         st.image(uploaded_file, use_container_width=True)
+        caption = stage1_generate_caption(uploaded_file)
+        st.write(f"✨ Detected Theme: **{caption}**")
         
-        with st.status("Processing"):
-            # Stage 1
-            img = Image.open(uploaded_file).convert("RGB")
-            inputs = img_processor(images=img, return_tensors="pt")
-            caption = img_processor.decode(img_model.generate(**inputs)[0], skip_special_tokens=True)
-            
-            # Stage 2
-            prompt = f"Children's story about {caption}:"
-            inputs = text_tokenizer(prompt, return_tensors="pt")
-            story = text_tokenizer.decode(
-                text_model.generate(inputs.input_ids, max_length=200)[0],
-                skip_special_tokens=True
-            ).replace(prompt, "")
-            
-            # Stage 3
-            tts = gTTS(text=story[:250], lang='en')
-            audio = io.BytesIO()
-            tts.write_to_fp(audio)
-            audio.seek(0)
+        # Stage2
+        story = stage2_generate_story(caption)
+        st.subheader("Magic Story")
+        st.write(story)
         
-        st.write(f"**Caption:** {caption}")
-        st.write(f"**Story:** {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")
 
 if __name__ == "__main__":
     main()