Update app.py

app.py

@@ -1,75 +1,158 @@
+import gradio as gr
+import os
+import shutil
+import asyncio
+import librosa
+import librosa.display
+import soundfile as sf
+import numpy as np
+import time
+import zipfile
+import tempfile
+import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.use('Agg')
+
+def update_output_visibility(choice):
+    if "2 Stems" in choice:
+        return { 
+            vocals_output: gr.update(visible=True), 
+            drums_output: gr.update(visible=False), 
+            bass_output: gr.update(visible=False), 
+            other_output: gr.update(visible=True, label="Instrumental (No Vocals)") 
+        }
+    elif "4 Stems" in choice:
+        return { 
+            vocals_output: gr.update(visible=True), 
+            drums_output: gr.update(visible=True), 
+            bass_output: gr.update(visible=True), 
+            other_output: gr.update(visible=True, label="Other") 
+        }
+
+async def separate_stems(audio_file_path, stem_choice, progress=gr.Progress(track_tqdm=True)):
+    if audio_file_path is None: raise gr.Error("No audio file uploaded!")
+    progress(0, desc="Starting...")
+    try:
+        progress(0.05, desc="Preparing audio file...")
+        original_filename_base = os.path.basename(audio_file_path).rsplit('.', 1)[0]
+        stable_input_path = f"stable_input_{original_filename_base}.wav"
+        shutil.copy(audio_file_path, stable_input_path)
+        
+        model_arg = "--two-stems=vocals" if "2 Stems" in stem_choice else ""
+        output_dir = "separated"
+        if os.path.exists(output_dir): shutil.rmtree(output_dir)
+        
+        command = f"python3 -m demucs {model_arg} -o \"{output_dir}\" \"{stable_input_path}\""
+        progress(0.2, desc="Running Demucs (this can take a minute)...")
+        
+        process = await asyncio.create_subprocess_shell(
+            command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE)
+        
+        stdout, stderr = await process.communicate()
+
+        if process.returncode != 0: 
+            raise gr.Error(f"Demucs failed to run. Error: {stderr.decode()[:500]}")
+
+        progress(0.8, desc="Locating separated stem files...")
+        stable_filename_base = os.path.basename(stable_input_path).rsplit('.', 1)[0]
+        model_folder_name = next(os.walk(output_dir))[1][0]
+        stems_path = os.path.join(output_dir, model_folder_name, stable_filename_base)
+        
+        if not os.path.exists(stems_path): 
+            raise gr.Error(f"Demucs finished, but the output directory was not found!")
+            
+        vocals_path = os.path.join(stems_path, "vocals.wav") if os.path.exists(os.path.join(stems_path, "vocals.wav")) else None
+        drums_path = os.path.join(stems_path, "drums.wav") if os.path.exists(os.path.join(stems_path, "drums.wav")) else None
+        bass_path = os.path.join(stems_path, "bass.wav") if os.path.exists(os.path.join(stems_path, "bass.wav")) else None
+        other_filename = "no_vocals.wav" if "2 Stems" in stem_choice else "other.wav"
+        other_path = os.path.join(stems_path, other_filename) if os.path.exists(os.path.join(stems_path, other_filename)) else None
+        
+        os.remove(stable_input_path)
+        return vocals_path, drums_path, bass_path, other_path
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        raise gr.Error(str(e))
+
+def visualize_slices(stem_audio_data, progress=gr.Progress(track_tqdm=True)):
+    if stem_audio_data is None:
+        gr.Warning("This stem is empty. Cannot visualize.")
+        return None, None, None
+    
+    sample_rate, y_int = stem_audio_data
+    y = librosa.util.buf_to_float(y_int)
+    
+    progress(0.3, desc="Finding transients...")
+    onset_frames = librosa.onset.onset_detect(y=librosa.to_mono(y.T) if y.ndim > 1 else y, sr=sample_rate, wait=1, pre_avg=1, post_avg=1, post_max=1, delta=0.05)
+    onset_times = librosa.frames_to_time(onset_frames, sr=sample_rate)
+    
+    progress(0.7, desc="Generating waveform plot...")
+    fig, ax = plt.subplots(figsize=(10, 3))
+    fig.patch.set_facecolor('#1f2937')
+    ax.set_facecolor('#111827')
+    librosa.display.waveshow(y, sr=sample_rate, ax=ax, color='#32f6ff', alpha=0.7)
+    for t in onset_times:
+        ax.axvline(x=t, color='#ff3b3b', linestyle='--', linewidth=1)
+    ax.tick_params(colors='gray'); ax.xaxis.label.set_color('gray'); ax.yaxis.label.set_color('gray')
+    ax.set_xlabel("Time (s)"); ax.set_ylabel("Amplitude"); ax.set_title("Detected Slices", color='white')
+    plt.tight_layout()
+    
+    progress(1, desc="Done!")
+    return fig, onset_times, stem_audio_data
+
+def preview_slice(active_stem_audio, onset_times, evt: gr.SelectData):
+    if active_stem_audio is None or onset_times is None: return None
+    sample_rate, y = active_stem_audio; clicked_time = evt.index[0]
+    start_time = 0; end_time = len(y) / sample_rate
+    for i, t in enumerate(onset_times):
+        if t > clicked_time:
+            end_time = t; break
+        start_time = t
+    start_sample = librosa.time_to_samples(start_time, sr=sample_rate)
+    end_sample = librosa.time_to_samples(end_time, sr=sample_rate)
+    sliced_audio = y[start_sample:end_sample]
+    return (sample_rate, sliced_audio)
+
 with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="red")) as demo:
     gr.Markdown("# 🎵 Loop Architect")
     onset_times_state = gr.State(value=None)
     active_stem_state = gr.State(value=None)
