feat(bs_roformer): 支持零直流分量处理并优化频率点计算

  -【优化】在 roformer_stft_freq_bins 中通过 torch.stft 计算频率点数，以保持模型初始化时的随机种子兼容性
  -【功能】在 MelBandRoformer 中引入 zero_dc 参数，支持强制置零直流分量
  -【功能】在 istft 过程中实现 PyTorch 和 MLX 的直流置零逻辑
  -【测试】新增针对 zero_dc 功能的单元测试用例

Author: baicai-1145

pymss_core/modules/bs_roformer/common.py

@@ -175,7 +175,15 @@ def init_roformer_stft(module, stft_n_fft, stft_hop_length, stft_win_length, stf
 
 
 def roformer_stft_freq_bins(module, window_length):
-    return int(module.stft_kwargs["n_fft"]) // 2 + 1
+    # The original training code computed this shape through torch.stft on a
+    # random probe tensor during model construction. Preserve that RNG-consuming
+    # behavior so scratch initialization remains seed-compatible.
+    return torch.stft(
+        torch.randn(1, 4096),
+        **module.stft_kwargs,
+        window=torch.ones(window_length),
+        return_complex=True,
+    ).shape[1]
 
 
 def roformer_freqs_per_bands_with_complex(module, freqs_per_bands, freqs):
@@ -354,6 +362,8 @@ def istft_roformer(module, stft_repr, context, length):
         .permute(0, 1, 3, 2, 4)
         .reshape(b * n * context.channels, context.freq_bins, t)
     )
+    if getattr(module, "zero_dc", False):
+        stft_repr = stft_repr.index_fill(1, torch.tensor(0, device=stft_repr.device), 0.0)
 
     try:
         recon_audio = torch.istft(

pymss_core/modules/bs_roformer/mel_band_roformer.py

@@ -42,6 +42,7 @@ def __init__(
         stft_win_length=2048,
         stft_normalized=False,
         stft_window_fn: Optional[Callable] = None,
+        zero_dc=True,
         mask_estimator_depth=1,
         match_input_audio_length=False,
         mlp_expansion_factor=4,
@@ -116,6 +117,7 @@ def __init__(
             mask_estimator_kwargs={"mlp_hidden_layers": mlp_hidden_layers},
         )
 
+        self.zero_dc = zero_dc
         self.match_input_audio_length = match_input_audio_length
 
     def _forward_mask_core(self, selected_stft_repr):

pymss_core/modules/bs_roformer/mlx_roformer.py

@@ -116,6 +116,8 @@ def _istft_roformer(module, stft_repr, context, length):
     stft_repr = stft_repr.reshape(b, n, freq_bins, channels, t, 2)
     stft_repr = mx.transpose(stft_repr, (0, 1, 3, 2, 4, 5)).reshape(b * n * channels, freq_bins, t, 2)
     complex_stft = stft_repr[..., 0] + (1j * stft_repr[..., 1])
+    if getattr(module, "zero_dc", False):
+        complex_stft = complex_stft.at[:, 0, :].set(0)
     complex_stft = mx.moveaxis(complex_stft, -2, -1)
     if normalized:
         complex_stft = complex_stft * np.sqrt(n_fft)

tests/test_core_api.py

@@ -5,6 +5,7 @@
 import pymss_core
 from pymss_core import AttrDict, load_config, unwrap_state_dict
 from pymss_core.modules._dsp import mel_filterbank
+from pymss_core.modules.bs_roformer.common import SpectralContext, istft_roformer
 
 
 def test_public_api_exports_core_functions():
@@ -49,6 +50,36 @@ def test_model_internal_dsp_helpers_do_not_require_librosa():
     assert filters.shape == (60, 1025)
 
 
+def test_roformer_istft_zero_dc(monkeypatch):
+    class DummyModule:
+        stft_kwargs = {}
+        zero_dc = True
+
+    captured = {}
+
+    def fake_istft(stft_repr, **kwargs):
+        captured["stft_repr"] = stft_repr.detach().clone()
+        return torch.zeros(stft_repr.shape[0], kwargs["length"])
+
+    monkeypatch.setattr(torch, "istft", fake_istft)
+
+    context = SpectralContext(
+        batch=1,
+        channels=1,
+        freq_bins=3,
+        audio_length=4,
+        stft_window=torch.ones(4),
+        x_is_mps=False,
+    )
+    stft_repr = torch.ones(1, 1, 3, 2, dtype=torch.complex64)
+
+    output = istft_roformer(DummyModule(), stft_repr, context, length=4)
+
+    assert output.shape == (1, 1, 4)
+    assert torch.equal(captured["stft_repr"][:, 0], torch.zeros_like(captured["stft_repr"][:, 0]))
+    assert torch.equal(captured["stft_repr"][:, 1:], torch.ones_like(captured["stft_repr"][:, 1:]))
+
+
 def test_vr_network_structures_remain_importable():
     from pymss_core.modules.vocal_remover import CascadedASPPNet, CascadedNet, ModelParameters