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, batch_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, 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."""
     base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model_key)
@@ -159,7 +159,7 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, batch_size
     except Exception as e:
         raise RuntimeError(f"Roformer separation failed: {e}") from e
 
-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)):
+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)):
     """Separate audio using MDX23C model."""
     base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
@@ -194,7 +194,7 @@ def mdx23c_separator(audio, model, seg_size, override_seg_size, batch_size, over
     except Exception as e:
         raise RuntimeError(f"MDX23C separation failed: {e}") from e
 
-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)):
+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)):
     """Separate audio using MDX-NET model."""
     base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
@@ -229,7 +229,7 @@ def mdx_separator(audio, model, hop_length, seg_size, overlap, batch_size, denoi
     except Exception as e:
         raise RuntimeError(f"MDX-NET separation failed: {e}") from e
 
-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)):
+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)):
     """Separate audio using VR ARCH model."""
     base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
@@ -330,6 +330,8 @@ with gr.Blocks(
             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.", interactive=False)
 
     with gr.Tab("Roformer"):
         with gr.Group():
@@ -338,7 +340,6 @@ 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_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():
@@ -356,7 +357,6 @@ 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_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():
@@ -375,7 +375,6 @@ with gr.Blocks(
                 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", 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")
@@ -390,7 +389,6 @@ with gr.Blocks(
             with gr.Row():
                 vr_model = gr.Dropdown(label="Select the Model", choices=VR_ARCH_MODELS)
             with gr.Row():
-                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.")
@@ -437,7 +435,6 @@ with gr.Blocks(
             roformer_model,
             roformer_seg_size,
             roformer_override_seg_size,
-            roformer_batch_size,
             roformer_overlap,
             roformer_pitch_shift,
             model_file_dir,
@@ -445,6 +442,7 @@ with gr.Blocks(
             output_format,
             norm_threshold,
             amp_threshold,
+            batch_size,
         ],
         outputs=[roformer_stem1, roformer_stem2],
     )
@@ -455,7 +453,6 @@ with gr.Blocks(
             mdx23c_model,
             mdx23c_seg_size,
             mdx23c_override_seg_size,
-            mdx23c_batch_size,
             mdx23c_overlap,
             mdx23c_pitch_shift,
             model_file_dir,
@@ -463,6 +460,7 @@ with gr.Blocks(
             output_format,
             norm_threshold,
             amp_threshold,
+            batch_size,
         ],
         outputs=[mdx23c_stem1, mdx23c_stem2],
     )
@@ -474,13 +472,13 @@ with gr.Blocks(
             mdx_hop_length,
             mdx_seg_size,
             mdx_overlap,
-            mdx_batch_size,
             mdx_denoise,
             model_file_dir,
             output_dir,
             output_format,
             norm_threshold,
             amp_threshold,
+            batch_size,
         ],
         outputs=[mdx_stem1, mdx_stem2],
     )
@@ -489,7 +487,6 @@ with gr.Blocks(
         inputs=[
             vr_audio,
             vr_model,
-            vr_batch_size,
             vr_window_size,
             vr_aggression,
             vr_tta,
@@ -501,6 +498,7 @@ with gr.Blocks(
             output_format,
             norm_threshold,
             amp_threshold,
+            batch_size,
         ],
         outputs=[vr_stem1, vr_stem2],
     )