Update app.py

app.py

@@ -4,7 +4,7 @@ import shutil
 import logging
 import gradio as gr
 
-from audio_separator.separator import Separator
+from PolUVR.separator import Separator
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 use_autocast = device == "cuda"
@@ -13,15 +13,15 @@ use_autocast = device == "cuda"
 #     Roformer Models     #
 #=========================#
 ROFORMER_MODELS = {
-    'BS-Roformer-De-Reverb': 'deverb_bs_roformer_8_384dim_10depth.ckpt',
     'BS-Roformer-Viperx-1053': 'model_bs_roformer_ep_937_sdr_10.5309.ckpt',
     'BS-Roformer-Viperx-1296': 'model_bs_roformer_ep_368_sdr_12.9628.ckpt',
     'BS-Roformer-Viperx-1297': 'model_bs_roformer_ep_317_sdr_12.9755.ckpt',
+    'BS-Roformer-De-Reverb': 'deverb_bs_roformer_8_384dim_10depth.ckpt',
+    'Mel-Roformer-Viperx-1143': 'model_mel_band_roformer_ep_3005_sdr_11.4360.ckpt',
     'Mel-Roformer-Crowd-Aufr33-Viperx': 'mel_band_roformer_crowd_aufr33_viperx_sdr_8.7144.ckpt',
+    'Mel-Roformer-Karaoke-Aufr33-Viperx': 'mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt',
     'Mel-Roformer-Denoise-Aufr33': 'denoise_mel_band_roformer_aufr33_sdr_27.9959.ckpt',
     'Mel-Roformer-Denoise-Aufr33-Aggr': 'denoise_mel_band_roformer_aufr33_aggr_sdr_27.9768.ckpt',
-    'Mel-Roformer-Karaoke-Aufr33-Viperx': 'mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt',
-    'Mel-Roformer-Viperx-1143': 'model_mel_band_roformer_ep_3005_sdr_11.4360.ckpt',
     'MelBand Roformer Kim | Inst V1 by Unwa': 'melband_roformer_inst_v1.ckpt',
     'MelBand Roformer Kim | Inst V2 by Unwa': 'melband_roformer_inst_v2.ckpt',
     'MelBand Roformer Kim | InstVoc Duality V1 by Unwa': 'melband_roformer_instvoc_duality_v1.ckpt',
@@ -39,7 +39,6 @@ MDX23C_MODELS = [
 #     MDXN-NET Models     #
 #=========================#
 MDXNET_MODELS = [
-    'UVR-MDX-NET-Crowd_HQ_1.onnx',
     'UVR-MDX-NET-Inst_1.onnx',
     'UVR-MDX-NET-Inst_2.onnx',
     'UVR-MDX-NET-Inst_3.onnx',
@@ -48,16 +47,17 @@ MDXNET_MODELS = [
     'UVR-MDX-NET-Inst_HQ_3.onnx',
     'UVR-MDX-NET-Inst_HQ_4.onnx',
     'UVR-MDX-NET-Inst_HQ_5.onnx',
-    'UVR-MDX-NET-Inst_full_292.onnx',
-    'UVR-MDX-NET-Voc_FT.onnx',
     'UVR-MDX-NET_Inst_82_beta.onnx',
     'UVR-MDX-NET_Inst_90_beta.onnx',
     'UVR-MDX-NET_Inst_187_beta.onnx',
+    'UVR-MDX-NET-Inst_full_292.onnx',
     'UVR-MDX-NET_Main_340.onnx',
     'UVR-MDX-NET_Main_390.onnx',
     'UVR-MDX-NET_Main_406.onnx',
     'UVR-MDX-NET_Main_427.onnx',
     'UVR-MDX-NET_Main_438.onnx',
+    'UVR-MDX-NET-Crowd_HQ_1.onnx',
+    'UVR-MDX-NET-Voc_FT.onnx',
     'UVR_MDXNET_1_9703.onnx',
     'UVR_MDXNET_2_9682.onnx',
     'UVR_MDXNET_3_9662.onnx',
@@ -124,7 +124,7 @@ def print_message(input_file, model_name):
     """Prints information about the audio separation process."""
     base_name = os.path.splitext(os.path.basename(input_file))[0]
     print("\n")
-    print("🎵 Audio-Separator 🎵")
+    print("🎵 PolUVR 🎵")
     print("Input audio:", base_name)
     print("Separation Model:", model_name)
     print("Audio Separation Process...")
@@ -141,9 +141,22 @@ def prepare_output_dir(input_file, output_dir):
         raise RuntimeError(f"Failed to prepare output directory {out_dir}: {e}")
     return out_dir
 
-def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, batch_size, progress=gr.Progress(track_tqdm=True)):
-    """Separate audio using Roformer model."""
+def rename_stems(audio, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem):
     base_name = os.path.splitext(os.path.basename(audio))[0]
+    stems = {
+        "Vocals": vocals_stem.replace("{base_name}", base_name),
+        "Instrumental": instrumental_stem.replace("{base_name}", base_name),
+        "Drums": drums_stem.replace("{base_name}", base_name),
+        "Bass": bass_stem.replace("{base_name}", base_name),
+        "Other": other_stem.replace("{base_name}", base_name),
+        "Guitar": guitar_stem.replace("{base_name}", base_name),
+        "Piano": piano_stem.replace("{base_name}", base_name),
+    }
+    return stems
+
+def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, batch_size, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem, progress=gr.Progress(track_tqdm=True)):
+    """Separate audio using Roformer model."""
+    stemname = rename_stems(audio, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem)
     print_message(audio, model_key)
     model = ROFORMER_MODELS[model_key]
     try:
@@ -169,17 +182,17 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
         separator.load_model(model_filename=model)
 
         progress(0.7, desc="Audio separated...")
-        separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
+        separation = separator.separate(audio, stemname)
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
         stems = [os.path.join(out_dir, file_name) for file_name in separation]
-        return stems[1], stems[0]
+        return stems[0], stems[1]
     except Exception as e:
         raise RuntimeError(f"Roformer separation failed: {e}") from e
 
-def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, batch_size, progress=gr.Progress(track_tqdm=True)):
+def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, batch_size, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using MDX23C model."""
-    base_name = os.path.splitext(os.path.basename(audio))[0]
+    stemname = rename_stems(audio, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem)
     print_message(audio, model)
     try:
         out_dir = prepare_output_dir(audio, out_dir)
@@ -204,17 +217,17 @@ def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_s
         separator.load_model(model_filename=model)
 
         progress(0.7, desc="Audio separated...")
-        separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
+        separation = separator.separate(audio, stemname)
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
         stems = [os.path.join(out_dir, file_name) for file_name in separation]
-        return stems[1], stems[0]
+        return stems[0], stems[1]
     except Exception as e:
         raise RuntimeError(f"MDX23C separation failed: {e}") from e
 
-def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_dir, out_dir, out_format, norm_thresh, amp_thresh, batch_size, progress=gr.Progress(track_tqdm=True)):
+def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_dir, out_dir, out_format, norm_thresh, amp_thresh, batch_size, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using MDX-NET model."""
-    base_name = os.path.splitext(os.path.basename(audio))[0]
+    stemname = rename_stems(audio, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem)
     print_message(audio, model)
     try:
         out_dir = prepare_output_dir(audio, out_dir)
@@ -239,7 +252,7 @@ def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_di
         separator.load_model(model_filename=model)
 
         progress(0.7, desc="Audio separated...")
-        separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
+        separation = separator.separate(audio, stemname)
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
         stems = [os.path.join(out_dir, file_name) for file_name in separation]
@@ -247,9 +260,9 @@ def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_di
     except Exception as e:
         raise RuntimeError(f"MDX-NET separation failed: {e}") from e
 
-def vr_separator(audio, model, window_size, aggression, tta, post_process, post_process_threshold, high_end_process, model_dir, out_dir, out_format, norm_thresh, amp_thresh, batch_size, progress=gr.Progress(track_tqdm=True)):
+def vr_separator(audio, model, window_size, aggression, tta, post_process, post_process_threshold, high_end_process, model_dir, out_dir, out_format, norm_thresh, amp_thresh, batch_size, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using VR ARCH model."""
-    base_name = os.path.splitext(os.path.basename(audio))[0]
+    stemname = rename_stems(audio, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem)
     print_message(audio, model)
     try:
         out_dir = prepare_output_dir(audio, out_dir)
@@ -276,7 +289,7 @@ def vr_separator(audio, model, window_size, aggression, tta, post_process, post_
         separator.load_model(model_filename=model)
 
         progress(0.7, desc="Audio separated...")
-        separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
+        separation = separator.separate(audio, stemname)
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
         stems = [os.path.join(out_dir, file_name) for file_name in separation]
@@ -284,8 +297,9 @@ def vr_separator(audio, model, window_size, aggression, tta, post_process, post_
     except Exception as e:
         raise RuntimeError(f"VR ARCH separation failed: {e}") from e
 
-def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True)):
+def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled, model_dir, out_dir, out_format, norm_thresh, amp_thresh, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using Demucs model."""
+    stemname = rename_stems(audio, vocals_stem, instrumental_stem, other_stem, drums_stem, bass_stem, guitar_stem, piano_stem)
     print_message(audio, model)
     try:
         out_dir = prepare_output_dir(audio, out_dir)
@@ -309,7 +323,7 @@ def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled,
         separator.load_model(model_filename=model)
 
         progress(0.7, desc="Audio separated...")
-        separation = separator.separate(audio)
+        separation = separator.separate(audio, stemname)
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
         stems = [os.path.join(out_dir, file_name) for file_name in separation]
@@ -328,15 +342,15 @@ def update_stems(model):
         return gr.update(visible=False)
 
 with gr.Blocks(
-    title="🎵 Audio-Separator 🎵",
+    title="🎵 PolUVR 🎵",
     css="footer{display:none !important}",
     theme=gr.themes.Default(
         spacing_size="sm",
         radius_size="lg",
     )
 ) as app:
-    gr.HTML("<h1> 🎵 Audio-Separator 🎵 </h1>")
-    
+    gr.HTML("<h1> 🎵 PolUVR 🎵 </h1>")
+
     with gr.Tab("Roformer"):
         with gr.Group():
             with gr.Row():
@@ -430,20 +444,32 @@ with gr.Blocks(
             demucs_stem5 = gr.Audio(label="Stem 5", type="filepath", interactive=False)
             demucs_stem6 = gr.Audio(label="Stem 6", type="filepath", interactive=False)
 
-
-    with gr.Tab("General settings"):
-        with gr.Group():
-            model_file_dir = gr.Textbox(value="/tmp/audio-separator-models/", label="Directory to cache model files", info="The directory where model files are stored.", placeholder="/tmp/audio-separator-models/")
+    with gr.Tab("Settings"):
+        with gr.Accordion("General settings", open=False):
+          with gr.Group():
+              model_file_dir = gr.Textbox(value="/tmp/PolUVR-models/", label="Directory to cache model files", info="The directory where model files are stored.", placeholder="/tmp/PolUVR-models/")
+              with gr.Row():
+                  output_dir = gr.Textbox(value="output", label="File output directory", info="The directory where output files will be saved.", placeholder="output")
+                  output_format = gr.Dropdown(value="wav", choices=["wav", "flac", "mp3"], label="Output Format", info="The format of the output audio file.")
+              with gr.Row():
+                  norm_threshold = gr.Slider(minimum=0.1, maximum=1, step=0.1, value=0.9, label="Normalization threshold", info="The threshold for audio normalization.")
+                  amp_threshold = gr.Slider(minimum=0.1, maximum=1, step=0.1, value=0.6, label="Amplification threshold", info="The threshold for audio amplification.")
+              with gr.Row():
+                  batch_size = gr.Slider(minimum=1, maximum=16, step=1, value=1, label="Batch Size", info="Larger consumes more RAM but may process slightly faster.")
+
+        with gr.Accordion("Rename Stems", open=False):
+            gr.HTML("<h3> `{base_name}` - input file name </h3>")
             with gr.Row():
-                output_dir = gr.Textbox(value="output", label="File output directory", info="The directory where output files will be saved.", placeholder="output")
-                output_format = gr.Dropdown(value="wav", choices=["wav", "flac", "mp3"], label="Output Format", info="The format of the output audio file.")
+                vocals_stem = gr.Textbox(value="{base_name}_(Vocals)", label="Vocals Stem", placeholder="{base_name}_(Vocals)")
+                instrumental_stem = gr.Textbox(value="{base_name}_(Instrumental)", label="Instrumental Stem", placeholder="{base_name}_(Instrumental)")
+                other_stem = gr.Textbox(value="{base_name}_(Other)", label="Other Stem", placeholder="{base_name}_(Other)")
             with gr.Row():
-                norm_threshold = gr.Slider(minimum=0.1, maximum=1, step=0.1, value=0.9, label="Normalization threshold", info="The threshold for audio normalization.")
-                amp_threshold = gr.Slider(minimum=0.1, maximum=1, step=0.1, value=0.6, label="Amplification threshold", info="The threshold for audio amplification.")
+                drums_stem = gr.Textbox(value="{base_name}_(Drums)", label="Drums Stem", placeholder="{base_name}_(Drums)")
+                bass_stem = gr.Textbox(value="{base_name}_(Bass)", label="Bass Stem", placeholder="{base_name}_(Bass)")
             with gr.Row():
-                batch_size = gr.Slider(minimum=1, maximum=16, step=1, value=1, label="Batch Size", info="Larger consumes more RAM but may process slightly faster.")
+                guitar_stem = gr.Textbox(value="{base_name}_(Guitar)", label="Guitar Stem", placeholder="{base_name}_(Guitar)")
+                piano_stem = gr.Textbox(value="{base_name}_(Piano)", label="Piano Stem", placeholder="{base_name}_(Piano)")
 
-            
     demucs_model.change(update_stems, inputs=[demucs_model], outputs=stem6)
 
     roformer_button.click(
@@ -461,6 +487,13 @@ with gr.Blocks(
             norm_threshold,
             amp_threshold,
             batch_size,
+            vocals_stem,
+            instrumental_stem,
+            other_stem,
+            drums_stem,
+            bass_stem,
+            guitar_stem,
+            piano_stem,
         ],
         outputs=[roformer_stem1, roformer_stem2],
     )
@@ -479,6 +512,13 @@ with gr.Blocks(
             norm_threshold,
             amp_threshold,
             batch_size,
+            vocals_stem,
+            instrumental_stem,
+            other_stem,
+            drums_stem,
+            bass_stem,
+            guitar_stem,
+            piano_stem,
         ],
         outputs=[mdx23c_stem1, mdx23c_stem2],
     )
@@ -497,6 +537,13 @@ with gr.Blocks(
             norm_threshold,
             amp_threshold,
             batch_size,
+            vocals_stem,
+            instrumental_stem,
+            other_stem,
+            drums_stem,
+            bass_stem,
+            guitar_stem,
+            piano_stem,
         ],
         outputs=[mdx_stem1, mdx_stem2],
     )
@@ -517,6 +564,13 @@ with gr.Blocks(
             norm_threshold,
             amp_threshold,
             batch_size,
+            vocals_stem,
+            instrumental_stem,
+            other_stem,
+            drums_stem,
+            bass_stem,
+            guitar_stem,
+            piano_stem,
         ],
         outputs=[vr_stem1, vr_stem2],
     )
@@ -534,6 +588,13 @@ with gr.Blocks(
             output_format,
             norm_threshold,
             amp_threshold,
+            vocals_stem,
+            instrumental_stem,
+            other_stem,
+            drums_stem,
+            bass_stem,
+            guitar_stem,
+            piano_stem,
         ],
         outputs=[demucs_stem1, demucs_stem2, demucs_stem3, demucs_stem4, demucs_stem5, demucs_stem6],
     )