Skip to content

Commit

Permalink
Fixed GPU mapping for roformer models
Browse files Browse the repository at this point in the history
  • Loading branch information
beveradb committed May 25, 2024
1 parent a3b7294 commit a581da7
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 4 deletions.
6 changes: 3 additions & 3 deletions audio_separator/separator/architectures/mdxc_separator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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...")
Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down

0 comments on commit a581da7

Please sign in to comment.