Update app.py

app.py

@@ -123,7 +123,7 @@ 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, progress=gr.Progress(track_tqdm=True)):
+def roformer_separator(audio, model_key, seg_size, override_seg_size, batch_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True, desc="Audio separated...")):
     """Separate audio using Roformer model."""
     base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model_key)
@@ -139,18 +139,15 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
             amplification_threshold=amp_thresh,
             use_autocast=use_autocast,
             mdxc_params={
-                "batch_size": 1,
                 "segment_size": seg_size,
                 "override_model_segment_size": override_seg_size,
+                "batch_size": batch_size,
                 "overlap": overlap,
                 "pitch_shift": pitch_shift,
             }
         )
 
-        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
-
-        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
@@ -159,7 +156,7 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
     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, progress=gr.Progress(track_tqdm=True)):
+def mdx23c_separator(audio, model, seg_size, override_seg_size, batch_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True, desc="Audio separated...")):
     """Separate audio using MDX23C model."""
     base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
@@ -174,18 +171,15 @@ def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_s
             amplification_threshold=amp_thresh,
             use_autocast=use_autocast,
             mdxc_params={
-                "batch_size": 1,
                 "segment_size": seg_size,
                 "override_model_segment_size": override_seg_size,
+                "batch_size": batch_size,
                 "overlap": overlap,
                 "pitch_shift": pitch_shift,
             }
         )
 
-        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
-
-        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
@@ -194,7 +188,7 @@ def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_s
     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, progress=gr.Progress(track_tqdm=True)):
+def mdx_separator(audio, model, hop_length, seg_size, overlap, batch_size, denoise, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True, desc="Audio separated...")):
     """Separate audio using MDX-NET model."""
     base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
@@ -209,18 +203,15 @@ def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_di
             amplification_threshold=amp_thresh,
             use_autocast=use_autocast,
             mdx_params={
-                "batch_size": 1,
                 "hop_length": hop_length,
                 "segment_size": seg_size,
                 "overlap": overlap,
+                "batch_size": batch_size,
                 "enable_denoise": denoise,
             }
         )
 
-        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
-
-        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
@@ -229,7 +220,7 @@ 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, progress=gr.Progress(track_tqdm=True)):
+def vr_separator(audio, model, batch_size, window_size, aggression, tta, post_process, post_process_threshold, high_end_process, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True, desc="Audio separated...")):
     """Separate audio using VR ARCH model."""
     base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
@@ -244,7 +235,7 @@ def vr_separator(audio, model, window_size, aggression, tta, post_process, post_
             amplification_threshold=amp_thresh,
             use_autocast=use_autocast,
             vr_params={
-                "batch_size": 1,
+                "batch_size": batch_size,
                 "window_size": window_size,
                 "aggression": aggression,
                 "enable_tta": tta,
@@ -254,10 +245,7 @@ def vr_separator(audio, model, window_size, aggression, tta, post_process, post_
             }
         )
 
-        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
-
-        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
@@ -266,7 +254,7 @@ 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, progress=gr.Progress(track_tqdm=True, desc="Audio separated...")):
     """Separate audio using Demucs model."""
     print_message(audio, model)
     try:
@@ -287,10 +275,7 @@ def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled,
             }
         )
 
-        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
-
-        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio)
         print(f"Separation complete!\nResults: {', '.join(separation)}")
 
