Update app.py

app.py

@@ -34,36 +34,45 @@ def stage1_process(uploaded_file):
 # ======================
 @st.cache_resource
 def load_story_model():
-    """Load story generation model"""
+    """Load reliable story model"""
     return (
-        AutoTokenizer.from_pretrained("prpappas/fairytale-gpt2"),
-        AutoModelForCausalLM.from_pretrained("prpappas/fairytale-gpt2")
+        AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium"),
+        AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
     )
 
 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)
+    prompt = f"""Write a children's story about {keyword} with animals in 100 words.
+Story: Once upon a time, there was a little rabbit named Fluffy who found"""
+    
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=100, truncation=True)
     outputs = model.generate(
         inputs.input_ids,
-        max_length=200,
+        max_length=300,
         temperature=0.85,
         top_k=50,
         repetition_penalty=1.2
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).replace(prompt, "")
+    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return full_text.replace(prompt, "").strip()
 
 # ======================
 # Stage 3: Text-to-Speech
 # ======================
 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
+    try:
+        clean_text = text.strip().replace('\n', ' ')[:200]
+        if len(clean_text) < 10:
+            return None
+        tts = gTTS(text=clean_text, lang='en')
+        audio = io.BytesIO()
+        tts.write_to_fp(audio)
+        audio.seek(0)
+        return audio
+    except:
+        return None
 
 # ======================
 # Main Application
@@ -73,44 +82,46 @@ def main():
     
     # 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
+        st.session_state.update({
+            'stage1_done': False,
+            'stage2_done': False,
+            'caption': "",
+            'story': ""
+        })
     
-    # File upload section
+    # File upload
     uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png"])
     
     if uploaded_file:
-        # Always show image and Stage 1 result
+        # Permanent display
         st.image(uploaded_file, width=300)
         
-        # Stage 1 Processing
+        # Stage 1
         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.caption = stage1_process(uploaded_file)
                 st.session_state.stage1_done = True
         st.success(f"Detected Theme: {st.session_state.caption}")
         
-        # Stage 2 Processing
+        # Stage 2
         if not st.session_state.stage2_done:
-            with st.spinner("Creating story..."):
-                story = stage2_process(st.session_state.caption)
-                st.session_state.story = story
+            with st.spinner("Writing story..."):
+                st.session_state.story = stage2_process(st.session_state.caption)
                 st.session_state.stage2_done = True
-                
-        if st.session_state.stage2_done:
+        
+        # Display results
+        if st.session_state.story:
             st.subheader("Generated Story")
             st.write(st.session_state.story)
             
-            # Stage 3 Processing
+            # Stage 3
             with st.spinner("Generating audio..."):
                 audio = stage3_process(st.session_state.story)
-                st.audio(audio, format="audio/mp3")
-                st.download_button("Download Audio", 
-                                 data=audio.getvalue(),
-                                 file_name="story.mp3",
-                                 mime="audio/mp3")
+                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 short text")
 
 if __name__ == "__main__":
     main()