Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from colab import AIDancePartner
 import tempfile
 import os
+import time
 
 def main():
     st.title("AI Dance Partner")
@@ -11,9 +12,10 @@ def main():
     uploaded_file = st.file_uploader("Choose a video file", type=['mp4', 'avi', 'mov'])
     
     if uploaded_file is not None:
-        # Save uploaded file temporarily
-        tfile = tempfile.NamedTemporaryFile(delete=False)
-        tfile.write(uploaded_file.read())
+        # Create a temporary file
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tfile:
+            tfile.write(uploaded_file.read())
+            temp_input_path = tfile.name
         
         # Dance style selection
         style = st.selectbox(
@@ -22,25 +24,32 @@ def main():
         )
         
         if st.button("Generate Dance Partner"):
-            with st.spinner("Processing your dance video..."):
-                # Process video
-                dance_partner = AIDancePartner()
-                output_path = dance_partner.process_video(tfile.name, mode=style)
-                
-                # Display result
-                st.video(output_path)
-                
-                # Download button
-                with open(output_path, 'rb') as file:
-                    st.download_button(
-                        label="Download video",
-                        data=file,
-                        file_name="ai_dance_partner.mp4",
-                        mime="video/mp4"
-                    )
+            try:
+                with st.spinner("Processing your dance video..."):
+                    # Process video
+                    dance_partner = AIDancePartner()
+                    output_path = dance_partner.process_video(temp_input_path, mode=style)
                     
-        # Cleanup
-        os.unlink(tfile.name)
+                    # Display result
+                    st.video(output_path)
+                    
+                    # Download button
+                    with open(output_path, 'rb') as file:
+                        st.download_button(
+                            label="Download video",
+                            data=file,
+                            file_name="ai_dance_partner.mp4",
+                            mime="video/mp4"
+                        )
+                    
+                    # Cleanup temporary files
+                    os.unlink(temp_input_path)
+                    os.unlink(output_path)
+                    
+            except Exception as e:
+                st.error(f"An error occurred: {str(e)}")
+                if os.path.exists(temp_input_path):
+                    os.unlink(temp_input_path)
 
 if __name__ == "__main__":
     main()