diff --git a/.github/workflows/run-integration-tests.yaml b/.github/workflows/run-integration-tests.yaml index 6ff4c46..43fbe67 100644 --- a/.github/workflows/run-integration-tests.yaml +++ b/.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 diff --git a/README.md b/README.md index 8b6c73d..d7c031e 100644 --- a/README.md +++ b/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` diff --git a/pyproject.toml b/pyproject.toml index 3cee1c0..1f50d97 100644 --- a/pyproject.toml +++ b/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 "] license = "MIT" diff --git a/tests/integration/test_roformer_inference_timing.py b/tests/integration/test_roformer_inference_timing.py new file mode 100644 index 0000000..04e2efd --- /dev/null +++ b/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." + ) diff --git a/tests/unit/test_mdxc_roformer_pass_budget.py b/tests/unit/test_mdxc_roformer_pass_budget.py new file mode 100644 index 0000000..7b489d3 --- /dev/null +++ b/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