Update app.py

app.py

@@ -11,6 +11,8 @@ from concurrent.futures import Future
 import torch
 import gradio as gr
 import pydub
+from scipy.io.wavfile import write
+import subprocess
 
 from audiocraft.data.audio_utils import convert_audio
 from audiocraft.data.audio import audio_write
@@ -119,12 +121,10 @@ def _do_predictions(texts, melodies, duration, progress=False, **gen_kwargs):
             audio_write(
                 file.name, output, MODEL.sample_rate, strategy="loudness",
                 loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)
-            out_files.append(file.name)  # Store the filename as a string
+            out_files.append(pool.submit(make_waveform, file.name))
             file_cleaner.add(file.name)
-    res = [out_file for out_file in out_files]
+    res = [out_file.result() for out_file in out_files]
     for file in res:
-        if isinstance(file, Future):  # Check if it's a Future object
-            file = file.result()  # Extract the filename from the Future object
         file_cleaner.add(file)
     print("batch finished", len(texts), time.time() - be)
     print("Tempfiles currently stored: ", len(file_cleaner.files))
@@ -171,6 +171,28 @@ def toggle_audio_src(choice):
         return gr.update(source="upload", value=None, label="File")
 
 
+def inference(audio):
+    os.makedirs("out", exist_ok=True)
+    write('test.wav', audio[0], audio[1])
+
+    command = "python3 -m demucs.separate -n mdx_extra_q -d cpu test.wav -o out"
+    process = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    print("Demucs script output:", process.stdout.decode())
+
+    # Check if files exist before returning
+    files = ["./out/mdx_extra_q/test/vocals.wav",
+             "./out/mdx_extra_q/test/bass.wav",
+             "./out/mdx_extra_q/test/drums.wav",
+             "./out/mdx_extra_q/test/other.wav"]
+    for file in files:
+        if not os.path.isfile(file):
+            print(f"File not found: {file}")
+        else:
+            print(f"File exists: {file}")
+
+    return files
+
+
 def ui_full(launch_kwargs):
     with gr.Blocks() as interface:
         gr.Markdown(
@@ -193,6 +215,7 @@ def ui_full(launch_kwargs):
                 with gr.Row():
                     submit = gr.Button("Submit")
                     # Adapted from https://github.com/rkfg/audiocraft/blob/long/app.py, MIT license.
+                    _ = gr.Button("Interrupt").click(fn=interrupt, queue=False)
                 with gr.Row():
                     model = gr.Radio(["melody", "medium", "small", "large"],
                                      label="Model", value="melody", interactive=True)
@@ -204,7 +227,7 @@ def ui_full(launch_kwargs):
                     temperature = gr.Number(label="Temperature", value=1.0, interactive=True)
                     cfg_coef = gr.Number(label="Classifier Free Guidance", value=3.0, interactive=True)
             with gr.Column():
-                output = gr.Audio(label="Generated Music")
+                output = gr.Video(label="Generated Music")
         submit.click(predict_full,
                      inputs=[model, text, melody, duration, topk, topp, temperature, cfg_coef],
                      outputs=[output])
@@ -298,7 +321,7 @@ def ui_batched(launch_kwargs):
                 with gr.Row():
                     submit = gr.Button("Generate")
             with gr.Column():
-                output = gr.Audio(label="Generated Music")
+                output = gr.Video(label="Generated Music")
         submit.click(predict_batched, inputs=[text, melody],
                      outputs=[output], batch=True, max_batch_size=MAX_BATCH_SIZE)
         radio.change(toggle_audio_src, radio, [melody], queue=False, show_progress=False)
@@ -343,6 +366,46 @@ def ui_batched(launch_kwargs):
         demo.queue(max_size=8 * 4).launch(**launch_kwargs)
 
 
+def convert_to_audio(file_path):
+    sound = pydub.AudioSegment.from_file(file_path)
+    samples = sound.get_array_of_samples()
+    sample_rate = sound.frame_rate
+    audio = np.array(samples).reshape((-1, 2))
+    audio = audio.T / 32768.0
+    return audio, sample_rate
+
+
+def separate_audio(source_file_path, target_folder_path):
+    command = f"python3 -m demucs.separate -n demucs_extra -d cpu {source_file_path} -o {target_folder_path}"
+    subprocess.call(command, shell=True)
+
+
+def separate_music_audio(audio_path):
+    audio, sample_rate = convert_to_audio(audio_path)
+    audio_file_name = os.path.basename(audio_path)
+    temp_dir = f"./temp/{audio_file_name}"
+    os.makedirs(temp_dir, exist_ok=True)
+    temp_audio_path = os.path.join(temp_dir, audio_file_name)
+    write(temp_audio_path, sample_rate, audio.T)
+    separate_audio(temp_audio_path, temp_dir)
+    return temp_dir
+
+
+def postprocess_audio(audio_dir):
+    separated_audio_files = {}
+    for file in os.listdir(audio_dir):
+        if file.endswith(".wav"):
+            instrument = file.split(".")[0]
+            separated_audio_files[instrument] = os.path.join(audio_dir, file)
+    return separated_audio_files
+
+
+def separate_music(audio_path):
+    audio_dir = separate_music_audio(audio_path)
+    separated_audio_files = postprocess_audio(audio_dir)
+    return separated_audio_files
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(