Create music separator class

modules/uvr/music_separator.py

@@ -13,35 +13,68 @@ class MusicSeparator:
                  output_dir: Optional[str] = None):
         self.model = None
         self.device = self.get_device()
+        self.available_devices = ["cpu", "cuda"]
         self.model_dir = model_dir
         self.output_dir = output_dir
+        self.sample_rate = 16000
+        self.available_models = ["UVR-MDX-NET-Inst_1", "UVR-MDX-NET-Inst_HQ_1"]
+        self.default_model = self.available_models[0]
+        self.current_model_size = self.default_model
+        self.model_config = {
+            "segment": 256,
+            "split": True
+        }
 
     def update_model(self,
                      model_name: str = "UVR-MDX-NET-Inst_1",
+                     device: Optional[str] = None,
                      segment_size: int = 256):
-        self.model = MDX(name="UVR-MDX-NET-Inst_1",
-                         other_metadata={"segment": segment_size, "split": True},
+        if device is None:
+            device = self.device
+
+        self.device = device
+        self.model_config = {
+            "segment": segment_size,
+            "split": True
+        }
+        self.model = MDX(name=model_name,
+                         other_metadata=self.model_config,
                          device=self.device,
                          logger=None,
-                         model_path="models\UVR\MDX_Net_Models\UVR-MDX-NET-Inst_HQ_1.onnx")
+                         model_dir=self.model_dir)
 
     def separate(self,
                  audio_file_path: str,
-                 sample_rate: int = 44100,
+                 model_name: str,
+                 device: Optional[str] = None,
+                 segment_size: int = 256,
                  save_file: bool = True):
-        if self.model is None:
-            self.model = self.update_model()
+        if device is None:
+            device = self.device
+
+        model_config = {
+            "segment": segment_size,
+            "split": True
+        }
+
+        if self.model is None or self.current_model_size != model_name or self.model_config != model_config:
+            self.update_model(
+                model_name=model_name,
+                device=device,
+                segment_size=segment_size
+            )
 
         result = self.model(audio_file_path)
         instrumental, vocals = result["instrumental"].T, result["vocals"].T
 
         if save_file:
-            filename = audio_file_path
+            filename, ext = os.path.splitext(audio_file_path)
+            filename = os.path.basename(filename) + ".wav"
             instrumental_output_path = os.path.join(self.output_dir, "instrumental", filename)
             vocals_output_path = os.path.join(self.output_dir, "vocals", filename)
 
-            sf.write(instrumental_output_path, instrumental, sample_rate, format="WAV")
-            sf.write(vocals_output_path, vocals, sample_rate, format="WAV")
+            sf.write(instrumental_output_path, instrumental, self.sample_rate, format="WAV")
+            sf.write(vocals_output_path, vocals, self.sample_rate, format="WAV")
 
         return instrumental, vocals