Update app.py

app.py

@@ -9,6 +9,7 @@ import io
 import time
 from streamlit.components.v1 import html
 import asyncio
+
 if not asyncio.get_event_loop().is_running():
     asyncio.set_event_loop(asyncio.new_event_loop())
     
@@ -25,11 +26,12 @@ def timer():
     return """
     <div id="timer" style="font-size:16px;color:#666;margin-bottom:10px;">⏱️ Elapsed: 00:00</div>
     <script>
+    var timerInterval;
     function updateTimer() {
         var start = Date.now();
         var timerElement = document.getElementById('timer');
         
-        var interval = setInterval(function() {
+        timerInterval = setInterval(function() {
             var elapsed = Date.now() - start;
             var minutes = Math.floor(elapsed / 60000);
             var seconds = Math.floor((elapsed % 60000) / 1000);
@@ -37,17 +39,18 @@ def timer():
                 (minutes < 10 ? '0' : '') + minutes + ':' + 
                 (seconds < 10 ? '0' : '') + seconds;
         }, 1000);
-        
-        // Cleanup when component is removed
-        return function() {
-            clearInterval(interval);
-        }
     }
-    var cleanup = updateTimer();
+    updateTimer();
+    
+    // Listen for custom stop event
+    document.addEventListener('stopTimer', function() {
+        clearInterval(timerInterval);
+        document.getElementById('timer').style.color = '#00cc00';
+    });
     
-    // Handle Streamlit's component cleanup
+    // Cleanup on error or page unload
     window.addEventListener('beforeunload', function() {
-        cleanup();
+        clearInterval(timerInterval);
     });
     </script>
     """
@@ -61,7 +64,19 @@ st.header("Turn Your Image to a Short Audio Story for Children")
 def load_models():
     return {
         "img_model": pipeline("image-to-text", "cnmoro/tiny-image-captioning"),
-        "story_model": pipeline("text-generation", "Qwen/Qwen2.5-0.5B-Instruct")
+        "story_model": pipeline(
+            "text-generation", 
+            "Qwen/Qwen2.5-0.5B-Instruct",
+            device_map="auto",
+            model_kwargs={
+                "load_in_4bit": True,
+                "quantization_config": {
+                    "load_in_4bit": True,
+                    "bnb_4bit_compute_dtype": torch.float16,
+                    "bnb_4bit_quant_type": "nf4"
+                }
+            }
+        )
     }
 
 models = load_models()
@@ -136,9 +151,13 @@ if uploaded_file is not None:
             )
             progress_bar.progress(100)
 
-            # Final status
+            # Final status and stop timer
             status_text.success("**✅ Generation complete!**")
-            html("<script>document.getElementById('timer').style.color = '#00cc00';</script>")
+            html("""
+                <script>
+                    document.dispatchEvent(new Event('stopTimer'));
+                </script>
+            """)
 
         # Show results
         st.subheader("Results")
@@ -146,14 +165,16 @@ if uploaded_file is not None:
         st.write("**Story:**", st.session_state.processed_data['story'])
 
     except Exception as e:
-        html("<script>document.getElementById('timer').remove();</script>")
+        html("""
+            <script>
+                document.dispatchEvent(new Event('stopTimer'));
+                document.getElementById('timer').remove();
+            </script>
+        """)
         status_text.error(f"**❌ Error:** {str(e)}")
         progress_bar.empty()
         raise e
 
-    finally:
-        pass  # Timer cleanup handled by JavaScript
-
 # Audio playback
 if st.button("Play Audio of the Story Generated"):
     if st.session_state.processed_data.get('audio'):