Upload app.py

app.py

@@ -33,8 +33,8 @@ class AIDancePartner:
         # Add playback controls
         play_speed = st.sidebar.slider("Playback Speed", 0.1, 2.0, 1.0, 0.1)
         
-        # Add play/pause button
-        is_playing = st.sidebar.button("Play/Pause")
+        # Add play/pause button with unique key
+        is_playing = st.sidebar.button("Play/Pause", key="play_pause_button")
         
         if video_file:
             self.process_video(video_file, mode, play_speed, is_playing)
@@ -59,6 +59,10 @@ class AIDancePartner:
         # Initialize progress bar
         progress_bar = st.progress(0)
         
+        # Initialize play state in session state
+        if 'play_state' not in st.session_state:
+            st.session_state.play_state = is_playing
+
         # Store video frames and AI sequences
         frames = []
         ai_sequences = []
@@ -81,12 +85,18 @@ class AIDancePartner:
                 frames.append(frame)
                 ai_sequences.append(ai_frame)
         
+        # Playback controls in sidebar
+        col1, col2 = st.sidebar.columns(2)
+        with col1:
+            stop_button = st.button("Stop", key="stop_button")
+        with col2:
+            replay_button = st.button("Replay", key="replay_button")
+
         # Playback loop
         frame_count = 0
-        play = True if is_playing else False
         
-        while True:
-            if play:
+        while not stop_button and frame_count < len(frames):
+            if st.session_state.play_state:
                 if frame_count >= len(frames):
                     frame_count = 0
                     
@@ -114,15 +124,15 @@ class AIDancePartner:
                 # Control playback speed
                 time.sleep(1 / (fps * play_speed))
                 frame_count += 1
-                
-            # Check for play/pause button
-            if st.sidebar.button("Stop"):
-                break
-                
-            # Add replay button
-            if frame_count >= len(frames):
-                if st.sidebar.button("Replay"):
-                    frame_count = 0
+            
+            # Handle replay
+            if replay_button:
+                frame_count = 0
+                st.session_state.play_state = True
+            
+            # Toggle play state when play/pause button is pressed
+            if is_playing:
+                st.session_state.play_state = not st.session_state.play_state
         
         # Cleanup
         cap.release()