diff --git a/audio_separator/separator/roformer/roformer_loader.py b/audio_separator/separator/roformer/roformer_loader.py index 1df1c7b..95c62dd 100644 --- a/audio_separator/separator/roformer/roformer_loader.py +++ b/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: diff --git a/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py b/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py index 8d96671..0a9adad 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py +++ b/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) diff --git a/audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py b/audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py index 61660c5..a4ecda1 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py +++ b/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,6 +413,12 @@ 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) @@ -401,7 +426,7 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False): 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) diff --git a/pyproject.toml b/pyproject.toml index cde5b4f..401dda3 100644 --- a/pyproject.toml +++ b/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 "] license = "MIT" diff --git a/tests/unit/test_directml.py b/tests/unit/test_directml.py index 2bbf862..86b87b0 100644 --- a/tests/unit/test_directml.py +++ b/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=) 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" diff --git a/tests/unit/test_roformer_dml_forward.py b/tests/unit/test_roformer_dml_forward.py new file mode 100644 index 0000000..ba6d55c --- /dev/null +++ b/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"