Update app.py

app.py

@@ -3,12 +3,12 @@ from PIL import Image
 from transformers import (
     BlipProcessor, 
     BlipForConditionalGeneration,
-    AutoTokenizer,
-    AutoModelForCausalLM
+    pipeline
 )
 from gtts import gTTS
 import io
 import logging
+import torch
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -19,7 +19,7 @@ logger = logging.getLogger(__name__)
 # ======================
 @st.cache_resource
 def load_image_model():
-    """Load official image captioning model"""
+    """Load image captioning model"""
     try:
         processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
         model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
@@ -34,7 +34,7 @@ def stage1_generate_caption(uploaded_file):
     processor, model = load_image_model()
     try:
         img = Image.open(uploaded_file).convert("RGB")
-        img.thumbnail((512, 512))  # Optimize image size
+        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)
@@ -47,57 +47,56 @@ def stage1_generate_caption(uploaded_file):
 # ======================
 @st.cache_resource
 def load_story_model():
-    """Load reliable story generation model"""
+    """Load Mistral story model"""
     try:
-        tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-        model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+        pipe = pipeline(
+            "text-generation",
+            model="ajibawa-2023/Young-Children-Storyteller-Mistral-7B",
+            device_map="auto",
+            torch_dtype=torch.float16
+        )
         logger.info("Stage 2 model loaded")
-        return tokenizer, model
+        return pipe
     except Exception as e:
         st.error("❌ Failed to load story model")
         raise
 
 def stage2_generate_story(keyword):
-    """Generate structured story"""
-    tokenizer, model = load_story_model()
-    
-    # Optimized prompt template
-    prompt = f"""Write a children's story with:
-    - Theme: {keyword}
-    - Characters: Animals
-    - Length: 100 words
-    
-    Story: Once upon a time, a little bear named Honey found"""
-    
+    """Generate story with chat format"""
+    pipe = load_story_model()
     try:
-        inputs = tokenizer(prompt, return_tensors="pt", max_length=100, truncation=True)
-        outputs = model.generate(
-            inputs.input_ids,
-            max_length=300,
-            temperature=0.85,
-            top_k=50,
-            repetition_penalty=1.2,
-            pad_token_id=tokenizer.eos_token_id
+        messages = [{
+            "role": "user", 
+            "content": f"Write a children's story about {keyword} with animals under 100 words"
+        }]
+        outputs = pipe(
+            messages,
+            max_new_tokens=200,
+            temperature=0.7,
+            do_sample=True
         )
-        full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return full_text.replace(prompt, "").strip()
+        return outputs[0]['generated_text'][-1]['content']
     except Exception as e:
         st.error(f"Story generation failed: {str(e)}")
-        return "The animals had a wonderful day playing together!"
+        return "The animals had a great time playing together!"
 
 # ======================
 # Stage 3: Text-to-Speech
 # ======================
 def stage3_generate_audio(text):
-    """Convert text to audio"""
+    """Generate audio with validation"""
     try:
-        tts = gTTS(text=text[:300], lang='en')
+        clean_text = text.strip().replace('\n', ' ')[:300]
+        if len(clean_text) < 10:
+            raise ValueError("Text too short")
+            
+        tts = gTTS(text=clean_text, 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"Audio generation failed: {str(e)}")
+        st.error(f"Audio Error: {str(e)}")
         return None
 
 # ======================
@@ -121,13 +120,14 @@ def main():
             st.subheader("Generated Story")
             st.write(story)
         
-        # Stage 3
-        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("Download Audio", audio.getvalue(), "story.mp3")
+        # Stage 3 (强制显示)
+        with st.spinner("Creating audio..."):
+            audio = stage3_generate_audio(story)
+            if audio:
+                st.audio(audio, format="audio/mp3")
+                st.download_button("Download Audio", audio.getvalue(), "story.mp3")
+            else:
+                st.warning("Audio generation skipped due to invalid input")
 
 if __name__ == "__main__":
     main()