test: pin RoFormer forward-pass budget + GPU per-pass timing guard (v0.47.1)

.github/workflows/run-integration-tests.yaml

@@ -80,7 +80,7 @@ jobs:
     needs: changes
     if: needs.changes.outputs.should_run == 'true'
     runs-on: [self-hosted, linux, gpu]
-    timeout-minutes: 15
+    timeout-minutes: 20
     env:
       AUDIO_SEPARATOR_MODEL_DIR: /opt/audio-separator-models
     steps:
@@ -103,7 +103,7 @@ jobs:
           cache: poetry
       - name: Install Poetry dependencies (GPU)
         run: poetry install -E gpu
-      - name: "Run: 24-bit, CLI, output, and roformer tests (~7 min)"
+      - name: "Run: 24-bit, CLI, output, roformer and timing tests (~8 min)"
         run: |
           poetry run pytest -sv \
             tests/integration/test_24bit_preservation.py \
@@ -115,7 +115,8 @@ jobs:
             tests/integration/test_roformer_e2e.py \
             tests/integration/test_roformer_fallback_mechanism.py \
             tests/integration/test_roformer_model_switching.py \
-            tests/integration/test_roformer_new_parameters.py
+            tests/integration/test_roformer_new_parameters.py \
+            tests/integration/test_roformer_inference_timing.py
       - name: Upload test artifacts
         if: always()
         uses: actions/upload-artifact@v4

README.md

