Update app.py

app.py

@@ -111,21 +111,6 @@ def print_message(input_file, model_name):
     print("Separation Model:", model_name)
     print("Audio Separation Process...")
 
-def rename_stems(input_file, output_dir, stems, output_format):
-    """Rename stems to the format of the input file name with __(StemX) suffix."""
-    base_name = os.path.splitext(os.path.basename(input_file))[0]
-    renamed_stems = []
-    for i, stem in enumerate(stems):
-        new_name = f"{base_name}_(Stem{i+1}).{output_format}"
-        new_path = os.path.join(output_dir, new_name)
-        try:
-            print("Separation result:", new_name)
-            os.rename(os.path.join(output_dir, stem), new_path)
-            renamed_stems.append(new_path)
-        except Exception as e:
-            logging.error(f"Failed to rename stem {stem}: {e}")
-    return renamed_stems
-
 def prepare_output_dir(input_file, output_dir):
     """Create a directory for the output files and clean it if it already exists."""
     base_name = os.path.splitext(os.path.basename(input_file))[0]
@@ -135,12 +120,12 @@ def prepare_output_dir(input_file, output_dir):
             shutil.rmtree(out_dir)
         os.makedirs(out_dir)
     except Exception as e:
-        logging.error(f"Failed to prepare output directory {out_dir}: {e}")
-        raise
+        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)):
     """Separate audio using Roformer model."""
+    base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model_key)
     model = ROFORMER_MODELS[model_key]
     try:
@@ -156,7 +141,7 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
             mdxc_params={
                 "batch_size": 1,
                 "segment_size": seg_size,
-                "override_model_segment_size": False,
+                "override_model_segment_size": override_seg_size,
                 "overlap": overlap,
                 "pitch_shift": pitch_shift,
             }
@@ -166,18 +151,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)
-
-        progress(0.9, desc="Stems renamed...")
-        stems = rename_stems(audio, out_dir, separation, out_format)
+        separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
+        print(f"Separation complete!\nResults: {', '.join(separation)}")
 
+        stems = [os.path.join(out_dir, file_name) for file_name in separation]
         return stems[0], stems[1]
     except Exception as e:
-        logging.error(f"Roformer separation failed: {e}")
-        return None, None
+        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)):
     """Separate audio using MDX23C model."""
+    base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
     try:
         out_dir = prepare_output_dir(audio, out_dir)
@@ -202,18 +186,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)
-
-        progress(0.9, desc="Stems renamed...")
-        stems = rename_stems(audio, out_dir, separation, out_format)
+        separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
+        print(f"Separation complete!\nResults: {', '.join(separation)}")
 
+        stems = [os.path.join(out_dir, file_name) for file_name in separation]
         return stems[0], stems[1]
     except Exception as e:
-        logging.error(f"MDX23C separation failed: {e}")
-        return None, None
+        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)):
     """Separate audio using MDX-NET model."""
+    base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
     try:
         out_dir = prepare_output_dir(audio, out_dir)
@@ -238,18 +221,17 @@ 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)
-
-        progress(0.9, desc="Stems renamed...")
-        stems = rename_stems(audio, out_dir, separation, out_format)
+        separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
+        print(f"Separation complete!\nResults: {', '.join(separation)}")
 
+        stems = [os.path.join(out_dir, file_name) for file_name in separation]
         return stems[0], stems[1]
     except Exception as e:
-        logging.error(f"MDX-NET separation failed: {e}")
-        return None, None
+        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)):
     """Separate audio using VR ARCH model."""
+    base_name = os.path.splitext(os.path.basename(audio))[0]
     print_message(audio, model)
     try:
         out_dir = prepare_output_dir(audio, out_dir)
@@ -276,15 +258,13 @@ 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)
-
-        progress(0.9, desc="Stems renamed...")
-        stems = rename_stems(audio, out_dir, separation, out_format)
+        separation = separator.separate(audio, f"{base_name}_(Stem1)", f"{base_name}_(Stem2)")
+        print(f"Separation complete!\nResults: {', '.join(separation)}")
 
+        stems = [os.path.join(out_dir, file_name) for file_name in separation]
         return stems[0], stems[1]
     except Exception as e:
-        logging.error(f"VR ARCH separation failed: {e}")
-        return None, None
+        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)):
     """Separate audio using Demucs model."""
@@ -312,14 +292,22 @@ def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled,
 
         progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio)
-
-        progress(0.9, desc="Stems renamed...")
-        stems = rename_stems(audio, out_dir, separation, out_format)
-
-        return stems[0], stems[1], stems[2], stems[3]
+        print(f"Separation complete!\nResults: {', '.join(separation)}")
+
+        stems = [os.path.join(out_dir, file_name) for file_name in separation]
+        
+        if model == "htdemucs_6s.yaml":
+            return stems[0], stems[1], stems[2], stems[3], stems[4], stems[5]
+        else:
+            return stems[0], stems[1], stems[2], stems[3], None, None
     except Exception as e:
-        logging.error(f"Demucs separation failed: {e}")
-        return None, None, None, None
+        raise RuntimeError(f"Demucs separation failed: {e}") from e
+
+def update_stems(model):
+    if model == "htdemucs_6s.yaml":
+        return gr.update(visible=True)
+    else:
+        return gr.update(visible=False)
 
 with gr.Blocks(
     title="🎵 PolUVR - Politrees 🎵",
@@ -432,6 +420,11 @@ with gr.Blocks(
         with gr.Row():
             demucs_stem3 = gr.Audio(label="Stem 3", type="filepath", interactive=False)
             demucs_stem4 = gr.Audio(label="Stem 4", type="filepath", interactive=False)
+        with gr.Row(visible=False) as stem6:
+            demucs_stem5 = gr.Audio(label="Stem 5", type="filepath", interactive=False)
+            demucs_stem6 = gr.Audio(label="Stem 6", type="filepath", interactive=False)
+
+    demucs_model.change(update_stems, inputs=[demucs_model], outputs=stem6)
 
     roformer_button.click(
         roformer_separator,
@@ -518,7 +511,7 @@ with gr.Blocks(
             norm_threshold,
             amp_threshold,
         ],
-        outputs=[demucs_stem1, demucs_stem2, demucs_stem3, demucs_stem4],
+        outputs=[demucs_stem1, demucs_stem2, demucs_stem3, demucs_stem4, demucs_stem5, demucs_stem6],
     )
 
 app.launch(share=True)