Update app.py

app.py

@@ -3,131 +3,114 @@ from PIL import Image
 from transformers import (
     BlipProcessor, 
     BlipForConditionalGeneration,
-    pipeline
+    AutoTokenizer,
+    AutoModelForCausalLM
 )
 from gtts import gTTS
 import io
-import logging
 import torch
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
 # ======================
 # Stage 1: Image Captioning
 # ======================
 @st.cache_resource
 def load_image_model():
     """Load 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("❌ Failed to load image model")
-        raise
+    return (
+        BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base"),
+        BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+    )
 
-def stage1_generate_caption(uploaded_file):
+def stage1_process(uploaded_file):
     """Generate image caption"""
     processor, model = load_image_model()
-    try:
-        img = Image.open(uploaded_file).convert("RGB")
-        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"Image processing failed: {str(e)}")
-        return "children playing"
+    img = Image.open(uploaded_file).convert("RGB")
+    inputs = processor(images=img, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    return processor.decode(outputs[0], skip_special_tokens=True)
 
 # ======================
 # Stage 2: Story Generation
 # ======================
 @st.cache_resource
 def load_story_model():
-    """Load Mistral story model"""
-    try:
-        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 pipe
-    except Exception as e:
-        st.error("❌ Failed to load story model")
-        raise
+    """Load story generation model"""
+    return (
+        AutoTokenizer.from_pretrained("prpappas/fairytale-gpt2"),
+        AutoModelForCausalLM.from_pretrained("prpappas/fairytale-gpt2")
+    )
 
-def stage2_generate_story(keyword):
-    """Generate story with chat format"""
-    pipe = load_story_model()
-    try:
-        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
-        )
-        return outputs[0]['generated_text'][-1]['content']
-    except Exception as e:
-        st.error(f"Story generation failed: {str(e)}")
-        return "The animals had a great time playing together!"
+def stage2_process(keyword):
+    """Generate children's story"""
+    tokenizer, model = load_story_model()
+    prompt = f"Write a children's story about {keyword} in 100 words:\n"
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=50, truncation=True)
+    outputs = model.generate(
+        inputs.input_ids,
+        max_length=200,
+        temperature=0.85,
+        top_k=50,
+        repetition_penalty=1.2
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).replace(prompt, "")
 
 # ======================
 # Stage 3: Text-to-Speech
 # ======================
-def stage3_generate_audio(text):
-    """Generate audio with validation"""
-    try:
-        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 Error: {str(e)}")
-        return None
+def stage3_process(text):
+    """Convert text to audio"""
+    tts = gTTS(text=text[:200], lang='en')
+    audio = io.BytesIO()
+    tts.write_to_fp(audio)
+    audio.seek(0)
+    return audio
 
 # ======================
 # Main Application
 # ======================
 def main():
-    st.title("📚 Smart Story Generator")
+    st.title("📖 Children's Story Generator")
+    
+    # Initialize session state
+    if 'stage1_done' not in st.session_state:
+        st.session_state.stage1_done = False
+    if 'stage2_done' not in st.session_state:
+        st.session_state.stage2_done = False
     
-    uploaded_file = st.file_uploader("Upload Photo (JPG/PNG)", type=["jpg", "png", "jpeg"])
+    # File upload section
+    uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png"])
     
     if uploaded_file:
-        # Stage 1
-        st.image(uploaded_file, use_container_width=True)
-        with st.spinner("Analyzing image..."):
-            caption = stage1_generate_caption(uploaded_file)
-            st.write(f"✨ Detected Theme: **{caption}**")
+        # Always show image and Stage 1 result
+        st.image(uploaded_file, width=300)
         
-        # Stage 2
-        with st.spinner("Generating story..."):
-            story = stage2_generate_story(caption)
-            st.subheader("Generated Story")
-            st.write(story)
+        # Stage 1 Processing
+        if not st.session_state.stage1_done:
+            with st.spinner("Analyzing image..."):
+                caption = stage1_process(uploaded_file)
+                st.session_state.caption = caption
+                st.session_state.stage1_done = True
+        st.success(f"Detected Theme: {st.session_state.caption}")
         
-        # Stage 3 (强制显示)
-        with st.spinner("Creating audio..."):
-            audio = stage3_generate_audio(story)
-            if audio:
+        # Stage 2 Processing
+        if not st.session_state.stage2_done:
+            with st.spinner("Creating story..."):
+                story = stage2_process(st.session_state.caption)
+                st.session_state.story = story
+                st.session_state.stage2_done = True
+                
+        if st.session_state.stage2_done:
+            st.subheader("Generated Story")
+            st.write(st.session_state.story)
+            
+            # Stage 3 Processing
+            with st.spinner("Generating audio..."):
+                audio = stage3_process(st.session_state.story)
                 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")
+                st.download_button("Download Audio", 
+                                 data=audio.getvalue(),
+                                 file_name="story.mp3",
+                                 mime="audio/mp3")
 
 if __name__ == "__main__":
     main()