-    vocals_bar_times_state = gr.State(value=None)
-    drums_bar_times_state = gr.State(value=None)
-    bass_bar_times_state = gr.State(value=None)
-    other_bar_times_state = gr.State(value=None)
-
-
+    
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### 1. Separate Stems")
             audio_input = gr.Audio(type="filepath", label="Upload a Track")
             stem_options = gr.Radio(["4 Stems (Vocals, Drums, Bass, Other)", "2 Stems (Vocals + Instrumental)"], label="Separation Type", value="4 Stems (Vocals, Drums, Bass, Other)")
             submit_button = gr.Button("Separate Stems")
-
+            
         with gr.Column(scale=2):
             with gr.Accordion("Separated Stems", open=True):
-                with gr.Row():
-                    vocals_output = gr.Audio(label="Vocals", scale=2)
-                    with gr.Column(scale=1):
-                        slice_vocals_btn = gr.Button("Visualize Slices")
-                        vocals_loop_length = gr.Dropdown(choices=["4 Bars", "8 Bars", "16 Bars"], label="Loop Length", value="4 Bars")
-                        create_vocals_loop_btn = gr.Button("Create Loop")
-                    vocals_loop_output = gr.Audio(label="Vocals Loop", visible=False, scale=2)
-                    vocals_loop_download_btn = gr.DownloadButton(value="Download Loop", visible=False)
-                with gr.Row():
-                    drums_output = gr.Audio(label="Drums", scale=2)
-                    with gr.Column(scale=1):
-                        slice_drums_btn = gr.Button("Visualize Slices")
-                        drums_loop_length = gr.Dropdown(choices=["4 Bars", "8 Bars", "16 Bars"], label="Loop Length", value="4 Bars")
-                        create_drums_loop_btn = gr.Button("Create Loop")
-                    drums_loop_output = gr.Audio(label="Drums Loop", visible=False, scale=2)
-                    drums_loop_download_btn = gr.DownloadButton(value="Download Loop", visible=False)
-                with gr.Row():
-                    bass_output = gr.Audio(label="Bass", scale=2)
-                    with gr.Column(scale=1):
-                        slice_bass_btn = gr.Button("Visualize Slices")
-                        bass_loop_length = gr.Dropdown(choices=["4 Bars", "8 Bars", "16 Bars"], label="Loop Length", value="4 Bars")
-                        create_bass_loop_btn = gr.Button("Create Loop")
-                    bass_loop_output = gr.Audio(label="Bass Loop", visible=False, scale=2)
-                    bass_loop_download_btn = gr.DownloadButton(value="Download Loop", visible=False)
-                with gr.Row():
-                    other_output = gr.Audio(label="Other / Instrumental", scale=2)
-                    with gr.Column(scale=1):
-                        slice_other_btn = gr.Button("Visualize Slices")
-                        other_loop_length = gr.Dropdown(choices=["4 Bars", "8 Bars", "16 Bars"], label="Loop Length", value="4 Bars")
-                        create_other_loop_btn = gr.Button("Create Loop")
-                    other_loop_output = gr.Audio(label="Other Loop", visible=False, scale=2)
-                    other_loop_download_btn = gr.DownloadButton(value="Download Loop", visible=False)
-
-
+                with gr.Row(): 
+                    vocals_output = gr.Audio(label="Vocals", scale=4)
+                    slice_vocals_btn = gr.Button("Visualize Slices", scale=1)
+                with gr.Row(): 
+                    drums_output = gr.Audio(label="Drums", scale=4)
+                    slice_drums_btn = gr.Button("Visualize Slices", scale=1)
+                with gr.Row(): 
+                    bass_output = gr.Audio(label="Bass", scale=4)
+                    slice_bass_btn = gr.Button("Visualize Slices", scale=1)
+                with gr.Row(): 
+                    other_output = gr.Audio(label="Other / Instrumental", scale=4)
+                    slice_other_btn = gr.Button("Visualize Slices", scale=1)
+            
             gr.Markdown("### Slice Editor")
