Add BS-RoFormer vocal separation mode

2026-09-08 19:23:29 +02:00
parent a4e894fe70
commit 0069b61dbb
15 changed files with 421 additions and 52 deletions
@@ -41,7 +41,40 @@ def music_item(state="exited"):
            "Labels": {controller.MUSIC_LABEL_KEY: "acestep"}}


+def separator_item(state="exited"):
+    return {"Id": "id-separator", "State": state,
+            "Labels": {controller.SEPARATOR_LABEL_KEY: "bs-roformer"}}
+
+
 class ProfileControllerTests(unittest.TestCase):
+    def test_separator_start_exclusively_stops_gpu_workers(self):
+        profiles = {name: item(name) for name in controller.ALLOWED}
+        profiles["large"] = item("large", "running")
+        calls = []
+
+        def request(method, path):
+            calls.append((method, path))
+            return 204, b""
+
+        with patch.object(controller, "SEPARATOR_WORKER", "bs-roformer"), \
+             patch.object(controller, "MUSIC_WORKER", "acestep"), \
+             patch.object(controller, "containers", return_value=profiles), \
+             patch.object(controller, "separator_container", return_value=separator_item()), \
+             patch.object(controller, "music_container", return_value=music_item("running")), \
+             patch.object(controller, "image_containers", return_value=[image_item("running")]), \
+             patch.object(controller, "tts_container", return_value=tts_item()), \
+             patch.object(controller, "docker_request", side_effect=request):
+            result = controller.set_separator_worker(True)
+
+        self.assertEqual(result, {"separator_worker": "bs-roformer", "state": "running"})
+        self.assertEqual(calls, [
+            ("POST", "/containers/id-large/stop?t=120"),
+            ("POST", "/containers/id-flux/stop?t=20"),
+            ("POST", "/containers/id-tts/stop?t=30"),
+            ("POST", "/containers/id-music/stop?t=30"),
+            ("POST", "/containers/id-separator/start"),
+        ])
+
    def test_music_start_exclusively_stops_gpu_workers(self):
        profiles = {name: item(name) for name in controller.ALLOWED}
        profiles["ultra"] = item("ultra", "running")