From 174e5f0c4fdcede4122c2d960847b37ea661de97 Mon Sep 17 00:00:00 2001 From: Blaise Date: Mon, 19 Oct 2026 09:40:20 +0200 Subject: [PATCH 1/3] perf(mdx): batch inference windows and reduce device transfers --- README.md | 2 +- .../separator/architectures/mdx_separator.py | 321 +++++++++--------- audio_separator/separator/uvr_lib_v5/stft.py | 20 +- tests/unit/test_mdx_inference.py | 287 ++++++++++++++++ tests/unit/test_stft.py | 43 +++ 5 files changed, 505 insertions(+), 168 deletions(-) create mode 100644 tests/unit/test_mdx_inference.py diff --git a/README.md b/README.md index 8b6c73d3..4eef6e7b 100644 --- a/README.md +++ b/README.md @@ -748,7 +748,7 @@ You can also rename specific stems: - **`use_native_fp16`:** (Optional) Convert a verified model to native float16 inference. Currently supported for MelBand RoFormer and BS-RoFormer on MPS and CUDA. Mutually exclusive with `use_autocast=True`; unsupported combinations warn and continue in float32. `Default: False` - **`use_torch_compile`:** (Optional) Compile verified repeated model blocks. This can be combined with float32 or autocast for MelBand RoFormer and BS-RoFormer on CPU, MPS, and CUDA, and with native float16 on MPS and CUDA. A fresh compiler cache can make the first run slower; unsupported combinations warn and continue in eager mode. `Default: False` - **`use_directml`:** (Optional) Flag to use DirectML for hardware-accelerated inference on Windows AMD/Intel GPUs (experimental; requires the `dml` extra and only takes effect when CUDA and Apple Silicon MPS are unavailable). `Default: False` -- **`mdx_params`:** (Optional) MDX Architecture Specific Attributes & Defaults. `Default: {"hop_length": 1024, "segment_size": 256, "overlap": 0.25, "batch_size": 1, "enable_denoise": False}` +- **`mdx_params`:** (Optional) MDX Architecture Specific Attributes & Defaults. `Default: {"hop_length": 1024, "segment_size": 256, "overlap": 0.25, "batch_size": 1, "enable_denoise": False}` `batch_size` groups independent windows for converted Torch models and dynamic-batch ONNX graphs; fixed-batch ONNX graphs retain their declared batch size. Larger batches can improve throughput at the cost of memory; overlap and denoise settings remain independent. - **`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.) diff --git a/audio_separator/separator/architectures/mdx_separator.py b/audio_separator/separator/architectures/mdx_separator.py index af5e311e..5dd87ee1 100644 --- a/audio_separator/separator/architectures/mdx_separator.py +++ b/audio_separator/separator/architectures/mdx_separator.py @@ -3,10 +3,8 @@ import os import platform import torch -import onnx import onnxruntime as ort import numpy as np -import onnx2torch from tqdm import tqdm from audio_separator.separator.uvr_lib_v5 import spec_utils from audio_separator.separator.uvr_lib_v5.stft import STFT @@ -120,10 +118,37 @@ def load_model(self): else: ort_session_options.log_severity_level = 0 - ort_inference_session = ort.InferenceSession(self.model_path, providers=self.onnx_execution_provider, sess_options=ort_session_options) + execution_providers = list(self.onnx_execution_provider) + provider_names = [provider if isinstance(provider, str) else provider[0] for provider in execution_providers] + self.onnx_stream = None + if ( + self.torch_device.type == "cuda" + and not torch.version.hip + and "CUDAExecutionProvider" in provider_names + and "CUDAExecutionProvider" in ort.get_available_providers() + ): + # A null/default stream pointer is treated by ORT as no user stream. + # Retain an explicit stream and order Torch producers/consumers around it. + self.onnx_stream = torch.cuda.Stream(device=self.torch_device) + execution_providers = [ + ( + ( + "CUDAExecutionProvider", + { + **(provider[1] if not isinstance(provider, str) else {}), + "device_id": self.onnx_stream.device.index, + "user_compute_stream": str(self.onnx_stream.cuda_stream), + }, + ) + if name == "CUDAExecutionProvider" + else provider + ) + for name, provider in zip(provider_names, execution_providers) + ] + ort_inference_session = ort.InferenceSession(self.model_path, providers=execution_providers, sess_options=ort_session_options) session_providers = ort_inference_session.get_providers() - requested_provider = self.onnx_execution_provider[0] if self.onnx_execution_provider else None + requested_provider = provider_names[0] if provider_names else None if requested_provider and requested_provider not in session_providers: self.logger.warning( f"ONNX Runtime could not activate requested provider {requested_provider}; " @@ -133,18 +158,84 @@ def load_model(self): else: self.logger.debug(f"ONNX Runtime session providers: {session_providers}") - self.model_run = lambda spek: ort_inference_session.run(None, {"input": spek.cpu().numpy()})[0] + self.onnx_session = ort_inference_session + self.onnx_input = ort_inference_session.get_inputs()[0] + self.onnx_output = ort_inference_session.get_outputs()[0] + if "CUDAExecutionProvider" not in session_providers: + self.onnx_stream = None + batch_dimension = self.onnx_input.shape[0] + self.fixed_batch_size = batch_dimension if isinstance(batch_dimension, int) and batch_dimension > 0 else None + self.inference_batch_size = self.fixed_batch_size or self.batch_size + if self.fixed_batch_size and self.fixed_batch_size != self.batch_size: + self.logger.info(f"ONNX model uses a fixed batch size of {self.fixed_batch_size}.") self.logger.debug("Model loaded successfully using ONNXruntime inferencing session.") else: - if platform.system() == 'Windows': + import onnx + import onnx2torch + + self.uses_pytorch_inference = True + self.fixed_batch_size = None + self.inference_batch_size = self.batch_size + if platform.system() == "Windows": onnx_model = onnx.load(self.model_path) self.model_run = onnx2torch.convert(onnx_model) else: self.model_run = onnx2torch.convert(self.model_path) - + self.model_run.to(self.torch_device).eval() self.logger.warning("Model converted from onnx to pytorch due to segment size not matching dim_t, processing may be slower.") + def model_run(self, spectrum): + """Run native ONNX inference; converted models replace this with a Torch module.""" + # Only bind shapes established by model metadata. Custom graphs with + # unknown output dimensions retain ONNX Runtime's allocating CPU path. + dimensions = {name: size for name, size in zip(self.onnx_input.shape, spectrum.shape) if isinstance(name, str)} + output_shape = [ + dimension if isinstance(dimension, int) else dimensions.get(dimension) for dimension in (self.onnx_output.shape or []) + ] + if ( + self.onnx_stream is None + or spectrum.device != self.onnx_stream.device + or spectrum.dtype != torch.float32 + or self.onnx_input.type != "tensor(float)" + or self.onnx_output.type != "tensor(float)" + or not output_shape + or not all(isinstance(size, int) and size > 0 for size in output_shape) + ): + return self.onnx_session.run( + [self.onnx_output.name], + {self.onnx_input.name: spectrum.cpu().numpy()}, + )[0] + + caller_stream = torch.cuda.current_stream(spectrum.device) + self.onnx_stream.wait_stream(caller_stream) + with torch.cuda.stream(self.onnx_stream): + contiguous = spectrum.contiguous() + output = torch.empty(output_shape, dtype=spectrum.dtype, device=spectrum.device) + binding = self.onnx_session.io_binding() + binding.bind_input( + self.onnx_input.name, + "cuda", + spectrum.device.index, + np.float32, + tuple(contiguous.shape), + contiguous.data_ptr(), + ) + binding.bind_output( + self.onnx_output.name, + "cuda", + spectrum.device.index, + np.float32, + tuple(output.shape), + output.data_ptr(), + ) + self.onnx_session.run_with_iobinding(binding) + # Retain allocator ownership until ORT's stream has consumed the input. + spectrum.record_stream(self.onnx_stream) + caller_stream.wait_stream(self.onnx_stream) + output.record_stream(caller_stream) + return output + def separate(self, audio_file_path, custom_output_names=None): """ Separates the audio file into primary and secondary sources based on the model's configuration. @@ -157,6 +248,8 @@ def separate(self, audio_file_path, custom_output_names=None): Returns: list: A list of paths to the output files generated by the separation process. """ + self.primary_source = None + self.secondary_source = None self.audio_file_path = audio_file_path self.audio_file_base = os.path.splitext(os.path.basename(audio_file_path))[0] @@ -172,7 +265,6 @@ def separate(self, audio_file_path, custom_output_names=None): source = self.demix(mix) * peak self.logger.debug("Demixing completed.") - if not isinstance(self.primary_source, np.ndarray): self.primary_source = source.T @@ -183,11 +275,11 @@ def separate(self, audio_file_path, custom_output_names=None): self.logger.debug("Processing output files...") # Process the secondary source if not already an array - if not isinstance(self.secondary_source, np.ndarray): - self.logger.debug("Producing secondary source: demixing in match_mix mode") - raw_mix = self.demix(mix, is_match_mix=True) - + needs_secondary = not self.output_single_stem or self.output_single_stem.lower() == self.secondary_stem_name.lower() + if needs_secondary and not isinstance(self.secondary_source, np.ndarray): if self.invert_using_spec: + self.logger.debug("Producing secondary source: demixing in match_mix mode") + raw_mix = self.demix(mix, is_match_mix=True) self.logger.debug("Inverting secondary stem using spectogram as invert_using_spec is set to True") self.secondary_source = spec_utils.invert_stem(raw_mix, self.primary_source * self.compensate) else: @@ -223,6 +315,11 @@ def initialize_model_settings(self): """ self.logger.debug("Initializing model settings...") + if not 0 <= self.overlap < 1: + raise ValueError("MDX overlap must be at least 0 and less than 1") + if self.batch_size < 1: + raise ValueError("MDX batch size must be greater than zero") + # n_bins is half the FFT size plus one (self.n_fft // 2 + 1). self.n_bins = self.n_fft // 2 + 1 @@ -234,6 +331,8 @@ def initialize_model_settings(self): # gen_size is the chunk size minus twice the trim size self.gen_size = self.chunk_size - 2 * self.trim + if self.gen_size <= 0: + raise ValueError("MDX segment size and hop length must exceed the FFT window") self.stft = STFT(self.logger, self.n_fft, self.hop_length, self.dim_f, self.torch_device) @@ -267,10 +366,10 @@ def initialize_mix(self, mix, is_ckpt=False): ( np.zeros((2, self.trim), dtype="float32"), # Pad at the start mix, - np.zeros((2, pad), dtype="float32"), # Pad in the middle (to match chunk size) + np.zeros((2, pad), dtype="float32"), # Pad in the middle (to match chunk size) np.zeros((2, self.trim), dtype="float32"), # Pad at the end ), - 1 + 1, ) # Determine the number of chunks based on the mixture's length num_chunks = mixture.shape[-1] // self.gen_size @@ -304,160 +403,60 @@ def initialize_mix(self, mix, is_ckpt=False): return mix_waves_tensor, pad def demix(self, mix, is_match_mix=False): - """ - Demixes the input mix into its constituent sources. If is_match_mix is True, the function adjusts the processing - to better match the mix, affecting chunk sizes and overlaps. The demixing process involves padding the mix, - processing it in chunks, applying windowing for overlaps, and accumulating the results to separate the sources. - """ - self.logger.debug(f"Starting demixing process with is_match_mix: {is_match_mix}...") + """Batch independent windows and normalize their overlap coverage.""" self.initialize_model_settings() - - # Preserves the original mix for later use. - # In UVR, this is used for the pitch fix and VR denoise processes, which aren't yet implemented here. - org_mix = mix - self.logger.debug(f"Original mix stored. Shape: {org_mix.shape}") - - # Initializes a list to store the separated waveforms. - tar_waves_ = [] - - # Handling different chunk sizes and overlaps based on the matching requirement. - if is_match_mix: - # Sets a smaller chunk size specifically for matching the mix. - chunk_size = self.hop_length * (self.segment_size - 1) - # Sets a small overlap for the chunks. - overlap = 0.02 - self.logger.debug(f"Chunk size for matching mix: {chunk_size}, Overlap: {overlap}") - else: - # Uses the regular chunk size defined in model settings. - chunk_size = self.chunk_size - # Uses the overlap specified in the model settings. - overlap = self.overlap - self.logger.debug(f"Standard chunk size: {chunk_size}, Overlap: {overlap}") - - # Calculates the generated size after subtracting the trim from both ends of the chunk. - gen_size = chunk_size - 2 * self.trim - self.logger.debug(f"Generated size calculated: {gen_size}") - - # Calculates padding to make the mix length a multiple of the generated size. - pad = gen_size + self.trim - ((mix.shape[-1]) % gen_size) - # Prepares the mixture with padding at the beginning and the end. - mixture = np.concatenate((np.zeros((2, self.trim), dtype="float32"), mix, np.zeros((2, pad), dtype="float32")), 1) - self.logger.debug(f"Mixture prepared with padding. Mixture shape: {mixture.shape}") - - # Calculates the step size for processing chunks based on the overlap. - step = int((1 - overlap) * chunk_size) - self.logger.debug(f"Step size for processing chunks: {step} as overlap is set to {overlap}.") - - # Initializes arrays to store the results and to account for overlap. - result = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32) - divider = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32) - - # Initializes counters for processing chunks. - total = 0 - total_chunks = (mixture.shape[-1] + step - 1) // step - self.logger.debug(f"Total chunks to process: {total_chunks}") - - # Processes each chunk of the mixture. - for i in tqdm(range(0, mixture.shape[-1], step)): - total += 1 - start = i - end = min(i + chunk_size, mixture.shape[-1]) - self.logger.debug(f"Processing chunk {total}/{total_chunks}: Start {start}, End {end}") - - # Handles windowing for overlapping chunks. - chunk_size_actual = end - start - window = None - if overlap != 0: - window = np.hanning(chunk_size_actual) - window = np.tile(window[None, None, :], (1, 2, 1)) - self.logger.debug("Window applied to the chunk.") - - # Zero-pad the chunk to prepare it for processing. - mix_part_ = mixture[:, start:end] - if end != i + chunk_size: - pad_size = (i + chunk_size) - end - mix_part_ = np.concatenate((mix_part_, np.zeros((2, pad_size), dtype="float32")), axis=-1) - - # Converts the chunk to a tensor for processing. - mix_part = torch.tensor([mix_part_], dtype=torch.float32).to(self.torch_device) - # Splits the chunk into smaller batches if necessary. - mix_waves = mix_part.split(self.batch_size) - total_batches = len(mix_waves) - self.logger.debug(f"Mix part split into batches. Number of batches: {total_batches}") - - with torch.no_grad(): - # Processes each batch in the chunk. - batches_processed = 0 - for mix_wave in mix_waves: - batches_processed += 1 - self.logger.debug(f"Processing mix_wave batch {batches_processed}/{total_batches}") - - # Runs the model to separate the sources. - tar_waves = self.run_model(mix_wave, is_match_mix=is_match_mix) - - # Applies windowing if needed and accumulates the results. - if window is not None: - tar_waves[..., :chunk_size_actual] *= window - divider[..., start:end] += window + chunk_size = self.chunk_size + overlap = 0.02 if is_match_mix else self.overlap + pad = self.gen_size + self.trim - (mix.shape[-1] % self.gen_size) + mixture = np.pad(mix, ((0, 0), (self.trim, pad))) + step = max(1, int((1 - overlap) * chunk_size)) + result = np.zeros_like(mixture, dtype=np.float32) + # Every channel uses the same window; keep a single coverage buffer. + divider = np.zeros(mixture.shape[-1], dtype=np.float32) + batch_size = self.inference_batch_size + starts = range(0, mixture.shape[-1], step) + full_window = np.hanning(chunk_size).astype(np.float32) if overlap else None + + with torch.inference_mode(), tqdm(total=len(starts), desc="Separating") as progress: + for offset in range(0, len(starts), batch_size): + locations = starts[offset : offset + batch_size] + actual_batch = self.fixed_batch_size or len(locations) + batch = np.zeros((actual_batch, 2, chunk_size), dtype=np.float32) + for index, start in enumerate(locations): + length = min(chunk_size, mixture.shape[-1] - start) + batch[index, :, :length] = mixture[:, start : start + length] + predictions = self.run_model( + torch.from_numpy(batch).to(self.torch_device), + is_match_mix=is_match_mix, + ) + for index, start in enumerate(locations): + length = min(chunk_size, mixture.shape[-1] - start) + window = full_window + if overlap and length != chunk_size: + window = np.hanning(length).astype(np.float32) + if window is None: + result[:, start : start + length] += predictions[index, :, :length] + divider[start : start + length] += 1 else: - divider[..., start:end] += 1 - - result[..., start:end] += tar_waves[..., : end - start] - - # Normalizes the results by the divider to account for overlap. - self.logger.debug("Normalizing result by dividing result by divider.") - tar_waves = result / divider - tar_waves_.append(tar_waves) - - # Reshapes the results to match the original dimensions. - tar_waves_ = np.vstack(tar_waves_)[:, :, self.trim : -self.trim] - tar_waves = np.concatenate(tar_waves_, axis=-1)[:, : mix.shape[-1]] - - # Extracts the source from the results. - source = tar_waves[:, 0:None] - self.logger.debug(f"Concatenated tar_waves. Shape: {tar_waves.shape}") - - # TODO: In UVR, pitch changing happens here. Consider implementing this as a feature. + result[:, start : start + length] += predictions[index, :, :length] * window + divider[start : start + length] += window + progress.update(len(locations)) - # TODO: In UVR, VR denoise model gets applied here. Consider implementing this as a feature. - - self.logger.debug("Demixing process completed.") - return source + np.divide(result, divider, out=result, where=divider > 0) + return result[:, self.trim : self.trim + mix.shape[-1]] def run_model(self, mix, is_match_mix=False): - """ - Processes the input mix through the model to separate the sources. - Applies STFT, handles spectrum modifications, and runs the model for source separation. - """ - # Applying the STFT to the mix. The mix is moved to the specified device (e.g., GPU) before processing. - # self.logger.debug(f"Running STFT on the mix. Mix shape before STFT: {mix.shape}") + """Infer spectra and keep device tensors through inverse STFT.""" spek = self.stft(mix.to(self.torch_device)) - self.logger.debug(f"STFT applied on mix. Spectrum shape: {spek.shape}") - - # Zeroing out the first 3 bins of the spectrum. This is often done to reduce low-frequency noise. spek[:, :, :3, :] *= 0 - # self.logger.debug("First 3 bins of the spectrum zeroed out.") - - # Handling the case where the mix needs to be matched (is_match_mix = True) if is_match_mix: - # self.logger.debug("Match mix mode is enabled. Converting spectrum to NumPy array.") - spec_pred = spek.cpu().numpy() - self.logger.debug("is_match_mix: spectrum prediction obtained directly from STFT output.") + spec_pred = spek + elif self.enable_denoise: + spec_pred = (self.model_run(-spek) * -0.5) + (self.model_run(spek) * 0.5) else: - # If denoising is enabled, the model is run on both the negative and positive spectrums. - if self.enable_denoise: - # Assuming spek is a tensor and self.model_run can process it directly - spec_pred_neg = self.model_run(-spek) # Ensure this line correctly negates spek and runs the model - spec_pred_pos = self.model_run(spek) - # Ensure both spec_pred_neg and spec_pred_pos are tensors before applying operations - spec_pred = (spec_pred_neg * -0.5) + (spec_pred_pos * 0.5) # [invalid-unary-operand-type] - self.logger.debug("Model run on both negative and positive spectrums for denoising.") - else: - spec_pred = self.model_run(spek) - self.logger.debug("Model run on the spectrum without denoising.") - - # Applying the inverse STFT to convert the spectrum back to the time domain. - result = self.stft.inverse(torch.tensor(spec_pred).to(self.torch_device)).cpu().detach().numpy() - self.logger.debug(f"Inverse STFT applied. Returning result with shape: {result.shape}") + spec_pred = self.model_run(spek) - return result + # Converted Torch models and bound ONNX outputs already reside on the + # inference device; only the generic NumPy fallback needs a transfer. + spec_pred = torch.as_tensor(spec_pred, device=self.torch_device) + return self.stft.inverse(spec_pred).cpu().detach().numpy() diff --git a/audio_separator/separator/uvr_lib_v5/stft.py b/audio_separator/separator/uvr_lib_v5/stft.py index 5ff9f865..aab0128f 100644 --- a/audio_separator/separator/uvr_lib_v5/stft.py +++ b/audio_separator/separator/uvr_lib_v5/stft.py @@ -19,6 +19,12 @@ def __init__(self, logger, n_fft, hop_length, dim_f, device): # Create a Hann window tensor for use in the STFT. self.hann_window = torch.hann_window(window_length=self.n_fft, periodic=True) + def _window_for(self, input_tensor): + # Cache on the effective device, including the complex-operation CPU fallback. + if self.hann_window.device != input_tensor.device: + self.hann_window = self.hann_window.to(input_tensor.device) + return self.hann_window + def __call__(self, input_tensor): original_device = input_tensor.device should_fallback = should_fallback_to_cpu_for_complex_ops(original_device) @@ -28,7 +34,7 @@ def __call__(self, input_tensor): input_tensor = input_tensor.cpu() # Transfer the pre-defined window tensor to the same device as the input tensor. - stft_window = self.hann_window.to(input_tensor.device) + stft_window = self._window_for(input_tensor) # Extract batch dimensions (all dimensions except the last two which are channel and time). batch_dimensions = input_tensor.shape[:-2] @@ -40,10 +46,12 @@ def __call__(self, input_tensor): reshaped_tensor = input_tensor.reshape([-1, time_dim]) # Perform the Short-Time Fourier Transform (STFT) on the reshaped tensor. - stft_output = torch.stft(reshaped_tensor, n_fft=self.n_fft, hop_length=self.hop_length, window=stft_window, center=True, return_complex=False) + stft_output = torch.stft( + reshaped_tensor, n_fft=self.n_fft, hop_length=self.hop_length, window=stft_window, center=True, return_complex=True + ) # Rearrange the dimensions of the STFT output to bring the frequency dimension forward. - permuted_stft_output = stft_output.permute([0, 3, 1, 2]) + permuted_stft_output = torch.view_as_real(stft_output).permute([0, 3, 1, 2]) # Reshape the output to restore the original batch and channel dimensions, while keeping the newly formed frequency and time dimensions. final_output = permuted_stft_output.reshape([*batch_dimensions, channel_dim, 2, -1, permuted_stft_output.shape[-1]]).reshape( @@ -62,7 +70,7 @@ def pad_frequency_dimension(self, input_tensor, batch_dimensions, channel_dim, f Adds zero padding to the frequency dimension of the input tensor. """ # Create a padding tensor for the frequency dimension - freq_padding = torch.zeros([*batch_dimensions, channel_dim, num_freq_bins - freq_dim, time_dim]).to(input_tensor.device) + freq_padding = input_tensor.new_zeros([*batch_dimensions, channel_dim, num_freq_bins - freq_dim, time_dim]) # Concatenate the padding to the input tensor along the frequency dimension. padded_tensor = torch.cat([input_tensor, freq_padding], -2) @@ -94,7 +102,7 @@ def prepare_for_istft(self, padded_tensor, batch_dimensions, channel_dim, num_fr permuted_tensor = flattened_tensor.permute([0, 2, 3, 1]) # Combine real and imaginary parts into a complex tensor. - complex_tensor = permuted_tensor[..., 0] + permuted_tensor[..., 1] * 1.0j + complex_tensor = torch.view_as_complex(permuted_tensor.contiguous()) return complex_tensor @@ -107,7 +115,7 @@ def inverse(self, input_tensor): input_tensor = input_tensor.cpu() # Transfer the pre-defined Hann window tensor to the same device as the input tensor. - stft_window = self.hann_window.to(input_tensor.device) + stft_window = self._window_for(input_tensor) batch_dimensions, channel_dim, freq_dim, time_dim, num_freq_bins = self.calculate_inverse_dimensions(input_tensor) diff --git a/tests/unit/test_mdx_inference.py b/tests/unit/test_mdx_inference.py new file mode 100644 index 00000000..999540ae --- /dev/null +++ b/tests/unit/test_mdx_inference.py @@ -0,0 +1,287 @@ +"""MDX window batching, device transfers, and per-file state contracts.""" + +import subprocess +import sys +import logging +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import onnx +from onnx import TensorProto, helper +import onnxruntime as ort +import pytest +import torch + +from audio_separator.separator.architectures.mdx_separator import MDXSeparator + +DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +def make_separator(tmp_path, batch_dimension="batch", batch_size=3, device="cpu", segment_size=16, providers=None): + shape = [batch_dimension, 4, 17, 16] + graph = helper.make_graph( + [helper.make_node("Neg", ["spectrogram"], ["prediction"])], + "mdx-test", + [helper.make_tensor_value_info("spectrogram", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("prediction", TensorProto.FLOAT, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model.ir_version = 9 + path = tmp_path / "mdx.onnx" + onnx.save(model, str(path)) + config = { + "logger": logging.getLogger(__name__), + "log_level": logging.INFO, + "torch_device": torch.device(device), + "torch_device_cpu": torch.device("cpu"), + "onnx_execution_provider": providers or ["CPUExecutionProvider"], + "model_path": str(path), + "model_name": "mdx-test", + "model_data": {"primary_stem": "Vocals", "compensate": 1.0, "mdx_dim_f_set": 17, "mdx_dim_t_set": 4, "mdx_n_fft_scale_set": 32}, + "normalization_threshold": 0.9, + "amplification_threshold": 0.0, + "sample_rate": 44100, + "output_format": "WAV", + } + params = {"segment_size": segment_size, "hop_length": 8, "overlap": 0.25, "batch_size": batch_size, "enable_denoise": False} + return MDXSeparator(config, params) + + +@pytest.mark.parametrize("batch_dimension", [1, 2, "batch", None]) +def test_onnx_batch_metadata_and_input_names(tmp_path, batch_dimension): + separator = make_separator(tmp_path, batch_dimension=batch_dimension) + batch = batch_dimension if isinstance(batch_dimension, int) else 3 + spectrum = torch.randn(batch, 4, 17, 16) + np.testing.assert_array_equal(separator.model_run(spectrum), -spectrum.numpy()) + assert separator.inference_batch_size == batch + assert separator.fixed_batch_size == (batch_dimension if isinstance(batch_dimension, int) else None) + + +@pytest.mark.parametrize("overlap", [0, 0.25, 0.75, 0.99]) +@pytest.mark.parametrize("length", [1, 1009]) +@pytest.mark.parametrize("fixed_batch", [None, 2]) +def test_batched_demix_preserves_samples_and_pads_last_batch(tmp_path, overlap, length, fixed_batch): + separator = make_separator(tmp_path, batch_dimension=fixed_batch or "batch") + mix = np.random.default_rng(42).normal(0, 0.1, (2, length)).astype(np.float32) + separator.overlap = overlap + calls = [] + + def predict(batch, is_match_mix=False): + calls.append(batch.shape[0]) + return batch.cpu().numpy() * 0.3 + + separator.run_model = predict + actual = separator.demix(mix) + assert np.isfinite(actual).all() + assert actual.shape == mix.shape + np.testing.assert_allclose(actual, mix * 0.3, atol=1e-7) + if fixed_batch: + assert set(calls) == {fixed_batch} + else: + assert all(1 <= batch <= 3 for batch in calls) + if length > separator.chunk_size: + assert max(calls) == 3 + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("batch_dimension", [1, 2, "batch"]) +def test_real_onnx_demix_matches_single_window_batches(tmp_path, device, batch_dimension): + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if device == "cuda" and CUDA_ONNX_AVAILABLE else ["CPUExecutionProvider"] + separator = make_separator(tmp_path, batch_dimension=batch_dimension, device=device, providers=providers) + mix = np.random.default_rng(42).normal(0, 0.1, (2, 1009)).astype(np.float32) + batched = separator.demix(mix) + separator.inference_batch_size = separator.fixed_batch_size or 1 + np.testing.assert_allclose(separator.demix(mix), batched, atol=2e-7) + + +@pytest.mark.parametrize("denoise", [False, True]) +@pytest.mark.parametrize("match_mix", [False, True]) +def test_torch_predictions_remain_on_device(tmp_path, denoise, match_mix): + separator = make_separator(tmp_path) + separator.enable_denoise = denoise + prediction = torch.randn(2, 4, 17, 16) + separator.stft = Mock(return_value=prediction.clone()) + predictions = [] + + def predict(value): + result = value * 0.3 + predictions.append(result) + return result + + separator.model_run = Mock(side_effect=predict) + seen = [] + + def inverse(value): + seen.append(value) + return torch.zeros(2, 2, 120) + + separator.stft.inverse.side_effect = inverse + separator.run_model(torch.zeros(2, 2, 120), is_match_mix=match_mix) + expected = prediction.clone() + expected[:, :, :3] = 0 + if not match_mix: + expected *= 0.3 + torch.testing.assert_close(seen[0], expected) + assert separator.model_run.call_count == (0 if match_mix else 2 if denoise else 1) + if match_mix: + assert seen[0] is separator.stft.return_value + elif not denoise: + assert seen[0] is predictions[0] + + +@pytest.mark.parametrize( + "selection,expected_names", [(None, ["Instrumental", "Vocals"]), ("vocals", ["Vocals"]), ("instrumental", ["Instrumental"])] +) +@pytest.mark.parametrize("invert", [False, True]) +def test_secondary_reconstruction_only_runs_when_needed(tmp_path, monkeypatch, selection, expected_names, invert): + separator = make_separator(tmp_path) + separator.output_single_stem = selection + separator.invert_using_spec = invert + separator.prepare_mix = Mock(return_value=np.ones((2, 32), dtype=np.float32) * 0.2) + separator.demix = Mock(side_effect=lambda mix, is_match_mix=False: mix * (1 if is_match_mix else 0.3)) + separator.get_stem_output_path = Mock(side_effect=lambda name, custom: name + ".wav") + separator.final_process = Mock() + invert_stem = Mock(side_effect=lambda raw, primary: raw.T - primary) + monkeypatch.setattr("audio_separator.separator.architectures.mdx_separator.spec_utils.invert_stem", invert_stem) + assert separator.separate("input.wav") == [name + ".wav" for name in expected_names] + needs_match = invert and selection != "vocals" + assert separator.demix.call_count == (2 if needs_match else 1) + assert invert_stem.call_count == int(needs_match) + + +def test_reused_separator_does_not_export_previous_audio(tmp_path): + separator = make_separator(tmp_path) + separator.prepare_mix = Mock(side_effect=[np.ones((2, 32), dtype=np.float32) * level for level in (0.1, 0.2)]) + separator.demix = Mock(side_effect=lambda mix: mix * 0.3) + separator.get_stem_output_path = Mock(side_effect=lambda name, custom: name + ".wav") + saved = [] + separator.final_process = lambda path, source, name: saved.append((name, source.copy())) + separator.separate("first.wav") + separator.separate("second.wav") + assert [name for name, _ in saved] == ["Instrumental", "Vocals"] * 2 + assert not np.array_equal(saved[0][1], saved[2][1]) + assert not np.array_equal(saved[1][1], saved[3][1]) + + +def test_native_module_does_not_import_torch_converter(): + subprocess.run( + [ + sys.executable, + "-c", + "import sys; import audio_separator.separator.architectures.mdx_separator; assert 'onnx2torch' not in sys.modules", + ], + check=True, + capture_output=True, + text=True, + ) + + +def test_converted_model_keeps_requested_batch_size(tmp_path, monkeypatch): + import sys + + model = torch.nn.Identity() + monkeypatch.setitem(sys.modules, "onnx2torch", SimpleNamespace(convert=Mock(return_value=model))) + separator = make_separator(tmp_path, segment_size=32, batch_size=3) + assert separator.uses_pytorch_inference + assert separator.fixed_batch_size is None + assert separator.inference_batch_size == 3 + assert separator.model_run is model + + +def test_invalid_chunk_size_fails_before_allocating_windows(tmp_path): + separator = make_separator(tmp_path) + separator.segment_size = 2 + with pytest.raises(ValueError, match="FFT window"): + separator.demix(np.zeros((2, 32), dtype=np.float32)) + + +CUDA_ONNX_AVAILABLE = torch.cuda.is_available() and "CUDAExecutionProvider" in ort.get_available_providers() + + +@pytest.mark.skipif(not CUDA_ONNX_AVAILABLE, reason="CUDA Torch and ONNX Runtime are required") +@pytest.mark.parametrize("batch_dimension", [1, 2, "batch"]) +def test_cuda_binding_handles_noncontiguous_inputs_and_nondefault_streams(tmp_path, batch_dimension): + separator = make_separator( + tmp_path, + device="cuda", + batch_dimension=batch_dimension, + providers=[("CUDAExecutionProvider", {"do_copy_in_default_stream": True}), "CPUExecutionProvider"], + ) + assert "CUDAExecutionProvider" in separator.onnx_session.get_providers() + assert separator.onnx_stream is not None + assert separator.onnx_session.get_provider_options()["CUDAExecutionProvider"]["has_user_compute_stream"] == "1" + batch_sizes = [1, 3] if batch_dimension == "batch" else [batch_dimension] + for batch in batch_sizes: + stream = torch.cuda.Stream() + with torch.cuda.stream(stream), torch.inference_mode(): + spectrum = torch.randn(batch, 4, 16, 17, device="cuda").transpose(-1, -2) + actual = separator.model_run(spectrum) + assert actual.device == spectrum.device + torch.testing.assert_close(actual, -spectrum, rtol=0, atol=0) + saved = actual.clone() + separator.model_run(spectrum * 2) + torch.testing.assert_close(actual, saved, rtol=0, atol=0) + stream.synchronize() + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("output_shape", [None, [None, 4, 17, 16], ["unknown", 4, 17, 16]]) +def test_unknown_output_shapes_use_allocating_fallback(tmp_path, device, output_shape): + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if device == "cuda" and CUDA_ONNX_AVAILABLE else ["CPUExecutionProvider"] + separator = make_separator(tmp_path, device=device, providers=providers) + metadata = separator.onnx_output + separator.onnx_output = SimpleNamespace(name=metadata.name, type=metadata.type, shape=output_shape) + spectrum = torch.randn(3, 4, 17, 16, device=device) + result = separator.model_run(spectrum) + assert isinstance(result, np.ndarray) + np.testing.assert_array_equal(result, -spectrum.cpu().numpy()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA Torch is required") +def test_cuda_provider_failure_retains_numpy_fallback(tmp_path, monkeypatch, caplog): + session_factory = ort.InferenceSession + + def cpu_session(*args, **kwargs): + kwargs["providers"] = ["CPUExecutionProvider"] + return session_factory(*args, **kwargs) + + monkeypatch.setattr(ort, "InferenceSession", cpu_session) + separator = make_separator(tmp_path, device="cuda", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) + assert separator.onnx_stream is None + spectrum = torch.randn(3, 4, 17, 16, device="cuda") + result = separator.model_run(spectrum) + assert isinstance(result, np.ndarray) + np.testing.assert_array_equal(result, -spectrum.cpu().numpy()) + assert "could not activate requested provider CUDAExecutionProvider" in caplog.text + + +@pytest.mark.parametrize("overlap", [-0.1, 1.0, float("inf"), float("nan")]) +def test_invalid_overlap_fails_before_scheduling_windows(tmp_path, overlap): + separator = make_separator(tmp_path) + separator.overlap = overlap + with pytest.raises(ValueError, match="MDX overlap"): + separator.demix(np.zeros((2, 32), dtype=np.float32)) + + +@pytest.mark.parametrize("batch_size", [0, -1]) +def test_nonpositive_batch_size_is_rejected(tmp_path, batch_size): + separator = make_separator(tmp_path, batch_size=batch_size) + with pytest.raises(ValueError, match="MDX batch size"): + separator.demix(np.zeros((2, 32), dtype=np.float32)) + + +def test_native_separator_releases_without_waiting_for_cycle_collection(tmp_path): + import gc + import weakref + + separator = make_separator(tmp_path) + reference = weakref.ref(separator) + collection_was_enabled = gc.isenabled() + gc.disable() + try: + del separator + assert reference() is None + finally: + if collection_was_enabled: + gc.enable() diff --git a/tests/unit/test_stft.py b/tests/unit/test_stft.py index bc1fc978..33273b93 100644 --- a/tests/unit/test_stft.py +++ b/tests/unit/test_stft.py @@ -1,5 +1,6 @@ import unittest import numpy as np +import pytest import torch from unittest.mock import Mock from audio_separator.separator.uvr_lib_v5.stft import STFT @@ -162,5 +163,47 @@ def debug(self, message): pass +@pytest.mark.parametrize("device", ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])) +@pytest.mark.parametrize("force_cpu_fallback", [False, True]) +def test_stft_preserves_mdx_channel_order_and_round_trip(device, force_cpu_fallback, monkeypatch): + if force_cpu_fallback: + monkeypatch.setenv("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX", "1") + else: + monkeypatch.delenv("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX", raising=False) + generator = torch.Generator().manual_seed(17) + audio = torch.randn(3, 2, 2048, generator=generator) * 0.1 + transform = STFT(Mock(), n_fft=256, hop_length=64, dim_f=129, device=torch.device(device)) + window = torch.hann_window(256) + complex_spectrum = torch.stft(audio.reshape(-1, 2048), n_fft=256, hop_length=64, window=window, return_complex=True) + complex_spectrum = complex_spectrum.reshape(3, 2, 129, -1) + expected = torch.stack( + [complex_spectrum[:, 0].real, complex_spectrum[:, 0].imag, complex_spectrum[:, 1].real, complex_spectrum[:, 1].imag], dim=1 + ) + actual = transform(audio.to(device)) + torch.testing.assert_close(actual.cpu(), expected, rtol=2e-5, atol=2e-6) + window_pointer = transform.hann_window.data_ptr() + restored = transform.inverse(actual) + torch.testing.assert_close(restored.cpu(), audio, rtol=2e-5, atol=2e-6) + assert restored.device.type == device + assert transform.hann_window.device.type == ("cpu" if force_cpu_fallback else device) + assert transform.hann_window.data_ptr() == window_pointer + transform(audio.to(device)) + assert transform.hann_window.data_ptr() == window_pointer + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_frequency_padding_preserves_dtype_and_zeroes_omitted_bins(dtype): + transform = STFT(Mock(), n_fft=256, hop_length=64, dim_f=100, device=torch.device("cpu")) + spectrum = torch.randn(3, 4, 100, 17, dtype=dtype) + dimensions = transform.calculate_inverse_dimensions(spectrum) + padded = transform.pad_frequency_dimension(spectrum, *dimensions) + assert padded.dtype == dtype + torch.testing.assert_close(padded[..., :100, :], spectrum) + assert torch.count_nonzero(padded[..., 100:, :]) == 0 + actual = transform.prepare_for_istft(padded, dimensions[0], dimensions[1], dimensions[-1], dimensions[-2]) + real, imaginary = padded.reshape(-1, 2, 129, 17).unbind(dim=1) + torch.testing.assert_close(actual, torch.complex(real, imaginary), rtol=0, atol=0) + + if __name__ == "__main__": unittest.main() From 7aa1e873ae7e3103d69ac50b726ac4035648c032 Mon Sep 17 00:00:00 2001 From: Blaise Date: Mon, 19 Oct 2026 09:58:49 +0200 Subject: [PATCH 2/3] docs: document MDX and STFT functions and regression tests --- .../separator/architectures/mdx_separator.py | 1 + audio_separator/separator/uvr_lib_v5/stft.py | 5 +++++ tests/unit/test_mdx_inference.py | 20 +++++++++++++++++++ tests/unit/test_stft.py | 13 ++++++++++++ 4 files changed, 39 insertions(+) diff --git a/audio_separator/separator/architectures/mdx_separator.py b/audio_separator/separator/architectures/mdx_separator.py index 5dd87ee1..48472915 100644 --- a/audio_separator/separator/architectures/mdx_separator.py +++ b/audio_separator/separator/architectures/mdx_separator.py @@ -18,6 +18,7 @@ class MDXSeparator(CommonSeparator): """ def __init__(self, common_config, arch_config): + """Configure MDX windowing and load the native ONNX or converted Torch model.""" # Any configuration values which can be shared between architectures should be set already in CommonSeparator, # e.g. user-specified functionality choices (self.output_single_stem) or common model parameters (self.primary_stem_name) super().__init__(config=common_config) diff --git a/audio_separator/separator/uvr_lib_v5/stft.py b/audio_separator/separator/uvr_lib_v5/stft.py index aab0128f..aa708f58 100644 --- a/audio_separator/separator/uvr_lib_v5/stft.py +++ b/audio_separator/separator/uvr_lib_v5/stft.py @@ -11,6 +11,7 @@ class STFT: """ def __init__(self, logger, n_fft, hop_length, dim_f, device): + """Configure spectral dimensions and create the reusable Hann window.""" self.logger = logger self.n_fft = n_fft self.hop_length = hop_length @@ -20,12 +21,14 @@ def __init__(self, logger, n_fft, hop_length, dim_f, device): self.hann_window = torch.hann_window(window_length=self.n_fft, periodic=True) def _window_for(self, input_tensor): + """Return the cached Hann window on the tensor's effective compute device.""" # Cache on the effective device, including the complex-operation CPU fallback. if self.hann_window.device != input_tensor.device: self.hann_window = self.hann_window.to(input_tensor.device) return self.hann_window def __call__(self, input_tensor): + """Pack channel spectra as alternating real and imaginary planes on the input device.""" original_device = input_tensor.device should_fallback = should_fallback_to_cpu_for_complex_ops(original_device) @@ -78,6 +81,7 @@ def pad_frequency_dimension(self, input_tensor, batch_dimensions, channel_dim, f return padded_tensor def calculate_inverse_dimensions(self, input_tensor): + """Return the batch, channel, spectral dimensions and full FFT bin count.""" # Extract batch dimensions and frequency-time dimensions. batch_dimensions = input_tensor.shape[:-3] channel_dim, freq_dim, time_dim = input_tensor.shape[-3:] @@ -107,6 +111,7 @@ def prepare_for_istft(self, padded_tensor, batch_dimensions, channel_dim, num_fr return complex_tensor def inverse(self, input_tensor): + """Reconstruct stereo audio and restore the caller's input device.""" original_device = input_tensor.device should_fallback = should_fallback_to_cpu_for_complex_ops(original_device) diff --git a/tests/unit/test_mdx_inference.py b/tests/unit/test_mdx_inference.py index 999540ae..86615109 100644 --- a/tests/unit/test_mdx_inference.py +++ b/tests/unit/test_mdx_inference.py @@ -19,6 +19,7 @@ def make_separator(tmp_path, batch_dimension="batch", batch_size=3, device="cpu", segment_size=16, providers=None): + """Create a real ONNX test graph and configure a separator without downloading weights.""" shape = [batch_dimension, 4, 17, 16] graph = helper.make_graph( [helper.make_node("Neg", ["spectrogram"], ["prediction"])], @@ -50,6 +51,7 @@ def make_separator(tmp_path, batch_dimension="batch", batch_size=3, device="cpu" @pytest.mark.parametrize("batch_dimension", [1, 2, "batch", None]) def test_onnx_batch_metadata_and_input_names(tmp_path, batch_dimension): + """Honor declared batch dimensions and resolve nonstandard ONNX input names.""" separator = make_separator(tmp_path, batch_dimension=batch_dimension) batch = batch_dimension if isinstance(batch_dimension, int) else 3 spectrum = torch.randn(batch, 4, 17, 16) @@ -62,12 +64,14 @@ def test_onnx_batch_metadata_and_input_names(tmp_path, batch_dimension): @pytest.mark.parametrize("length", [1, 1009]) @pytest.mark.parametrize("fixed_batch", [None, 2]) def test_batched_demix_preserves_samples_and_pads_last_batch(tmp_path, overlap, length, fixed_batch): + """Preserve overlap-add samples and honor batch limits for short and partial windows.""" separator = make_separator(tmp_path, batch_dimension=fixed_batch or "batch") mix = np.random.default_rng(42).normal(0, 0.1, (2, length)).astype(np.float32) separator.overlap = overlap calls = [] def predict(batch, is_match_mix=False): + """Record the submitted batch size and return deterministic scaled audio.""" calls.append(batch.shape[0]) return batch.cpu().numpy() * 0.3 @@ -87,6 +91,7 @@ def predict(batch, is_match_mix=False): @pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_dimension", [1, 2, "batch"]) def test_real_onnx_demix_matches_single_window_batches(tmp_path, device, batch_dimension): + """Match serial-window reconstruction when real ONNX inference runs in batches.""" providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if device == "cuda" and CUDA_ONNX_AVAILABLE else ["CPUExecutionProvider"] separator = make_separator(tmp_path, batch_dimension=batch_dimension, device=device, providers=providers) mix = np.random.default_rng(42).normal(0, 0.1, (2, 1009)).astype(np.float32) @@ -98,6 +103,7 @@ def test_real_onnx_demix_matches_single_window_batches(tmp_path, device, batch_d @pytest.mark.parametrize("denoise", [False, True]) @pytest.mark.parametrize("match_mix", [False, True]) def test_torch_predictions_remain_on_device(tmp_path, denoise, match_mix): + """Preserve prediction values and tensor identity through the inverse STFT boundary.""" separator = make_separator(tmp_path) separator.enable_denoise = denoise prediction = torch.randn(2, 4, 17, 16) @@ -105,6 +111,7 @@ def test_torch_predictions_remain_on_device(tmp_path, denoise, match_mix): predictions = [] def predict(value): + """Retain each generated prediction so the test can detect unnecessary tensor copies.""" result = value * 0.3 predictions.append(result) return result @@ -113,6 +120,7 @@ def predict(value): seen = [] def inverse(value): + """Capture the prediction received by inverse STFT and return a valid audio-shaped tensor.""" seen.append(value) return torch.zeros(2, 2, 120) @@ -135,6 +143,7 @@ def inverse(value): ) @pytest.mark.parametrize("invert", [False, True]) def test_secondary_reconstruction_only_runs_when_needed(tmp_path, monkeypatch, selection, expected_names, invert): + """Run mixture reconstruction only when a requested secondary stem needs spectral inversion.""" separator = make_separator(tmp_path) separator.output_single_stem = selection separator.invert_using_spec = invert @@ -151,6 +160,7 @@ def test_secondary_reconstruction_only_runs_when_needed(tmp_path, monkeypatch, s def test_reused_separator_does_not_export_previous_audio(tmp_path): + """Replace cached sources when one separator processes successive input files.""" separator = make_separator(tmp_path) separator.prepare_mix = Mock(side_effect=[np.ones((2, 32), dtype=np.float32) * level for level in (0.1, 0.2)]) separator.demix = Mock(side_effect=lambda mix: mix * 0.3) @@ -165,6 +175,7 @@ def test_reused_separator_does_not_export_previous_audio(tmp_path): def test_native_module_does_not_import_torch_converter(): + """Keep native MDX imports independent of the optional Torch conversion module.""" subprocess.run( [ sys.executable, @@ -178,6 +189,7 @@ def test_native_module_does_not_import_torch_converter(): def test_converted_model_keeps_requested_batch_size(tmp_path, monkeypatch): + """Use the requested window batch size when ONNX conversion produces a Torch model.""" import sys model = torch.nn.Identity() @@ -190,6 +202,7 @@ def test_converted_model_keeps_requested_batch_size(tmp_path, monkeypatch): def test_invalid_chunk_size_fails_before_allocating_windows(tmp_path): + """Reject segment geometry that cannot accommodate the FFT trimming window.""" separator = make_separator(tmp_path) separator.segment_size = 2 with pytest.raises(ValueError, match="FFT window"): @@ -202,6 +215,7 @@ def test_invalid_chunk_size_fails_before_allocating_windows(tmp_path): @pytest.mark.skipif(not CUDA_ONNX_AVAILABLE, reason="CUDA Torch and ONNX Runtime are required") @pytest.mark.parametrize("batch_dimension", [1, 2, "batch"]) def test_cuda_binding_handles_noncontiguous_inputs_and_nondefault_streams(tmp_path, batch_dimension): + """Synchronize CUDA buffers across streams and keep results independent between calls.""" separator = make_separator( tmp_path, device="cuda", @@ -228,6 +242,7 @@ def test_cuda_binding_handles_noncontiguous_inputs_and_nondefault_streams(tmp_pa @pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("output_shape", [None, [None, 4, 17, 16], ["unknown", 4, 17, 16]]) def test_unknown_output_shapes_use_allocating_fallback(tmp_path, device, output_shape): + """Use NumPy allocation when output metadata cannot establish safe buffer dimensions.""" providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if device == "cuda" and CUDA_ONNX_AVAILABLE else ["CPUExecutionProvider"] separator = make_separator(tmp_path, device=device, providers=providers) metadata = separator.onnx_output @@ -240,9 +255,11 @@ def test_unknown_output_shapes_use_allocating_fallback(tmp_path, device, output_ @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA Torch is required") def test_cuda_provider_failure_retains_numpy_fallback(tmp_path, monkeypatch, caplog): + """Retain correct inference and a warning when the requested CUDA provider is unavailable.""" session_factory = ort.InferenceSession def cpu_session(*args, **kwargs): + """Force a real CPU session to simulate CUDA provider activation failure.""" kwargs["providers"] = ["CPUExecutionProvider"] return session_factory(*args, **kwargs) @@ -258,6 +275,7 @@ def cpu_session(*args, **kwargs): @pytest.mark.parametrize("overlap", [-0.1, 1.0, float("inf"), float("nan")]) def test_invalid_overlap_fails_before_scheduling_windows(tmp_path, overlap): + """Reject invalid overlap values before they can schedule excessive or invalid window work.""" separator = make_separator(tmp_path) separator.overlap = overlap with pytest.raises(ValueError, match="MDX overlap"): @@ -266,12 +284,14 @@ def test_invalid_overlap_fails_before_scheduling_windows(tmp_path, overlap): @pytest.mark.parametrize("batch_size", [0, -1]) def test_nonpositive_batch_size_is_rejected(tmp_path, batch_size): + """Reject batch sizes that cannot advance the window-processing loop.""" separator = make_separator(tmp_path, batch_size=batch_size) with pytest.raises(ValueError, match="MDX batch size"): separator.demix(np.zeros((2, 32), dtype=np.float32)) def test_native_separator_releases_without_waiting_for_cycle_collection(tmp_path): + """Release native separator ownership immediately without relying on cyclic garbage collection.""" import gc import weakref diff --git a/tests/unit/test_stft.py b/tests/unit/test_stft.py index 33273b93..639055af 100644 --- a/tests/unit/test_stft.py +++ b/tests/unit/test_stft.py @@ -22,6 +22,7 @@ class TestSTFT(unittest.TestCase): def setUp(self): + """Configure a CPU STFT instance with a complete frequency spectrum.""" self.n_fft = 2048 self.hop_length = 512 self.dim_f = 1025 @@ -29,12 +30,14 @@ def setUp(self): self.stft = STFT(logger=Mock(), n_fft=self.n_fft, hop_length=self.hop_length, dim_f=self.dim_f, device=self.device) def create_mock_tensor(self, shape, device=None): + """Create random test input on the requested device.""" tensor = torch.rand(shape) if device: tensor = tensor.to(device) return tensor def test_stft_initialization(self): + """Retain the configured transform dimensions, device and Hann window.""" self.assertEqual(self.stft.n_fft, self.n_fft) self.assertEqual(self.stft.hop_length, self.hop_length) self.assertEqual(self.stft.dim_f, self.dim_f) @@ -42,6 +45,7 @@ def test_stft_initialization(self): self.assertIsInstance(self.stft.hann_window, torch.Tensor) def test_stft_call(self): + """Produce the expected frequency-bin and time-frame dimensions.""" input_tensor = self.create_mock_tensor((1, 16000)) # Apply STFT @@ -76,6 +80,7 @@ def test_stft_call(self): self.assertEqual(stft_result.shape[-2:], expected_shape) def test_calculate_inverse_dimensions(self): + """Recover inverse-transform dimensions from the packed spectrum shape.""" # Create a sample input tensor sample_input = torch.randn(1, 2, 500, 32) # Batch, Channel, Frequency, Time dimensions batch_dims, channel_dim, freq_dim, time_dim, num_freq_bins = self.stft.calculate_inverse_dimensions(sample_input) @@ -91,6 +96,7 @@ def test_calculate_inverse_dimensions(self): self.assertEqual(num_freq_bins, expected_num_freq_bins) def test_pad_frequency_dimension(self): + """Extend a truncated spectrum to the full inverse-transform frequency dimension.""" # Create a sample input tensor sample_input = torch.randn(1, 2, 500, 32) # Batch, Channel, Frequency, Time dimensions batch_dims, channel_dim, freq_dim, time_dim, num_freq_bins = self.stft.calculate_inverse_dimensions(sample_input) @@ -105,6 +111,7 @@ def test_pad_frequency_dimension(self): self.assertEqual(padded_output.shape[-2], expected_freq_dim) def test_prepare_for_istft(self): + """Combine packed real and imaginary planes into the expected complex spectrum shape.""" # Create a sample input tensor sample_input = torch.randn(1, 2, 500, 32) # Batch, Channel, Frequency, Time dimensions batch_dims, channel_dim, freq_dim, time_dim, num_freq_bins = self.stft.calculate_inverse_dimensions(sample_input) @@ -123,6 +130,7 @@ def test_prepare_for_istft(self): self.assertEqual(complex_tensor.shape, expected_shape) def test_inverse_stft(self): + """Preserve the legacy inverse output shape on CPU.""" # Create a mock tensor with the correct input shape input_tensor = torch.rand(1, 2, 1025, 32) # shape matching output of STFT @@ -140,6 +148,7 @@ def test_inverse_stft(self): @unittest.skipIf(not torch.backends.mps.is_available(), "MPS not available") def test_stft_with_mps_device(self): + """Transform input on an available MPS device.""" mps_device = torch.device("mps") self.stft.device = mps_device input_tensor = self.create_mock_tensor((1, 16000), device=mps_device) @@ -149,6 +158,7 @@ def test_stft_with_mps_device(self): @unittest.skipIf(not torch.backends.mps.is_available(), "MPS not available") def test_inverse_with_mps_device(self): + """Reconstruct audio from spectra supplied on an available MPS device.""" mps_device = torch.device("mps") self.stft.device = mps_device input_tensor = self.create_mock_tensor((1, 2, 1025, 32), device=mps_device) @@ -160,12 +170,14 @@ def test_inverse_with_mps_device(self): # Mock logger to use in tests class MockLogger: def debug(self, message): + """Accept debug messages without emitting output during STFT tests.""" pass @pytest.mark.parametrize("device", ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])) @pytest.mark.parametrize("force_cpu_fallback", [False, True]) def test_stft_preserves_mdx_channel_order_and_round_trip(device, force_cpu_fallback, monkeypatch): + """Preserve channel packing, reconstruction and window reuse through native or CPU fallback paths.""" if force_cpu_fallback: monkeypatch.setenv("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX", "1") else: @@ -193,6 +205,7 @@ def test_stft_preserves_mdx_channel_order_and_round_trip(device, force_cpu_fallb @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) def test_frequency_padding_preserves_dtype_and_zeroes_omitted_bins(dtype): + """Preserve spectral values and dtype while zero-filling omitted frequency bins.""" transform = STFT(Mock(), n_fft=256, hop_length=64, dim_f=100, device=torch.device("cpu")) spectrum = torch.randn(3, 4, 100, 17, dtype=dtype) dimensions = transform.calculate_inverse_dimensions(spectrum) From 3e76b6a89e625c498b9ef4c2f918fbf916f41d7d Mon Sep 17 00:00:00 2001 From: Blaise Date: Mon, 19 Oct 2026 16:09:09 +0200 Subject: [PATCH 3/3] ci: persist FFmpeg paths for Windows integration tests --- .github/workflows/run-integration-tests.yaml | 28 +++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/.github/workflows/run-integration-tests.yaml b/.github/workflows/run-integration-tests.yaml index 6ff4c467..f80e2dd9 100644 --- a/.github/workflows/run-integration-tests.yaml +++ b/.github/workflows/run-integration-tests.yaml @@ -195,9 +195,35 @@ jobs: python-version: '3.12' cache: poetry - name: Install ffmpeg - run: choco install ffmpeg -y --no-progress + shell: pwsh + run: | + choco install ffmpeg -y --no-progress + if ($LASTEXITCODE -ne 0) { throw "FFmpeg installation failed with exit code $LASTEXITCODE" } + # Chocolatey updates the persistent PATH, not this runner process. + $env:Path = [Environment]::GetEnvironmentVariable('Path', 'Machine') + ';' + [Environment]::GetEnvironmentVariable('Path', 'User') + ';' + $env:Path + foreach ($tool in @('ffmpeg.exe', 'ffprobe.exe')) { + $executable = (Get-Command $tool -CommandType Application -ErrorAction Stop).Source + Split-Path -Parent $executable | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append + } - name: Install Poetry dependencies (CPU) run: poetry install -E cpu + - name: Verify FFmpeg export in the Poetry environment + shell: pwsh + run: | + @' + import subprocess + import tempfile + from pathlib import Path + from pydub import AudioSegment + + for executable in ('ffmpeg', 'ffprobe'): + subprocess.run([executable, '-version'], check=True) + with tempfile.TemporaryDirectory() as directory: + output = Path(directory) / 'ffmpeg-check.flac' + AudioSegment.silent(duration=100).export(str(output), format='flac').close() + assert output.stat().st_size > 0 + assert len(AudioSegment.from_file(str(output))) == 100 + '@ | poetry run python - - name: Cache separation models uses: actions/cache@v4 with: