app.py

import streamlit as st
from colab import AIDancePartner
import tempfile
import os

def main():
    st.title("AI Dance Partner")
    st.write("Upload your dance video and get an AI dance partner!")
    
    # Video upload
    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())
        
        # Dance style selection
        style = st.selectbox(
            "Choose dance partner style",
            ["Sync Partner", "Creative Partner"]
        )
        
        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"
                    )
                    
        # Cleanup
        os.unlink(tfile.name)

if __name__ == "__main__":
    main()