-            sensitivity_slider = gr.Slider(minimum=0, maximum=1, value=0.5, label="Onset Sensitivity")
-            slice_plot = gr.Image(label="Click a region on the waveform to preview a slice")
+            slice_plot = gr.Plot(label="Click a region on the waveform to preview a slice")
             preview_player = gr.Audio(label="Slice Preview")
 
-    submit_button.click(fn=separate_stems, inputs=[audio_input, stem_options], outputs=[vocals_output, drums_output, bass_output, other_output, vocals_bar_times_state, drums_bar_times_state, bass_bar_times_state, other_bar_times_state])
+    submit_button.click(fn=separate_stems, inputs=[audio_input, stem_options], outputs=[vocals_output, drums_output, bass_output, other_output])
     stem_options.change(fn=update_output_visibility, inputs=stem_options, outputs=[vocals_output, drums_output, bass_output, other_output])
-
-    slice_vocals_btn.click(fn=visualize_slices, inputs=[vocals_output, sensitivity_slider], outputs=[slice_plot, onset_times_state, active_stem_state])
-    slice_drums_btn.click(fn=visualize_slices, inputs=[drums_output, sensitivity_slider], outputs=[slice_plot, onset_times_state, active_stem_state])
-    slice_bass_btn.click(fn=visualize_slices, inputs=[bass_output, sensitivity_slider], outputs=[slice_plot, onset_times_state, active_stem_state])
-    slice_other_btn.click(fn=visualize_slices, inputs=[other_output, sensitivity_slider], outputs=[slice_plot, onset_times_state, active_stem_state])
-
+    
+    slice_vocals_btn.click(fn=visualize_slices, inputs=vocals_output, outputs=[slice_plot, onset_times_state, active_stem_state])
+    slice_drums_btn.click(fn=visualize_slices, inputs=drums_output, outputs=[slice_plot, onset_times_state, active_stem_state])
+    slice_bass_btn.click(fn=visualize_slices, inputs=bass_output, outputs=[slice_plot, onset_times_state, active_stem_state])
+    slice_other_btn.click(fn=visualize_slices, inputs=other_output, outputs=[slice_plot, onset_times_state, active_stem_state])
+    
     slice_plot.select(fn=preview_slice, inputs=[active_stem_state, onset_times_state], outputs=preview_player)
 
-    create_vocals_loop_btn.click(fn=create_loop, inputs=[vocals_output, vocals_bar_times_state, vocals_loop_length], outputs=[vocals_loop_output, vocals_loop_download_btn])
-    create_drums_loop_btn.click(fn=create_loop, inputs=[drums_output, drums_bar_times_state, drums_loop_length], outputs=[drums_loop_output, drums_loop_download_btn])
-    create_bass_loop_btn.click(fn=create_loop, inputs=[bass_output, bass_bar_times_state, bass_loop_length], outputs=[bass_loop_output, bass_loop_download_btn])
-    create_other_loop_btn.click(fn=create_loop, inputs=[other_output, other_bar_times_state, other_loop_length], outputs=[other_loop_output, other_loop_download_btn])
-
-
 demo.launch()