Update app.py

app.py

@@ -30,32 +30,40 @@ def stage1_process(uploaded_file):
     return processor.decode(outputs[0], skip_special_tokens=True)
 
 # ======================
-# Stage 2: Story Generation
+# Stage 2: Story Generation (Optimized)
 # ======================
 @st.cache_resource
 def load_story_model():
-    """Load reliable story model"""
+    """Load optimized story model"""
     return (
-        AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium"),
-        AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+        AutoTokenizer.from_pretrained("gpt2-medium"),
+        AutoModelForCausalLM.from_pretrained("gpt2-medium")
     )
 
 def stage2_process(keyword):
-    """Generate children's story"""
+    """Generate structured story"""
     tokenizer, model = load_story_model()
-    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)
+    # Enhanced prompt template
+    prompt = f"""Write a children's story in 100-150 words with these elements:
+    - Theme: {keyword}
+    - Characters: Friendly animals
+    - Moral: Sharing is caring
+
+    Story begins: One sunny morning, a little rabbit named Cotton discovered"""
+    
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=150, truncation=True)
     outputs = model.generate(
         inputs.input_ids,
-        max_length=300,
-        temperature=0.85,
+        max_new_tokens=300,
+        temperature=0.9,
         top_k=50,
-        repetition_penalty=1.2
+        no_repeat_ngram_size=3,
+        repetition_penalty=1.2,
+        do_sample=True
     )
     full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return full_text.replace(prompt, "").strip()
+    return full_text.split("Story begins:")[-1].strip()
 
 # ======================
 # Stage 3: Text-to-Speech
@@ -63,8 +71,8 @@ Story: Once upon a time, there was a little rabbit named Fluffy who found"""
 def stage3_process(text):
     """Convert text to audio"""
     try:
-        clean_text = text.strip().replace('\n', ' ')[:200]
-        if len(clean_text) < 10:
+        clean_text = text.strip().replace('\n', ' ')[:300]
+        if len(clean_text) < 20:
             return None
         tts = gTTS(text=clean_text, lang='en')
         audio = io.BytesIO()
@@ -81,12 +89,11 @@ def main():
     st.title("📖 Children's Story Generator")
     
     # Initialize session state
-    if 'stage1_done' not in st.session_state:
+    if 'processing' not in st.session_state:
         st.session_state.update({
-            'stage1_done': False,
-            'stage2_done': False,
-            'caption': "",
-            'story': ""
+            'caption': None,
+            'story': None,
+            'audio': None
         })
     
     # File upload
@@ -97,31 +104,31 @@ def main():
         st.image(uploaded_file, width=300)
         
         # Stage 1
-        if not st.session_state.stage1_done:
+        if not st.session_state.caption:
             with st.spinner("Analyzing image..."):
                 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
-        if not st.session_state.stage2_done:
-            with st.spinner("Writing story..."):
+        if not st.session_state.story:
+            with st.spinner("Writing magical story..."):
                 st.session_state.story = stage2_process(st.session_state.caption)
-                st.session_state.stage2_done = True
         
-        # Display results
+        # Display story
         if st.session_state.story:
             st.subheader("Generated Story")
             st.write(st.session_state.story)
             
             # Stage 3
-            with st.spinner("Generating audio..."):
-                audio = stage3_process(st.session_state.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 short text")
+            if not st.session_state.audio:
+                with st.spinner("Generating audio..."):
+                    st.session_state.audio = stage3_process(st.session_state.story)
+            if st.session_state.audio:
+                st.audio(st.session_state.audio, format="audio/mp3")
+                st.download_button("Download Audio", 
+                                 st.session_state.audio.getvalue(),
+                                 "story.mp3",
+                                 mime="audio/mp3")
 
 if __name__ == "__main__":
     main()