Update app.py

app.py

@@ -51,7 +51,10 @@ HEADER = """
 
 CACHE_EXAMPLES = os.getenv('CACHE_EXAMPLES', '1') == '1'
 
-def compress_files(folder_path, dissector_file, original_audio):
+BASIC_PITCH_DIR = "/app/basic-pitch"
+
+
+def compress_files(demucs_output_path, dissector_file, original_audio):
     """Compresses specific files in the specified folder into a .zip file"""
     # List of the files to be compressed
     wav_files = ['bass.wav', 'drums.wav', 'other.wav', 'vocals.wav']
@@ -59,7 +62,7 @@ def compress_files(folder_path, dissector_file, original_audio):
 
     def convert_to_mp3(wav, mp3):
         """Utility function to run the ffmpeg command"""
-        command = ["ffmpeg", "-i", os.path.join(folder_path, wav), os.path.join(folder_path, mp3)]
+        command = ["ffmpeg", "-i", os.path.join(demucs_output_path, wav), os.path.join(demucs_output_path, mp3)]
         process = Popen(command)
         process.communicate()
 
@@ -67,11 +70,26 @@ def compress_files(folder_path, dissector_file, original_audio):
     with ThreadPoolExecutor(max_workers=8) as executor:
         list(executor.map(convert_to_mp3, wav_files, mp3_files))
 
+    if not os.path.exists(basic_pitch_dir):
+        os.makedirs(basic_pitch_dir)
+    command = [
+        "basic-pitch",
+        BASIC_PITCH_DIR,
+        os.path.join(demucs_output_path, 'bass.wav'),
+        os.path.join(demucs_output_path, 'vocals.wav'),
+        os.path.join(demucs_output_path, 'other.wav')
+    ]
+    process = Popen(command)
+    process.communicate()
+
     # Compress the files
-    zip_path = folder_path + ".zip"
+    zip_path = demucs_output_path + ".zip"
     with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
         for mp3 in mp3_files:
-            zipf.write(os.path.join(folder_path, mp3), arcname=mp3)
+            zipf.write(os.path.join(demucs_output_path, mp3), arcname=mp3)
+        zipf.write('/app/basic-pitch/bass_basic_pitch.mid', arcname='bass_basic_pitch.mid')
+        zipf.write('/app/basic-pitch/vocals_basic_pitch.mid', arcname='vocals_basic_pitch.mid')
+        zipf.write('/app/basic-pitch/other_basic_pitch.mid', arcname='other_basic_pitch.mid')
         zipf.write(dissector_file, arcname='dissector.json')  # Fixed name
         zipf.write(original_audio, arcname='mixdown.mp3')  # Added original audio