Remove empty 4stems and 5stems directories before Spleeter separation

JeffreyCA · Feb 2, 2024 · 1860156 · 1860156
1 parent 80a912e
commit 1860156
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/api/separators/spleeter_separator.py b/api/separators/spleeter_separator.py
@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import numpy as np
 from spleeter import *
 from spleeter.audio.adapter import AudioAdapter
@@ -28,6 +30,15 @@ def __init__(self, cpu_separation: bool, output_format=OutputFormat.MP3_256.valu
                                    multiprocess=False)
         self.audio_adapter = AudioAdapter.default()
 
+    def check_and_remove_empty_model_dirs(self):
+        model_paths = [
+            Path('pretrained_models', '4stems'),
+            Path('pretrained_models', '5stems')
+        ]
+        for model_path in model_paths:
+            if model_path.exists() and not any(model_path.iterdir()):
+                model_path.rmdir()
+
     def create_static_mix(self, parts, input_path, output_path):
         """Creates a static mix by performing source separation and adding the
            parts to be kept into a single track.
@@ -39,6 +50,7 @@ def create_static_mix(self, parts, input_path, output_path):
         """
         waveform, _ = self.audio_adapter.load(input_path,
                                               sample_rate=self.sample_rate)
+        self.check_and_remove_empty_model_dirs()
         prediction = self.separator.separate(waveform, '')
         out = np.zeros_like(prediction['vocals'])
 
@@ -56,6 +68,7 @@ def separate_into_parts(self, input_path, output_path):
         :param input_path: Input path
         :param output_path: Output path
         """
+        self.check_and_remove_empty_model_dirs()
         self.separator.separate_to_file(input_path,
                                         output_path,
                                         self.audio_adapter,