@@ -752,6 +752,7 @@ You can also rename specific stems:
 - **`vr_params`:** (Optional) VR Architecture Specific Attributes & Defaults. `Default: {"batch_size": 1, "window_size": 512, "aggression": 5, "enable_tta": False, "enable_post_process": False, "post_process_threshold": 0.2, "high_end_process": False}`
 - **`demucs_params`:** (Optional) Demucs Architecture Specific Attributes & Defaults. `Default: {"segment_size": "Default", "shifts": 2, "overlap": 0.25, "segments_enabled": True}` _(Note: `segment_size` "Default" uses the model's internal default, typically 40 for older Demucs models and 10 for Demucs v4/htdemucs)_
 - **`mdxc_params`:** (Optional) MDXC Architecture Specific Attributes & Defaults. `Default: {"segment_size": 256, "override_model_segment_size": False, "batch_size": None, "overlap": None, "pitch_shift": 0}` (`None` uses the model YAML's `inference` value, falling back to `1` for batch size and `8` for overlap.)
+  - **Performance note (since 0.47.0):** for RoFormer models, `overlap` is the number of overlapping prediction windows (step = chunk size / overlap), taken from the model YAML's `num_overlap` by default. Earlier releases used a fixed ~8 s step, which meant almost no overlap. RoFormer models whose YAML sets `num_overlap: 4` or `8` (for example the MelBand karaoke models) therefore run 4–6× more forward passes per track than in 0.46.x and take proportionally longer. To trade a little seam quality for speed, pass an explicit value, e.g. `mdxc_params={"overlap": 2, ...}` (Python) or `--mdxc_overlap=2` (CLI). `tests/unit/test_mdxc_roformer_pass_budget.py` pins the per-model pass counts.
 - **`ensemble_algorithm`:** (Optional) Algorithm to use for ensembling multiple models. `Default: 'avg_wave'`
 - **`ensemble_weights`:** (Optional) Weights for each model in the ensemble. `Default: None` (equal weights)
 - **`ensemble_preset`:** (Optional) Named ensemble preset (e.g. `'vocal_balanced'`, `'karaoke'`). Sets models, algorithm, and weights automatically. Use `Separator(info_only=True).list_ensemble_presets()` to see all. `Default: None`

pyproject.toml

@@ -52,7 +52,7 @@ dml = ["onnxruntime-directml>=1.17", "torch_directml"]
 
 [tool.poetry]
 requires-poetry = ">=2.0.0"
-version = "0.47.0"
+version = "0.47.1"
 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/integration/test_roformer_inference_timing.py

@@ -0,0 +1,72 @@
+"""
+GPU timing guard for RoFormer inference.
+
+Separation time = forward passes x seconds per pass. tests/unit/test_mdxc_roformer_pass_budget.py
+pins the pass count on every PR; this test measures the other factor on a real CUDA GPU, so a
+change that makes each pass slower (precision, attention backend, device transfers, compile)
+fails here instead of surfacing as slow production jobs.
+
+The ceiling is deliberately generous (CI runs on a T4) and can be tuned with
+ROFORMER_MAX_SECONDS_PER_PASS. The measured numbers are printed so drift is visible in the log.
+"""
+
+import os
+import statistics
+import time
+
+import pytest
+import torch
+
+from audio_separator.separator import Separator
+
+MODEL = "mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt"
+INPUT = os.path.join(os.path.dirname(__file__), "..", "inputs", "mardy20s.flac")
+MAX_SECONDS_PER_PASS = float(os.environ.get("ROFORMER_MAX_SECONDS_PER_PASS", "2.0"))
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="RoFormer timing guard needs a CUDA GPU")
+
+
+def _time_separation(tmp_path, overlap):
+    separator = Separator(output_dir=str(tmp_path), mdxc_params={"segment_size": 256, "override_model_segment_size": False, "batch_size": None, "overlap": overlap, "pitch_shift": 0})
+    separator.load_model(model_filename=MODEL)
+    model_instance = separator.model_instance
+
+    pass_seconds = []
+    run_pass = model_instance._run_roformer_model
+
+    def timed_pass(part):
+        torch.cuda.synchronize()
+        started = time.perf_counter()
+        output = run_pass(part)
+        torch.cuda.synchronize()
+        pass_seconds.append(time.perf_counter() - started)
+        return output
+
+    model_instance._run_roformer_model = timed_pass
+    separator.separate(INPUT)
+    return model_instance.overlap, pass_seconds
+
+
+@pytest.mark.parametrize(
+    "overlap,expected_overlap,expected_passes",
+    [
+        (None, 4, 8),  # model YAML num_overlap
+        (2, 2, 5),  # what karaoke-gen runs
+    ],
+)
+def test_roformer_karaoke_seconds_per_pass(tmp_path, overlap, expected_overlap, expected_passes):
+    effective_overlap, pass_seconds = _time_separation(tmp_path, overlap)
+
+    assert effective_overlap == expected_overlap
+    assert len(pass_seconds) == expected_passes
+
+    # The first pass includes CUDA/cuDNN warm-up, so judge the steady state.
+    steady = statistics.median(pass_seconds[1:])
+    print(
+        f"\n{MODEL} on {torch.cuda.get_device_name(0)} overlap={effective_overlap}: "
+        f"{len(pass_seconds)} passes, first {pass_seconds[0]:.3f}s, median steady {steady:.3f}s, total {sum(pass_seconds):.2f}s"
+    )
+    assert steady <= MAX_SECONDS_PER_PASS, (
+        f"Median RoFormer forward pass took {steady:.3f}s (> {MAX_SECONDS_PER_PASS}s). "
+        "Per-pass inference got slower; check precision, attention backend and device transfers."
+    )

tests/unit/test_mdxc_roformer_pass_budget.py

@@ -0,0 +1,90 @@
+"""
+Forward-pass budget for RoFormer inference on real model geometries.
+
+The RoFormer cost of a separation is (number of model forward passes) x (cost per pass).
+v0.47.0 (#299) changed how `overlap` maps to the chunk step, which silently made the
+karaoke models (YAML num_overlap 4/8) run 4-6x more forward passes than v0.44.2 for the
+same track. Nothing failed; downstream users just saw separations take 3-6x longer.
+
+These tests pin the pass count for a 3-minute stereo track. If a change to chunking,
+step or tail handling moves any of these numbers, that is a deliberate compute change:
+update the budget AND call it out in the release notes.
+"""
+
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pytest
+import torch
+from ml_collections import ConfigDict
+
+from audio_separator.separator.architectures.mdxc_separator import MDXCSeparator
+
+SAMPLE_RATE = 44100
+TRACK_SECONDS = 180
+
+
+class CountingIdentityModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.anchor = torch.nn.Parameter(torch.zeros(1))
+        self.calls = 0
+
+    def forward(self, audio):
+        self.calls += 1
+        return audio.unsqueeze(1)
+
+
+def _count_forward_passes(dim_t: int, overlap: int) -> int:
+    separator = MDXCSeparator.__new__(MDXCSeparator)
+    separator.pitch_shift = 0
+    separator.is_roformer = True
+    separator.override_model_segment_size = False
+    separator.model_data_cfgdict = ConfigDict(
+        {
+            "inference": {"dim_t": dim_t},
+            "training": {"target_instrument": "Vocals", "instruments": ["Vocals"]},
+            "model": {"stft_hop_length": 441},
+            "audio": {"hop_length": 441, "sample_rate": SAMPLE_RATE},
+        }
+    )
+    separator.overlap = overlap
+    separator.model_run = CountingIdentityModel()
+    separator.is_primary_stem_main_target = False
+    separator.logger = Mock()
+
+    mix = np.zeros((2, SAMPLE_RATE * TRACK_SECONDS), dtype=np.float32)
+    with patch("audio_separator.separator.architectures.mdxc_separator.tqdm", side_effect=lambda values: values):
+        separator.demix(mix)
+    return separator.model_run.calls
+
+
+# (model, inference.dim_t, YAML num_overlap, passes at YAML num_overlap, passes at overlap=2)
+# v0.44.2 ran every one of these at 23 passes (fixed 8 s step, clamped to chunk size).
+MODEL_BUDGETS = [
+    ("mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956", 801, 4, 87, 44),
+    ("mel_band_roformer_karaoke_gabox_v2", 1101, 8, 124, 32),
+    ("mel_band_roformer_karaoke_becruily", 1101, 8, 124, 32),
+    ("mel_band_roformer_instrumental_fv7z_gabox", 1101, 2, 32, 32),
+    ("bs_roformer_instrumental_resurrection_unwa", 1700, 2, 21, 21),
+]
+
+
+@pytest.mark.parametrize("model,dim_t,num_overlap,yaml_passes,overlap2_passes", MODEL_BUDGETS)
+def test_roformer_forward_pass_budget(model, dim_t, num_overlap, yaml_passes, overlap2_passes):
+    assert _count_forward_passes(dim_t, num_overlap) == yaml_passes, (
+        f"{model}: forward passes at YAML num_overlap={num_overlap} changed. "
+        "This changes separation time proportionally; update the budget and the release notes."
+    )
+    assert _count_forward_passes(dim_t, 2) == overlap2_passes, (
+        f"{model}: forward passes at overlap=2 changed. "
+        "This changes separation time proportionally; update the budget and the release notes."
+    )
+
+
+def test_pass_count_scales_with_overlap():
+    """Doubling overlap roughly doubles compute: the step is chunk_size // overlap."""
+    passes = [_count_forward_passes(801, overlap) for overlap in (1, 2, 4, 8)]
+    assert passes == sorted(passes)
+    for lower, higher in zip(passes, passes[1:]):
+        assert 1.7 <= higher / lower <= 2.1