@@ -323,13 +308,13 @@ with gr.Blocks(
     gr.Markdown("<h1> Audio-Separator by Politrees </h1>")
     with gr.Accordion("General settings", open=False):
         with gr.Group():
-            model_file_dir = gr.Textbox(value="/tmp/audio-separator-models/", label="Directory for storing model files", placeholder="/tmp/audio-separator-models/", interactive=False)
+            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/", interactive=False)
             with gr.Row():
-                output_dir = gr.Textbox(value="output", label="File output directory", placeholder="output", interactive=False)
-                output_format = gr.Dropdown(value="wav", choices=["wav", "flac", "mp3"], label="Output Format")
+                output_dir = gr.Textbox(value="output", label="File output directory", info="The directory where output files will be saved.", placeholder="output", interactive=False)
+                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", info="max peak amplitude to normalize input and output audio.")
-                amp_threshold = gr.Slider(minimum=0.1, maximum=1, step=0.1, value=0.6, label="Amplification", info="min peak amplitude to amplify input and output audio.")
+                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.Tab("Roformer"):
         with gr.Group():
@@ -338,7 +323,8 @@ with gr.Blocks(
             with gr.Row():
                 roformer_seg_size = gr.Slider(minimum=32, maximum=4000, step=32, value=256, label="Segment Size", info="Larger consumes more resources, but may give better results.")
                 roformer_override_seg_size = gr.Checkbox(value=False, label="Override segment size", info="Override model default segment size instead of using the model default value.")
-                roformer_overlap = gr.Slider(minimum=2, maximum=10, step=1, value=8, label="Overlap", info="Amount of overlap between prediction windows.")
+                roformer_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.", interactive=False)
+                roformer_overlap = gr.Slider(minimum=2, maximum=10, step=1, value=8, label="Overlap", info="Amount of overlap between prediction windows. Lower is better but slower.")
                 roformer_pitch_shift = gr.Slider(minimum=-12, maximum=12, step=1, value=0, label="Pitch shift", info="Shift audio pitch by a number of semitones while processing. may improve output for deep/high vocals.")
         with gr.Row():
             roformer_audio = gr.Audio(label="Input Audio", type="filepath")
@@ -355,7 +341,8 @@ with gr.Blocks(
             with gr.Row():
                 mdx23c_seg_size = gr.Slider(minimum=32, maximum=4000, step=32, value=256, label="Segment Size", info="Larger consumes more resources, but may give better results.")
                 mdx23c_override_seg_size = gr.Checkbox(value=False, label="Override segment size", info="Override model default segment size instead of using the model default value.")
-                mdx23c_overlap = gr.Slider(minimum=2, maximum=50, step=1, value=8, label="Overlap", info="Amount of overlap between prediction windows.")
+                mdx23c_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.", interactive=False)
+                mdx23c_overlap = gr.Slider(minimum=2, maximum=50, step=1, value=8, label="Overlap", info="Amount of overlap between prediction windows. Higher is better but slower.")
                 mdx23c_pitch_shift = gr.Slider(minimum=-12, maximum=12, step=1, value=0, label="Pitch shift", info="Shift audio pitch by a number of semitones while processing. may improve output for deep/high vocals.")
         with gr.Row():
             mdx23c_audio = gr.Audio(label="Input Audio", type="filepath")
@@ -370,10 +357,11 @@ with gr.Blocks(
             with gr.Row():
                 mdx_model = gr.Dropdown(label="Select the Model", choices=MDXNET_MODELS)
             with gr.Row():
-                mdx_hop_length = gr.Slider(minimum=32, maximum=2048, step=32, value=1024, label="Hop Length")
+                mdx_hop_length = gr.Slider(minimum=32, maximum=2048, step=32, value=1024, label="Hop Length", info"Usually called stride in neural networks; only change if you know what you're doing.")
                 mdx_seg_size = gr.Slider(minimum=32, maximum=4000, step=32, value=256, label="Segment Size", info="Larger consumes more resources, but may give better results.")
-                mdx_overlap = gr.Slider(minimum=0.001, maximum=0.999, step=0.001, value=0.25, label="Overlap")
-                mdx_denoise = gr.Checkbox(value=False, label="Denoise", info="Enable denoising during separation.")
+                mdx_overlap = gr.Slider(minimum=0.001, maximum=0.999, step=0.001, value=0.25, label="Overlap", info"Amount of overlap between prediction windows. Higher is better but slower.")
+                mdx_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.", interactive=False)
+                mdx_denoise = gr.Checkbox(value=False, label="Denoise", info="Enable denoising after separation.")
         with gr.Row():
             mdx_audio = gr.Audio(label="Input Audio", type="filepath")
         with gr.Row():
@@ -387,10 +375,11 @@ with gr.Blocks(
             with gr.Row():
                 vr_model = gr.Dropdown(label="Select the Model", choices=VR_ARCH_MODELS)
             with gr.Row():
-                vr_window_size = gr.Slider(minimum=320, maximum=1024, step=32, value=512, label="Window Size")
+                vr_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.", interactive=False)
+                vr_window_size = gr.Slider(minimum=320, maximum=1024, step=32, value=512, label="Window Size", info="Balance quality and speed. 1024 = fast but lower, 320 = slower but better quality.")
                 vr_aggression = gr.Slider(minimum=1, maximum=50, step=1, value=5, label="Agression", info="Intensity of primary stem extraction.")
                 vr_tta = gr.Checkbox(value=False, label="TTA", info="Enable Test-Time-Augmentation; slow but improves quality.")
-                vr_post_process = gr.Checkbox(value=False, label="Post Process", info="Enable post-processing.")
+                vr_post_process = gr.Checkbox(value=False, label="Post Process", info="Identify leftover artifacts within vocal output; may improve separation for some songs.")
                 vr_post_process_threshold = gr.Slider(minimum=0.1, maximum=0.3, step=0.1, value=0.2, label="Post Process Threshold", info="Threshold for post-processing.")
                 vr_high_end_process = gr.Checkbox(value=False, label="High End Process", info="Mirror the missing frequency range of the output.")
         with gr.Row():
@@ -406,10 +395,10 @@ with gr.Blocks(
             with gr.Row():
                 demucs_model = gr.Dropdown(label="Select the Model", choices=DEMUCS_MODELS)
             with gr.Row():
-                demucs_seg_size = gr.Slider(minimum=1, maximum=100, step=1, value=40, label="Segment Size")
+                demucs_seg_size = gr.Slider(minimum=1, maximum=100, step=1, value=40, label="Segment Size", info="Size of segments into which the audio is split. Higher = slower but better quality.")
                 demucs_shifts = gr.Slider(minimum=0, maximum=20, step=1, value=2, label="Shifts", info="Number of predictions with random shifts, higher = slower but better quality.")
-                demucs_overlap = gr.Slider(minimum=0.001, maximum=0.999, step=0.001, value=0.25, label="Overlap")
-                demucs_segments_enabled = gr.Checkbox(value=True, label="Segment-wise processing")
+                demucs_overlap = gr.Slider(minimum=0.001, maximum=0.999, step=0.001, value=0.25, label="Overlap", info="Overlap between prediction windows. Higher = slower but better quality.")
+                demucs_segments_enabled = gr.Checkbox(value=True, label="Segment-wise processing", info="Enable segment-wise processing.")
         with gr.Row():
             demucs_audio = gr.Audio(label="Input Audio", type="filepath")
         with gr.Row():
@@ -433,6 +422,7 @@ with gr.Blocks(
             roformer_model,
             roformer_seg_size,
             roformer_override_seg_size,
+            roformer_batch_size,
             roformer_overlap,
             roformer_pitch_shift,
             model_file_dir,
@@ -450,6 +440,7 @@ with gr.Blocks(
             mdx23c_model,
             mdx23c_seg_size,
             mdx23c_override_seg_size,
+            mdx23c_batch_size,
             mdx23c_overlap,
             mdx23c_pitch_shift,
             model_file_dir,
@@ -468,6 +459,7 @@ with gr.Blocks(
             mdx_hop_length,
             mdx_seg_size,
             mdx_overlap,
+            mdx_batch_size,
             mdx_denoise,
             model_file_dir,
             output_dir,
@@ -482,6 +474,7 @@ with gr.Blocks(
         inputs=[
             vr_audio,
             vr_model,
+            vr_batch_size,
             vr_window_size,
             vr_aggression,
             vr_tta,