Update app.py

app.py

@@ -1,106 +1,65 @@
-import os
-import re
-import random
-from scipy.io.wavfile import write, read
-import numpy as np
 import gradio as gr
-import yt_dlp
-
-# Model dictionaries and lists
-uvr5_models = {
-    'model_bs_roformer_ep_317_sdr_12.9755.ckpt',
-    'model_bs_roformer_ep_368_sdr_12.9628.ckpt',
-    'model_bs_roformer_ep_937_sdr_10.5309.ckpt',
-    'model_mel_band_roformer_ep_3005_sdr_11.4360.ckpt',
-    'MDX23C-8KFFT-InstVoc_HQ.ckpt',
-    'MDX23C-8KFFT-InstVoc_HQ_2.ckpt',
-    'Kim_Vocal_1.onnx',
-    'Kim_Vocal_2.onnx',
-    'Kim_Inst.onnx',
-    'Reverb_HQ_By_FoxJoy.onnx',
-    'UVR-MDX-NET_Crowd_HQ_1.onnx',
-    'UVR-De-Echo-Aggressive.pth',
-    'UVR-De-Echo-Normal.pth',
-    'UVR-DeEcho-DeReverb.pth',
-    'UVR-DeNoise-Lite.pth',
-    'UVR-DeNoise.pth',
-    'UVR-BVE-4B_SN-44100-1.pth',
-    'htdemucs_ft.yaml', 
-    'htdemucs.yaml',
-    'hdemucs_mmi.yaml',
-}
-
+import os
+from audio_separator import Separator
 
+def separate_audio(input_file, output_dir, model_name):
+    # Create output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
 
-# More model lists...
+    # Initialize the Separator
+    separator = Separator()
 
-output_format = ['wav', 'flac', 'mp3']
-mdxnet_overlap_values = ['0.25', '0.5', '0.75', '0.99']
-vrarch_window_size_values = ['320', '512', '1024']
-demucs_overlap_values = ['0.25', '0.50', '0.75', '0.99']
+    # Separate the audio
+    separator.separate_audio_file(
+        input_file,
+        output_dir,
+        model_name=model_name,
+        denoise=True,
+        output_format='wav',
+        normalization_threshold=0.9,
+        mdx_segment_size=256,
+        mdx_overlap=8,
+        primary_stem_only=False
+    )
 
-# Function to download audio
-def download_audio(url):
-    ydl_opts = {
-        'format': 'bestaudio/best',
-        'outtmpl': 'ytdl/%(title)s.%(ext)s',
-        'postprocessors': [{
-            'key': 'FFmpegExtractAudio',
-            'preferredcodec': 'wav',
-            'preferredquality': '192',
-        }],
-    }
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        info_dict = ydl.extract_info(url, download=True)
-        file_path = ydl.prepare_filename(info_dict).rsplit('.', 1)[0] + '.wav'
-        sample_rate, audio_data = read(file_path)
-        audio_array = np.asarray(audio_data, dtype=np.int16)
-        return sample_rate, audio_array
+    # Rename the output files to match the requested format
+    os.rename(os.path.join(output_dir, 'Vocals.wav'), os.path.join(output_dir, '1_main_vocal.wav'))
+    os.rename(os.path.join(output_dir, 'Other.wav'), os.path.join(output_dir, '2_backing_vocal.wav'))
+    os.rename(os.path.join(output_dir, 'Instrumental.wav'), os.path.join(output_dir, '3_instrumental.wav'))
 
-# Function to separate audio using Roformer
-def roformer_separator(audio, model, output_format, overlap, segment_size, denoise):
-    directory = "./outputs"
-    random_id = str(random.randint(10000, 99999))
-    os.makedirs("outputs", exist_ok=True)
-    write(f'{random_id}.wav', audio[0], audio[1])
-    full_roformer_model = uvr5_models[model]
-    
-    prompt = f"audio-separator {random_id}.wav --model_filename {full_roformer_model} --output_dir=./outputs --output_format={output_format} --normalization=0.9 --mdxc_overlap={overlap} --mdxc_segment_size={segment_size}"
-    if denoise:
-        prompt += " --mdx_enable_denoise"
-    
-    os.system(prompt)
+    return [
+        os.path.join(output_dir, '1_main_vocal.wav'),
+        os.path.join(output_dir, '2_backing_vocal.wav'),
+        os.path.join(output_dir, '3_instrumental.wav')
+    ]
 
-    files_list = [os.path.join(directory, file) for file in os.listdir(directory) if re.search(random_id, file)]
-    stem1_file, stem2_file, stem3_file = files_list[:3]  # Assuming the files are in the correct order
-    return stem1_file, stem2_file, stem3_file
+def process_audio(audio_file, model_name):
+    output_dir = "output"
+    return separate_audio(audio_file.name, output_dir, model_name)
 
-# Gradio interface
-def process_audio(url, model, output_format, overlap, segment_size, denoise):
-    sample_rate, audio_array = download_audio(url)
-    stems = roformer_separator((sample_rate, audio_array), model, output_format, overlap, segment_size, denoise)
-    return stems
+# Define the Gradio interface
+iface = gr.Blocks()
 
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("# Hex Audio Separator")
-    url_input = gr.Textbox(label="YouTube URL")
+with iface:
+    gr.Markdown("# Audio Separator")
     with gr.Row():
-        model_input = gr.Dropdown(choices=list(uvr5_models()), label="Roformer Model")
-        format_input = gr.Dropdown(choices=output_format, label="Output Format")
-        overlap_input = gr.Dropdown(choices=mdxnet_overlap_values, label="Overlap")
-    segment_input = gr.Slider(0, 100, label="Segment Size")
-    denoise_input = gr.Checkbox(label="Enable Denoise")
-
-    output1 = gr.Audio(label="Vocals")
-    output2 = gr.Audio(label="Instrumental")
-    output3 = gr.Audio(label="Backing Vocals")
+        with gr.Column():
+            audio_input = gr.Audio(type="filepath", label="Input Audio")
+            model_name = gr.Dropdown(
+                choices=["UVR-MDX-NET-Inst_HQ_3", "UVR_MDXNET_KARA_2", "UVR-MDX-NET-Inst_HQ_4"],
+                label="Model",
+                value="UVR-MDX-NET-Inst_HQ_3"
+            )
+            submit_btn = gr.Button("Separate Audio")
+        with gr.Column():
+            vocal_output = gr.Audio(label="Main Vocal")
+            backing_vocal_output = gr.Audio(label="Backing Vocal")
+            instrumental_output = gr.Audio(label="Instrumental")
 
-    submit_button = gr.Button("Process")
-    submit_button.click(
+    submit_btn.click(
         process_audio,
-        inputs=[url_input, model_input, format_input, overlap_input, segment_input, denoise_input],
-        outputs=[output1, output2, output3]
+        inputs=[audio_input, model_name],
+        outputs=[vocal_output, backing_vocal_output, instrumental_output]
     )
 
-demo.launch()
+iface.launch()