Update app.py

app.py

@@ -4,6 +4,7 @@ import torch
 from gtts import gTTS
 import io
 import time
+from streamlit.components.v1 import html
 
 # Initialize session state
 if 'processed_data' not in st.session_state:
@@ -13,29 +14,42 @@ if 'processed_data' not in st.session_state:
         'audio': None
     }
 
-if 'progress' not in st.session_state:
-    st.session_state.progress = {
-        'current_stage': '',
-        'percent': 0,
-        'start_time': None
+# JavaScript timer component
+def timer():
+    return """
+    <div id="timer" style="font-size:16px;color:#666;margin-bottom:10px;">⏱️ Elapsed: 00:00</div>
+    <script>
+    function updateTimer() {
+        var start = Date.now();
+        var timerElement = document.getElementById('timer');
+        
+        var interval = setInterval(function() {
+            var elapsed = Date.now() - start;
+            var minutes = Math.floor(elapsed / 60000);
+            var seconds = Math.floor((elapsed % 60000) / 1000);
+            timerElement.innerHTML = '⏱️ Elapsed: ' + 
+                (minutes < 10 ? '0' : '') + minutes + ':' + 
+                (seconds < 10 ? '0' : '') + seconds;
+        }, 1000);
+        
+        // Cleanup when component is removed
+        return function() {
+            clearInterval(interval);
+        }
     }
+    var cleanup = updateTimer();
+    
+    // Handle Streamlit's component cleanup
+    window.addEventListener('beforeunload', function() {
+        cleanup();
+    });
+    </script>
+    """
 
 # Page setup
 st.set_page_config(page_title="Your Image to Audio Story", page_icon="🦜")
 st.header("Turn Your Image to a Short Audio Story for Children")
 
-# Helper functions
-def update_progress(stage, percentage):
-    st.session_state.progress.update({
-        'current_stage': stage,
-        'percent': percentage,
-        'start_time': time.time() if percentage == 0 else st.session_state.progress['start_time']
-    })
-
-def get_elapsed_time():
-    elapsed = time.time() - st.session_state.progress['start_time']
-    return f"{int(elapsed//60):02d}:{int(elapsed%60):02d}"
-
 # Model loading
 @st.cache_resource
 def load_models():
@@ -58,7 +72,7 @@ def text2story(text):
     ]
     response = models["story_model"](
         messages,
-        max_new_tokens=100,
+        max_new_tokens=128,
         do_sample=True,
         temperature=0.7
     )[0]["generated_text"]
@@ -78,7 +92,9 @@ if uploaded_file is not None:
     # Initialize progress containers
     status_text = st.empty()
     progress_bar = st.progress(0)
-    time_text = st.empty()
+    
+    # Start JavaScript timer
+    html(timer(), height=50)
 
     try:
         # Save uploaded file
@@ -88,47 +104,35 @@ if uploaded_file is not None:
 
         if st.session_state.get('current_file') != uploaded_file.name:
             st.session_state.current_file = uploaded_file.name
-            update_progress("Starting processing...", 0)
 
             # Display image
             st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
 
             # Stage 1: Image to Text
-            update_progress("🖼️ Analyzing image...", 0)
-            status_text.markdown(f"""
-                **{st.session_state.progress['current_stage']}**  
-                ⏱️ Elapsed: {get_elapsed_time()}
-            """)
+            status_text.markdown("**🖼️ Analyzing image...**")
+            progress_bar.progress(0)
             st.session_state.processed_data['scenario'] = img2text(uploaded_file.name)
             progress_bar.progress(33)
 
             # Stage 2: Text to Story
-            update_progress("📖 Generating story...", 33)
-            status_text.markdown(f"""
-                **{st.session_state.progress['current_stage']}**  
-                ⏱️ Elapsed: {get_elapsed_time()}
-            """)
+            status_text.markdown("**📖 Generating story...**")
+            progress_bar.progress(33)
             st.session_state.processed_data['story'] = text2story(
                 st.session_state.processed_data['scenario']
             )
             progress_bar.progress(66)
 
             # Stage 3: Story to Audio
-            update_progress("🔊 Synthesizing audio...", 66)
-            status_text.markdown(f"""
-                **{st.session_state.progress['current_stage']}**  
-                ⏱️ Elapsed: {get_elapsed_time()}
-            """)
+            status_text.markdown("**🔊 Synthesizing audio...**")
+            progress_bar.progress(66)
             st.session_state.processed_data['audio'] = text2audio(
                 st.session_state.processed_data['story']
             )
             progress_bar.progress(100)
 
             # Final status
-            status_text.success(f"""
-                ✅ Generation complete!  
-                Total time: {get_elapsed_time()}
-            """)
+            status_text.success("**✅ Generation complete!**")
+            html("<script>document.getElementById('timer').style.color = '#00cc00';</script>")
 
         # Show results
         st.subheader("Results")
@@ -136,12 +140,13 @@ if uploaded_file is not None:
         st.write("**Story:**", st.session_state.processed_data['story'])
 
     except Exception as e:
-        status_text.error(f"❌ Error: {str(e)}")
+        html("<script>document.getElementById('timer').remove();</script>")
+        status_text.error(f"**❌ Error:** {str(e)}")
         progress_bar.empty()
         raise e
 
     finally:
-        time_text.empty()
+        pass  # Timer cleanup handled by JavaScript
 
 # Audio playback
 if st.button("Play Audio of the Story Generated"):