Fixed GPU mapping for roformer models

audio_separator/separator/architectures/mdxc_separator.py

@@ -94,7 +94,7 @@ def load_model(self):
                 checkpoint = torch.load(self.model_path, map_location="cpu")
                 self.model_run = model if not isinstance(model, torch.nn.DataParallel) else model.module
                 self.model_run.load_state_dict(checkpoint)
-                self.model_run.to(self.torch_device_cpu).eval()
+                self.model_run.to(self.torch_device).eval()
 
             else:
                 self.logger.debug("Loading TFC_TDF_net model...")
@@ -183,8 +183,8 @@ def overlap_add(self, result, x, weights, start, length):
         """
         Adds the overlapping part of the result to the result tensor.
         """
-        if self.torch_device == "mps":
-            x = x.to(self.torch_device_cpu)
+        x = x.to(result.device)
+        weights = weights.to(result.device)
         result[..., start : start + length] += x[..., :length] * weights[:length]
         return result
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "audio-separator"
-version = "0.17.1"
+version = "0.17.2"
 description = "Easy to use audio stem separation, using various models from UVR trained primarily by @Anjok07"
 authors = ["Andrew Beveridge <[email protected]>"]
 license = "MIT"