fix: RoFormer on DirectML — CPU state-dict load + complex-op CPU hops (#292)

audio_separator/separator/roformer/roformer_loader.py

@@ -10,6 +10,18 @@
 logger = logging.getLogger(__name__)
 
 
+def _is_dml_device(device) -> bool:
+    """True if `device` refers to a torch-directml device.
+
+    torch-directml registers itself on torch's out-of-tree backend slot, so
+    its devices stringify as "privateuseone[:N]". Accepts either a string or
+    a torch.device. Only reachable in practice when the Separator enabled
+    DirectML explicitly (use_directml=True), so this cannot misfire for
+    cuda/mps/cpu users.
+    """
+    return str(device).startswith("privateuseone")
+
+
 class RoformerLoader:
     """Main Roformer model loader (new implementation only)."""
 
@@ -95,7 +107,15 @@ def _load_with_new_implementation(self,
                 raise ValueError(f"Unknown model type: {model_type}")
 
             if os.path.exists(model_path):
-                state_dict = torch.load(model_path, map_location=device)
+                # torch-directml's deserialization hook expects integer device
+                # ids and raises TypeError ("'>=' not supported between
+                # instances of 'torch.device' and 'int'") when torch.load maps
+                # storages straight onto a privateuseone device. Loading on CPU
+                # and moving the model afterwards is equivalent and works
+                # everywhere, but is gated to DML so all other devices keep
+                # their exact existing behavior. (Issue #292)
+                map_location = "cpu" if _is_dml_device(device) else device
+                state_dict = torch.load(model_path, map_location=map_location)
                 if isinstance(state_dict, dict) and 'state_dict' in state_dict:
                     model.load_state_dict(state_dict['state_dict'])
                 elif isinstance(state_dict, dict) and 'model' in state_dict:

audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py

@@ -18,6 +18,15 @@
 # helper functions
 
 
+
+def _is_dml_device(device) -> bool:
+    """torch-directml devices use torch's out-of-tree backend slot (privateuseone).
+
+    Module-level so tests can patch it to exercise the DML CPU-hop branches
+    on CPU-only machines.
+    """
+    return device.type == "privateuseone"
+
 def exists(val):
     return val is not None
 
@@ -430,6 +439,11 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
 
         original_device = raw_audio.device
         x_is_mps = True if original_device.type == "mps" else False
+        # torch-directml (privateuseone) has no complex tensor support, so all
+        # complex ops (stft, view_as_complex, complex multiply, istft) hop to
+        # CPU; the transformer stack — the heavy compute — stays on the DML
+        # device. Gated so cuda/mps/cpu behavior is unchanged. (Issue #292)
+        x_is_dml = _is_dml_device(original_device)
 
         # if x_is_mps:
         #     raw_audio = raw_audio.cpu()
@@ -450,8 +464,12 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
 
         stft_window = self.stft_window_fn().to(device)
 
-        stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
-        stft_repr = torch.view_as_real(stft_repr)
+        if x_is_dml:
+            stft_repr = torch.stft(raw_audio.cpu(), **self.stft_kwargs, window=stft_window.cpu(), return_complex=True)
+            stft_repr = torch.view_as_real(stft_repr).to(device)
+        else:
+            stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
+            stft_repr = torch.view_as_real(stft_repr)
 
         stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c")
         stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c")  # merge stereo / mono into the frequency, with frequency leading dimension, for band splitting
@@ -500,6 +518,10 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
 
         # complex number multiplication
 
+        if x_is_dml:
+            stft_repr = stft_repr.cpu()
+            mask = mask.cpu()
+
         stft_repr = torch.view_as_complex(stft_repr)
         mask = torch.view_as_complex(mask)
 
@@ -509,7 +531,7 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
 
         stft_repr = rearrange(stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels)
 
-        recon_audio = torch.istft(stft_repr.cpu() if x_is_mps else stft_repr, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=False).to(device)
+        recon_audio = torch.istft(stft_repr.cpu() if x_is_mps else stft_repr, **self.stft_kwargs, window=stft_window.cpu() if (x_is_mps or x_is_dml) else stft_window, return_complex=False).to(device)
 
         recon_audio = rearrange(recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=self.num_stems)
 

audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py

@@ -17,6 +17,15 @@
 from librosa import filters
 
 
+
+def _is_dml_device(device) -> bool:
+    """torch-directml devices use torch's out-of-tree backend slot (privateuseone).
+
+    Module-level so tests can patch it to exercise the DML CPU-hop branches
+    on CPU-only machines.
+    """
+    return device.type == "privateuseone"
+
 def exists(val):
     return val is not None
 
@@ -339,6 +348,12 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
 
         original_device = raw_audio.device
         x_is_mps = True if original_device.type == "mps" else False
+        # torch-directml (privateuseone) has no complex tensor support, so all
+        # complex ops (stft, view_as_complex, scatter over complex, complex
+        # multiply, istft) hop to CPU; the transformer stack — the heavy
+        # compute — stays on the DML device. Gated so cuda/mps/cpu behavior
+        # is unchanged. (Issue #292)
+        x_is_dml = _is_dml_device(original_device)
 
         if x_is_mps:
             raw_audio = raw_audio.cpu()
@@ -360,8 +375,12 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
 
         stft_window = self.stft_window_fn().to(device)
 
-        stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
-        stft_repr = torch.view_as_real(stft_repr)
+        if x_is_dml:
+            stft_repr = torch.stft(raw_audio.cpu(), **self.stft_kwargs, window=stft_window.cpu(), return_complex=True)
+            stft_repr = torch.view_as_real(stft_repr).to(device)
+        else:
+            stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
+            stft_repr = torch.view_as_real(stft_repr)
 
         stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c")
         stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c")  # merge stereo / mono into the frequency, with frequency leading dimension, for band splitting
@@ -394,14 +413,20 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
         if x_is_mps:
             masks = masks.cpu()
 
+        # Everything from view_as_complex through istft is complex-dtype work;
+        # DML tensors must hop to CPU for it.
+        if x_is_dml:
+            stft_repr = stft_repr.cpu()
+            masks = masks.cpu()
+
         stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c")
 
         stft_repr = torch.view_as_complex(stft_repr)
         masks = torch.view_as_complex(masks)
 
         masks = masks.type(stft_repr.dtype)
 
-        if x_is_mps:
+        if x_is_mps or x_is_dml:
             scatter_indices = repeat(self.freq_indices.cpu(), "f -> b n f t", b=batch, n=self.num_stems, t=stft_repr.shape[-1])
         else:
             scatter_indices = repeat(self.freq_indices, "f -> b n f t", b=batch, n=self.num_stems, t=stft_repr.shape[-1])
@@ -410,12 +435,15 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
         masks_summed = (
             torch.zeros_like(stft_repr_expanded_stems.cpu() if x_is_mps else stft_repr_expanded_stems)
             .scatter_add_(2, scatter_indices.cpu() if x_is_mps else scatter_indices, masks.cpu() if x_is_mps else masks)
-            .to(device)
         )
+        if not x_is_dml:
+            # complex tensors cannot live on a DML device; keep them on CPU
+            # until after the istft
+            masks_summed = masks_summed.to(device)
 
         denom = repeat(self.num_bands_per_freq, "f -> (f r) 1", r=channels)
 
-        if x_is_mps:
+        if x_is_mps or x_is_dml:
             denom = denom.cpu()
 
         masks_averaged = masks_summed / denom.clamp(min=1e-8)
@@ -424,7 +452,10 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False):
 
         stft_repr = rearrange(stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels)
 
-        recon_audio = torch.istft(stft_repr.cpu() if x_is_mps else stft_repr, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=False, length=istft_length)
+        recon_audio = torch.istft(stft_repr.cpu() if x_is_mps else stft_repr, **self.stft_kwargs, window=stft_window.cpu() if (x_is_mps or x_is_dml) else stft_window, return_complex=False, length=istft_length)
+
+        if x_is_dml:
+            recon_audio = recon_audio.to(original_device)
 
         recon_audio = rearrange(recon_audio, "(b n s) t -> b n s t", b=batch, s=self.audio_channels, n=self.num_stems)
 

pyproject.toml

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

tests/unit/test_directml.py

@@ -36,3 +36,77 @@ def test_directml_hint_absent_when_no_packages(caplog):
     with caplog.at_level(logging.INFO):
         _run_setup(use_directml=False, dml_installed=False)
     assert not any(HINT in r.message for r in caplog.records)
+
+
+# ---------------------------------------------------------------------------
+# RoformerLoader map_location handling on DirectML (issue #292)
+#
+# torch-directml's deserialization hook expects integer device ids, so
+# torch.load(map_location=<privateuseone device>) raises TypeError and the
+# loader silently fell back to the legacy implementation. The fix loads the
+# state dict on CPU (DML only) and moves the model to the device afterwards.
+# ---------------------------------------------------------------------------
+
+from audio_separator.separator.roformer.roformer_loader import RoformerLoader, _is_dml_device
+
+
+class TestIsDmlDevice:
+    def test_dml_device_strings(self):
+        assert _is_dml_device("privateuseone")
+        assert _is_dml_device("privateuseone:0")
+        assert _is_dml_device("privateuseone:1")
+
+    def test_non_dml_devices(self):
+        assert not _is_dml_device("cpu")
+        assert not _is_dml_device("cuda")
+        assert not _is_dml_device("cuda:0")
+        assert not _is_dml_device("mps")
+
+    def test_torch_device_objects(self):
+        import torch
+
+        assert not _is_dml_device(torch.device("cpu"))
+        assert _is_dml_device(torch.device("privateuseone", 0))
+
+
+def _load_via_new_implementation(device):
+    """Drive _load_with_new_implementation with mocked model + torch.load,
+    returning (map_location_used, device_model_moved_to)."""
+    loader = RoformerLoader()
+    model = MagicMock(name="model")
+    seen = {}
+
+    def fake_torch_load(path, map_location=None):
+        seen["map_location"] = map_location
+        return {}
+
+    with patch.object(loader, "_create_bs_roformer", return_value=model), \
+         patch("torch.load", side_effect=fake_torch_load), \
+         patch("os.path.exists", return_value=True):
+        result = loader._load_with_new_implementation(
+            model_path="/fake/model.ckpt",
+            config={"dim": 1, "depth": 1, "freqs_per_bands": (2,)},
+            model_type="bs_roformer",
+            device=device,
+        )
+
+    assert result.success
+    model.to.assert_called_once_with(device)
+    return seen["map_location"]
+
+
+def test_new_implementation_loads_on_cpu_for_dml_device():
+    # State dict must be mapped to CPU; the model still moves to the DML device.
+    assert _load_via_new_implementation("privateuseone:0") == "cpu"
+
+
+def test_new_implementation_map_location_unchanged_for_cpu():
+    assert _load_via_new_implementation("cpu") == "cpu"
+
+
+def test_new_implementation_map_location_unchanged_for_cuda():
+    assert _load_via_new_implementation("cuda:0") == "cuda:0"
+
+
+def test_new_implementation_map_location_unchanged_for_mps():
+    assert _load_via_new_implementation("mps") == "mps"

tests/unit/test_roformer_dml_forward.py

@@ -0,0 +1,90 @@
+"""DML CPU-hop equivalence tests for the Roformer forward passes (issue #292).
+
+torch-directml has no complex tensor support, so bs_roformer/mel_band_roformer
+route their complex ops (stft, view_as_complex, complex multiply, istft) to CPU
+when running on a privateuseone device. We can't create DML tensors in CI, but
+we CAN force the DML branches on CPU tensors (where every hop is a no-op) and
+assert the output is identical to the untouched path — proving the hop plumbing
+itself doesn't alter results, reorder dims, or drop tensors.
+"""
+
+import pytest
+import torch
+from unittest.mock import patch
+
+from audio_separator.separator.uvr_lib_v5.roformer import bs_roformer as bs_mod
+from audio_separator.separator.uvr_lib_v5.roformer import mel_band_roformer as mel_mod
+
+
+def _tiny_bs_roformer():
+    torch.manual_seed(0)
+    return bs_mod.BSRoformer(
+        dim=32,
+        depth=1,
+        stereo=False,
+        num_stems=1,
+        time_transformer_depth=1,
+        freq_transformer_depth=1,
+        freqs_per_bands=(129, 128),  # sums to 512 // 2 + 1
+        stft_n_fft=512,
+        stft_hop_length=128,
+        stft_win_length=512,
+    ).eval()
+
+
+def _tiny_mel_band_roformer():
+    torch.manual_seed(0)
+    return mel_mod.MelBandRoformer(
+        dim=32,
+        depth=1,
+        stereo=False,
+        num_stems=1,
+        time_transformer_depth=1,
+        freq_transformer_depth=1,
+        num_bands=8,
+        stft_n_fft=512,
+        stft_hop_length=128,
+        stft_win_length=512,
+    ).eval()
+
+
+class TestIsDmlDeviceHelper:
+    def test_cpu_is_not_dml(self):
+        assert not bs_mod._is_dml_device(torch.device("cpu"))
+        assert not mel_mod._is_dml_device(torch.device("cpu"))
+
+    def test_privateuseone_is_dml(self):
+        assert bs_mod._is_dml_device(torch.device("privateuseone", 0))
+        assert mel_mod._is_dml_device(torch.device("privateuseone", 0))
+
+
+class TestBSRoformerDmlBranchEquivalence:
+    def test_forced_dml_branch_matches_normal_cpu_output(self):
+        model = _tiny_bs_roformer()
+        torch.manual_seed(1)
+        audio = torch.randn(1, 8192)
+
+        with torch.no_grad():
+            normal = model(audio)
+            with patch.object(bs_mod, "_is_dml_device", return_value=True):
+                hopped = model(audio)
+
+        assert hopped.device == audio.device
+        assert hopped.shape == normal.shape
+        assert torch.allclose(normal, hopped, atol=1e-6), "DML CPU-hop branch changed the output"
+
+
+class TestMelBandRoformerDmlBranchEquivalence:
+    def test_forced_dml_branch_matches_normal_cpu_output(self):
+        model = _tiny_mel_band_roformer()
+        torch.manual_seed(1)
+        audio = torch.randn(1, 8192)
+
+        with torch.no_grad():
+            normal = model(audio)
+            with patch.object(mel_mod, "_is_dml_device", return_value=True):
+                hopped = model(audio)
+
+        assert hopped.device == audio.device
+        assert hopped.shape == normal.shape
+        assert torch.allclose(normal, hopped, atol=1e-6), "DML CPU-hop branch changed the output"