diff --git a/.github/workflows/run-unit-tests.yaml b/.github/workflows/run-unit-tests.yaml index 338419f..429feda 100644 --- a/.github/workflows/run-unit-tests.yaml +++ b/.github/workflows/run-unit-tests.yaml @@ -82,7 +82,16 @@ jobs: cache: poetry # caching dependencies from poetry.lock - name: Install Poetry dependencies (CPU) - run: poetry install -E cpu + # Large wheels intermittently fail mid-download on hosted macOS runners, + # so retry the install; completed downloads are reused from the cache. + run: | + n=0 + until poetry install -E cpu; do + n=$((n+1)) + if [ "$n" -ge 3 ]; then exit 1; fi + echo "poetry install failed; retry $n/2 in 15s" + sleep 15 + done - name: Run unit tests with coverage run: | diff --git a/README.md b/README.md index ab0451e..51cda8e 100644 --- a/README.md +++ b/README.md @@ -25,9 +25,10 @@ The simplest (and probably most used) use case for this package is to separate a - [Installation ๐Ÿ› ๏ธ](#installation-%EF%B8%8F) - [๐Ÿณ Docker](#-docker) - [๐ŸŽฎ Nvidia GPU with CUDA or ๐Ÿงช Google Colab](#-nvidia-gpu-with-cuda-or--google-colab) - - [๏ฃฟ Apple Silicon, macOS Sonoma+ with M1 or newer CPU (CoreML acceleration)](#-apple-silicon-macos-sonoma-with-m1-or-newer-cpu-coreml-acceleration) + - [๏ฃฟ Apple Silicon, macOS Sonoma+ with M1 or newer (CoreML and MPS acceleration)](#-apple-silicon-macos-sonoma-with-m1-or-newer-coreml-and-mps-acceleration) - [๐Ÿข No hardware acceleration, CPU only](#-no-hardware-acceleration-cpu-only) - [๐ŸชŸ Windows AMD / Intel GPU with DirectML (experimental)](#-windows-amd--intel-gpu-with-directml-experimental) + - [Inference precision and compilation](#inference-precision-and-compilation) - [๐ŸŽฅ FFmpeg dependency](#-ffmpeg-dependency) - [GPU / CUDA specific installation steps with Pip](#gpu--cuda-specific-installation-steps-with-pip) - [Multiple CUDA library versions may be needed](#multiple-cuda-library-versions-may-be-needed) @@ -113,16 +114,41 @@ Docker: beveradb/audio-separator:gpu ``` -### ๏ฃฟ Apple Silicon, macOS Sonoma+ with M1 or newer CPU (CoreML acceleration) +### ๏ฃฟ Apple Silicon, macOS Sonoma+ with M1 or newer (CoreML and MPS acceleration) -๐Ÿ’ฌ If successfully configured, you should see this log message when running `audio-separator --env_info`: - `ONNXruntime has CoreMLExecutionProvider available, enabling acceleration` +PyTorch models use the MPS device, while ONNX models use the CoreML execution provider when it is available. Pip: ```sh pip install "audio-separator[cpu]" ``` +๐Ÿ’ฌ If successfully configured, `audio-separator --env_info` logs: + +```text +Apple Silicon MPS/CoreML is available in Torch and processor is ARM, setting Torch device to MPS +ONNXruntime has CoreMLExecutionProvider available, enabling acceleration +``` + +On Apple Silicon, audio-separator requires PyTorch 2.13 or newer (and earlier than PyTorch 3). Other platforms continue to support PyTorch 2.3 or newer. PyTorch 2.13 is the validated baseline for the fast paths below, and its Apple Silicon wheels target macOS 14 (Sonoma) or newer. + +If the runtime probe finds an unsupported complex operation, inference automatically uses the compatible CPU fallback for that spectral work. + +**Model architecture status on Apple Silicon:** + +| Architecture | Model types | Accelerator | +|---|---|---| +| MDX | `.onnx` | CoreML when the execution provider is available | +| VR | `.pth` | PyTorch MPS | +| Demucs | `.yaml` | PyTorch MPS; supported spectral operations stay on-device | +| MDXC / RoFormer | `.ckpt` / `.yaml` | PyTorch MPS; supported spectral operations and bounded overlap-add buffers stay on-device | + +For long inputs, the full-track working buffers can move to CPU automatically to leave Metal memory free for inference โ€” model inference itself always stays on MPS. The threshold adapts to the device's free Metal working set and is never below 1 GiB. + +Set `AUDIO_SEPARATOR_MPS_BUFFER_BUDGET_GIB` to override that threshold in GiB, for example `AUDIO_SEPARATOR_MPS_BUFFER_BUDGET_GIB=8`. + +Set `AUDIO_SEPARATOR_FORCE_CPU_COMPLEX=1` to force the legacy CPU path for complex spectral operations when diagnosing an MPS compatibility issue. + ### ๐Ÿข No hardware acceleration, CPU only Conda: @@ -171,6 +197,50 @@ separations on an NVIDIA T4 in WDDM mode. If `AUDIO_SEPARATOR_FORCE_DML_MDXC=1` works on your GPU, please [open an issue](https://github.com/nomadkaraoke/python-audio-separator/issues) with your `--env_info` output โ€” torch-directml behaves differently across vendors. +### Inference precision and compilation + +Three opt-in flags control PyTorch inference: + +- `--use_autocast` runs supported operations through PyTorch autocast. +- `--use_native_fp16` converts a verified model to native float16. It is mutually exclusive with `--use_autocast`. +- `--use_torch_compile` enables regional compilation for verified model, device, and precision combinations. It is orthogonal to precision, so it can be used with float32, autocast, or native float16 where supported. + +The verified combinations are intentionally conservative: + +| Device | Model family | Effective precision | Regional `torch.compile` | +|---|---|---|---| +| MPS or CUDA | MelBand RoFormer, BS-RoFormer | `fp32`, `autocast`, or `native_fp16` | Supported with all three precision modes | +| CPU | MelBand RoFormer, BS-RoFormer | `fp32`, or `autocast` when available | Supported with both precision modes | +| MPS, CUDA, or CPU | VR, Demucs, and other PyTorch models | `fp32`, or `autocast` when available | Not yet verified; a request warns and continues in eager mode | +| DirectML | PyTorch models | `fp32` | Not enabled for these optimizations; requests warn and continue with float32/eager inference | +| Any | ONNX models | Managed by the selected ONNX Runtime provider | Not applicable | + +Native float16 is currently verified only for MelBand RoFormer and BS-RoFormer on MPS and CUDA. Requesting it for any other combination logs a warning and safely continues in float32. Native float16 keeps numerically sensitive RoFormer operations, including rotary angles, normalization, STFT, and ISTFT, in float32. + +Regional compilation requires PyTorch 2.6 or newer. Older supported PyTorch releases keep the selected precision mode, warn, and continue with eager inference. + +Compilation has a cold-start cost: the first separation for a new model or input shape can be slower while PyTorch builds and caches graphs. Verified MPS and CUDA measurements showed that autocast plus compilation can improve warm, repeated same-shape MelBand and BS-RoFormer inference, so this is a useful starting point for that workload: + +```sh +audio-separator path/to/audio.wav --use_autocast --use_torch_compile +``` + +Results depend on the model, input shape, compiler-cache state, PyTorch version, and hardware; for one-shot or changing-shape workloads, benchmark eager inference as well. Float32 compilation can be selected explicitly on CPU, MPS, or CUDA: + +```sh +audio-separator path/to/audio.wav --use_torch_compile +``` + +CPU compilation also requires a PyTorch and Python combination supported by Torch Dynamo. If compilation is unavailable or fails at any point, audio-separator warns, falls back to eager inference, and reports `effective_torch_compile=False`. + +Native-float16 compilation is available only for verified MelBand RoFormer and BS-RoFormer models on MPS or CUDA: + +```sh +audio-separator path/to/audio.wav --use_native_fp16 --use_torch_compile +``` + +After `load_model()`, the Python API exposes the selected mode through `separator.effective_precision` (`"fp32"`, `"autocast"`, or `"native_fp16"`) and `separator.effective_torch_compile` (`True` only when regional compilation was activated). These properties make warning-based fallbacks observable to callers. When a multi-model ensemble is selected, they return `"fp32"` and `False` until an individual member is loaded, because the ensemble itself has no single effective execution mode. + ### ๐ŸŽฅ FFmpeg dependency ๐Ÿ’ฌ To test if `audio-separator` has been successfully configured to use FFmpeg, run `audio-separator --env_info`. The log will show `FFmpeg installed`. @@ -450,7 +520,7 @@ Presets are defined in `audio_separator/ensemble_presets.json` โ€” contributions ```sh usage: audio-separator [-h] [-v] [-d] [-e] [-l] [--log_level LOG_LEVEL] [--list_filter LIST_FILTER] [--list_limit LIST_LIMIT] [--list_format {pretty,json}] [-m MODEL_FILENAME] [--output_format OUTPUT_FORMAT] [--output_bitrate OUTPUT_BITRATE] [--output_dir OUTPUT_DIR] [--model_file_dir MODEL_FILE_DIR] [--download_model_only] [--invert_spect] [--normalization NORMALIZATION] - [--amplification AMPLIFICATION] [--single_stem SINGLE_STEM] [--sample_rate SAMPLE_RATE] [--use_soundfile] [--use_autocast] [--use_directml] [--custom_output_names CUSTOM_OUTPUT_NAMES] + [--amplification AMPLIFICATION] [--single_stem SINGLE_STEM] [--sample_rate SAMPLE_RATE] [--use_soundfile] [--use_autocast | --use_native_fp16] [--use_torch_compile] [--use_directml] [--custom_output_names CUSTOM_OUTPUT_NAMES] [--mdx_segment_size MDX_SEGMENT_SIZE] [--mdx_overlap MDX_OVERLAP] [--mdx_batch_size MDX_BATCH_SIZE] [--mdx_hop_length MDX_HOP_LENGTH] [--mdx_enable_denoise] [--vr_batch_size VR_BATCH_SIZE] [--vr_window_size VR_WINDOW_SIZE] [--vr_aggression VR_AGGRESSION] [--vr_enable_tta] [--vr_high_end_process] [--vr_enable_post_process] [--vr_post_process_threshold VR_POST_PROCESS_THRESHOLD] [--demucs_segment_size DEMUCS_SEGMENT_SIZE] [--demucs_shifts DEMUCS_SHIFTS] [--demucs_overlap DEMUCS_OVERLAP] @@ -491,7 +561,9 @@ Common Separation Parameters: --single_stem SINGLE_STEM Output only single stem, e.g. Instrumental, Vocals, Drums, Bass, Guitar, Piano, Other. Example: --single_stem=Instrumental --sample_rate SAMPLE_RATE Modify the sample rate of the output audio (default: 44100). Example: --sample_rate=44100 --use_soundfile Use soundfile to write audio output (default: False). Example: --use_soundfile - --use_autocast Use PyTorch autocast for faster inference (default: False). Do not use for CPU inference. Example: --use_autocast + --use_autocast Use PyTorch autocast when supported (default: False). Example: --use_autocast + --use_native_fp16 Use native float16 for verified model/device combinations (default: False). Mutually exclusive with --use_autocast. Example: --use_native_fp16 + --use_torch_compile Compile verified repeated model blocks when supported (default: False). Best for long inputs or repeated same-shape runs; a fresh compiler cache can make the first run slower. Example: --use_torch_compile --use_directml Use DirectML for hardware-accelerated inference on Windows AMD/Intel GPUs (experimental; requires the 'dml' extra). Example: --use_directml --custom_output_names CUSTOM_OUTPUT_NAMES Custom names for all output files in JSON format (default: None). Example: --custom_output_names='{"Vocals": "vocals_output", "Drums": "drums_output"}' @@ -550,6 +622,10 @@ You can process multiple files without reloading the model to save time and memo You only need to load a model when choosing or changing models. See example below: +Consecutive calls to `load_model()` with the same single model filename reuse the loaded instance. Call `load_model(..., force_reload=True)` after changing settings that are captured when the model is loaded. Multi-model ensembles keep their existing loading behavior. + +The reused model and its weights stay in memory until another model replaces them or the `Separator` is released. Demucs is the exception: it still loads and releases its internal network inside each `separate()` call. + ```python from audio_separator.separator import Separator @@ -648,7 +724,9 @@ You can also rename specific stems: - **`invert_using_spec`:** (Optional) Flag to invert using spectrogram. `Default: False` - **`sample_rate`:** (Optional) Set the sample rate of the output audio. `Default: 44100` - **`use_soundfile`:** (Optional) Use soundfile for output writing, can solve OOM issues, especially on longer audio. -- **`use_autocast`:** (Optional) Flag to use PyTorch autocast for faster inference. Do not use for CPU inference. `Default: False` +- **`use_autocast`:** (Optional) Use PyTorch autocast when the loaded model and device support it. Mutually exclusive with `use_native_fp16=True`. `Default: False` +- **`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}` - **`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}` @@ -658,6 +736,8 @@ You can also rename specific stems: - **`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` +After loading a model, inspect the read-only `Separator.effective_precision` and `Separator.effective_torch_compile` properties to see which requested modes were actually activated. + ## Remote API Usage ๐ŸŒ Audio Separator includes a remote API client that allows you to connect to a deployed Audio Separator API service, enabling you to perform audio separation without running the models locally. The API uses asynchronous processing with job polling for efficient handling of separation tasks. @@ -666,7 +746,7 @@ To deploy Audio Separator as an API on modal.com and use this for remote process ## Requirements ๐Ÿ“‹ -Python >= 3.10 +Python >= 3.10 (Python 3.14.1 is not supported) Libraries: torch, onnx, onnxruntime, numpy, librosa, requests, six, tqdm, pydub @@ -676,9 +756,12 @@ This project uses Poetry for dependency management and packaging. Follow these s ### Prerequisites -- Make sure you have Python 3.10 or newer installed on your machine. +- Make sure you have Python 3.10 or newer installed on your machine, excluding Python 3.14.1. +- Install Poetry 2.0.0 or newer. Poetry 2 is required for dependency resolution, installation, and builds. - Install Conda (I recommend Miniforge: [Miniforge GitHub](https://github.com/conda-forge/miniforge)) to manage your Python virtual environments +The contributor lock currently resolves PyTorch 2.13 with the CUDA 13.0 stack on Linux. [CUDA 13 requires an R580-or-newer NVIDIA driver](https://docs.nvidia.com/deploy/cuda-compatibility/minor-version-compatibility.html), so check `nvidia-smi` before using a self-hosted GPU runner. This applies to the contributor lock only, not to pip installations. + ### Clone the Repository Clone the repository to your local machine: @@ -785,4 +868,4 @@ For questions or feedback, please raise an issue or reach out to @beveradb ([And - \ No newline at end of file + diff --git a/audio_separator/separator/architectures/demucs_separator.py b/audio_separator/separator/architectures/demucs_separator.py index d1d62dc..7d618ca 100644 --- a/audio_separator/separator/architectures/demucs_separator.py +++ b/audio_separator/separator/architectures/demucs_separator.py @@ -8,6 +8,7 @@ from audio_separator.separator.uvr_lib_v5.demucs.hdemucs import HDemucs from audio_separator.separator.uvr_lib_v5.demucs.pretrained import get_model as get_demucs_model from audio_separator.separator.uvr_lib_v5 import spec_utils +from audio_separator.separator.uvr_lib_v5.device_utils import mps_accumulation_budget_bytes, should_accumulate_on_device DEMUCS_4_SOURCE = ["drums", "bass", "other", "vocals"] @@ -23,6 +24,17 @@ } +def _estimate_demucs_full_track_buffer_bytes(channels, samples, num_sources, shifts, num_bag_models): + """Estimate peak float32 storage retained across Demucs split and shift passes.""" + input_bytes = channels * samples * 4 + output_copies = 2 if shifts > 1 else 1 + if num_bag_models: + output_copies = max(output_copies, 2) + if num_bag_models > 1 and shifts > 1: + output_copies = 3 + return 3 * input_bytes + output_copies * num_sources * input_bytes + samples * 4 + + class DemucsSeparator(CommonSeparator): """ DemucsSeparator is responsible for separating audio sources using Demucs models. @@ -107,19 +119,34 @@ def separate(self, audio_file_path, custom_output_names=None): self.logger.debug("Loading model for demixing...") - self.demucs_model_instance = HDemucs(sources=DEMUCS_4_SOURCE) - self.demucs_model_instance = get_demucs_model(name=os.path.splitext(os.path.basename(self.model_path))[0], repo=Path(os.path.dirname(self.model_path))) - self.demucs_model_instance = demucs_segments(self.segment_size, self.demucs_model_instance) - self.demucs_model_instance.to(self.torch_device) - self.demucs_model_instance.eval() - - self.logger.debug("Model loaded and set to evaluation mode.") - - source = self.demix_demucs(mix) - - del self.demucs_model_instance - self.clear_gpu_cache() - self.logger.debug("Model and GPU cache cleared after demixing.") + separation_failed = False + try: + self.demucs_model_instance = HDemucs(sources=DEMUCS_4_SOURCE) + self.demucs_model_instance = get_demucs_model(name=os.path.splitext(os.path.basename(self.model_path))[0], repo=Path(os.path.dirname(self.model_path))) + self.demucs_model_instance = demucs_segments(self.segment_size, self.demucs_model_instance) + self.demucs_model_instance.to(self.torch_device) + self.demucs_model_instance.eval() + + self.logger.debug("Model loaded and set to evaluation mode.") + + source = self.demix_demucs(mix) + except BaseException: + separation_failed = True + raise + finally: + if hasattr(self, "demucs_model_instance"): + del self.demucs_model_instance + try: + self.clear_gpu_cache() + except Exception as cleanup_error: + if not separation_failed: + raise + self.logger.warning( + "Failed to clear the GPU cache after Demucs inference failed: %s", + cleanup_error, + exc_info=True, + ) + self.logger.debug("Model and GPU cache cleared after demixing.") output_files = [] self.logger.debug("Processing output files...") @@ -166,9 +193,28 @@ def demix_demucs(self, mix): self.logger.debug("Starting demixing process in demix_demucs...") processed = {} - mix = torch.tensor(mix, dtype=torch.float32) + num_sources = len(self.demucs_model_instance.sources) + estimated_buffer_bytes = _estimate_demucs_full_track_buffer_bytes( + channels=mix.shape[0], + samples=mix.shape[-1], + num_sources=num_sources, + shifts=self.shifts, + num_bag_models=len(getattr(self.demucs_model_instance, "models", ())), + ) + accumulate_on_device = should_accumulate_on_device(self.torch_device, estimated_buffer_bytes) + mix_device = self.torch_device if accumulate_on_device else torch.device("cpu") + if self.torch_device.type == "mps" and not accumulate_on_device: + self.logger.info( + "Keeping the full-track Demucs mix and source buffers on CPU for this input to limit MPS memory use; " + "model inference still runs on MPS " + f"(estimated buffers: {estimated_buffer_bytes / 1024**3:.2f} GiB, " + f"budget: {mps_accumulation_budget_bytes() / 1024**3:.2f} GiB)." + ) + mix = torch.tensor(mix, dtype=torch.float32, device=mix_device) ref = mix.mean(0) - mix = (mix - ref.mean()) / ref.std() + ref_mean = ref.mean() + ref_std = ref.std() + mix.sub_(ref_mean).div_(ref_std) mix_infer = mix with torch.no_grad(): @@ -185,7 +231,8 @@ def demix_demucs(self, mix): progress=True, )[0] - sources = (sources * ref.std() + ref.mean()).cpu().numpy() + sources.mul_(ref_std).add_(ref_mean) + sources = sources.cpu().numpy() sources[[0, 1]] = sources[[1, 0]] processed[mix] = sources[:, :, 0:None].copy() sources = list(processed.values()) diff --git a/audio_separator/separator/architectures/mdx_separator.py b/audio_separator/separator/architectures/mdx_separator.py index 0516d3b..61af1b9 100644 --- a/audio_separator/separator/architectures/mdx_separator.py +++ b/audio_separator/separator/architectures/mdx_separator.py @@ -113,6 +113,7 @@ def load_model(self): self.logger.debug("Loading ONNX model for inference...") if self.segment_size == self.dim_t: + self.uses_pytorch_inference = False ort_session_options = ort.SessionOptions() if self.log_level > 10: ort_session_options.log_severity_level = 3 diff --git a/audio_separator/separator/architectures/mdxc_separator.py b/audio_separator/separator/architectures/mdxc_separator.py index 702d0a4..3e713f9 100644 --- a/audio_separator/separator/architectures/mdxc_separator.py +++ b/audio_separator/separator/architectures/mdxc_separator.py @@ -4,16 +4,30 @@ import torch import numpy as np +from rotary_embedding_torch import RotaryEmbedding from tqdm import tqdm from ml_collections import ConfigDict from scipy import signal from audio_separator.separator.common_separator import CommonSeparator +from audio_separator.separator.execution_policy import NATIVE_FP16 from audio_separator.separator.uvr_lib_v5 import spec_utils +from audio_separator.separator.uvr_lib_v5.device_utils import mps_accumulation_budget_bytes, should_accumulate_on_device, supports_autocast from audio_separator.separator.uvr_lib_v5.tfc_tdf_v3 import TFC_TDF_net # Roformer direct constructors removed; loading handled via RoformerLoader in CommonSeparator. +def _estimate_roformer_full_track_buffer_bytes(num_instruments, channels, samples, chunk_size): + """Estimate float32 overlap-add results, counters, and the chunk window.""" + return (2 * num_instruments * channels * samples + chunk_size) * 4 + + +def _estimate_mdxc_full_track_buffer_bytes(num_stems, channels, padded_length): + """Estimate float32 padded inputs, concatenation temporaries, and outputs.""" + full_track_copies = 2 + max(num_stems, 1) + return full_track_copies * channels * padded_length * 4 + + def _mdxc_inference_device(torch_device, torch_device_cpu, logger): """Pick the inference device for MDXC-family (incl. RoFormer) models. @@ -140,15 +154,19 @@ def load_model(self): if self.is_roformer: # Use the RoformerLoader exclusively; no legacy fallback self.logger.debug("Loading Roformer model via RoformerLoader...") + load_device = "cpu" if self.torch_device.type == "mps" and self.use_native_fp16 else str(self.torch_device) result = self.roformer_loader.load_model( model_path=self.model_path, config=self.model_data, - device=str(self.torch_device), + device=load_device, ) if getattr(result, "success", False) and getattr(result, "model", None) is not None: self.model_run = result.model + self.roformer_model_type = getattr(result, "model_info", {}).get("model_type") + self._configure_model_precision() self.model_run.to(self.torch_device).eval() + self._configure_model_compilation() else: error_msg = getattr(result, "error_message", "RoformerLoader unsuccessful") self.logger.error(f"Failed to load Roformer model: {error_msg}") @@ -169,6 +187,106 @@ def load_model(self): self.logger.error(f"Please try deleting the model file from {self.model_path} and run audio-separator again to re-download it.") sys.exit(1) + def _configure_model_precision(self): + """Apply the resolved precision policy for the loaded RoFormer model.""" + model_family = self.roformer_model_type + if model_family is None: + class_name = self.model_run.__class__.__name__ + model_family = { + "MelBandRoformer": "mel_band_roformer", + "BSRoformer": "bs_roformer", + }.get(class_name, "roformer") + + self.resolve_execution_policy(model_family) + self.is_native_fp16 = self.effective_precision == NATIVE_FP16 + if not self.is_native_fp16: + return + + # Preserve full-precision rotary angles before converting the rest of the model. + rotary_frequencies = [ + (module, module.freqs.detach().float().clone()) + for module in self.model_run.modules() + if isinstance(module, RotaryEmbedding) + ] + self.model_run.half() + for rotary_embedding, frequencies in rotary_frequencies: + rotary_embedding.freqs.data = frequencies.to(rotary_embedding.freqs.device) + rotary_embedding.cached_freqs = None + + self.logger.info("Using native float16 for %s on %s.", model_family, self.torch_device.type) + + def _configure_model_compilation(self): + """Compile repeated RoFormer transformer blocks when the resolved policy allows it.""" + self.is_torch_compiled = False + self.effective_torch_compile = False + if not self._should_torch_compile: + return + + # Resolve this once before Dynamo traces the shared rotary helper. + supports_autocast(self.torch_device) + + transformer_blocks = self._regional_compile_targets() + if not all(callable(getattr(transformer, "compile", None)) for transformer in transformer_blocks): + self.logger.warning("Skipping regional torch.compile: this PyTorch build does not provide Module.compile().") + return + if not all(hasattr(transformer, "_compiled_call_impl") for transformer in transformer_blocks): + self.logger.warning( + "Skipping regional torch.compile: this PyTorch build cannot safely restore eager module calls." + ) + return + + self._regional_compile_original_calls = [ + (transformer, transformer._compiled_call_impl) for transformer in transformer_blocks + ] + + try: + for transformer in transformer_blocks: + transformer.compile() + except Exception as exc: + self._disable_model_compilation() + self.logger.warning(f"Regional torch.compile could not be enabled; continuing with eager inference: {exc}") + return + + self.is_torch_compiled = True + self.effective_torch_compile = True + model_family = getattr(self, "roformer_model_type", self.model_run.__class__.__name__) + device_type = getattr(getattr(self, "torch_device", None), "type", "unknown") + self.logger.info("Using regional torch.compile for %s on %s.", model_family, device_type) + + def _regional_compile_targets(self): + """Return the repeated transformer blocks used by regional compilation.""" + return [transformer for layer in self.model_run.layers for transformer in layer] + + def _disable_model_compilation(self): + """Restore regional compile targets to their pre-compilation call implementations.""" + original_calls = getattr(self, "_regional_compile_original_calls", None) + if original_calls is None: + original_calls = [ + (transformer, None) + for transformer in self._regional_compile_targets() + if hasattr(transformer, "_compiled_call_impl") + ] + for transformer, original_call in original_calls: + transformer._compiled_call_impl = original_call + self.is_torch_compiled = False + self.effective_torch_compile = False + + def _run_roformer_model(self, part): + """Run one RoFormer chunk and restore pre-compilation calls after a lazy compile failure.""" + try: + return self.model_run(part.unsqueeze(0))[0] + except Exception as exc: + if not getattr(self, "is_torch_compiled", False): + raise + + self._disable_model_compilation() + output = self.model_run(part.unsqueeze(0))[0] + self.logger.warning( + "Regional torch.compile failed during inference; retried this chunk successfully using the " + f"pre-compilation module calls: {exc}" + ) + 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. @@ -190,19 +308,15 @@ def separate(self, audio_file_path, custom_output_names=None): self.logger.debug(f"Preparing mix for input audio file {self.audio_file_path}...") mix = self.prepare_mix(self.audio_file_path) - # Check if audio is shorter than threshold + # Short inputs need the configured segment size, but this automatic + # override must not persist when the separator instance is reused. audio_duration_seconds = mix.shape[1] / self.sample_rate - if audio_duration_seconds < 10.0: - # Only change and warn if it wasn't already set by the user - if not self.override_model_segment_size: - self.override_model_segment_size = True - self.logger.warning(f"Audio duration ({audio_duration_seconds:.2f}s) is less than 10 seconds.") - self.logger.warning("Automatically enabling override_model_segment_size for better processing of short audio.") + override_model_segment_size = self._use_model_segment_override(audio_duration_seconds) self.logger.debug("Normalizing mix before demixing...") mix = spec_utils.normalize(wave=mix, max_peak=self.normalization_threshold, min_peak=self.amplification_threshold) - source = self.demix(mix=mix) + source = self.demix(mix=mix, override_model_segment_size=override_model_segment_size) self.logger.debug("Demixing completed.") output_files = [] @@ -281,6 +395,14 @@ def separate(self, audio_file_path, custom_output_names=None): return output_files + def _use_model_segment_override(self, audio_duration_seconds: float) -> bool: + """Resolve the segment-size override for one input without mutating the separator.""" + is_short_audio = audio_duration_seconds < 10.0 + if is_short_audio and not self.override_model_segment_size: + self.logger.warning(f"Audio duration ({audio_duration_seconds:.2f}s) is less than 10 seconds.") + self.logger.warning("Automatically enabling override_model_segment_size for better processing of short audio.") + return self.override_model_segment_size or is_short_audio + def pitch_fix(self, source, sr_pitched, orig_mix): """ Change the pitch of the source audio by a number of semitones. @@ -308,16 +430,39 @@ def overlap_add(self, result, x, weights, start, length): result[..., start : start + safe_len] += x[..., :safe_len] * weights[:safe_len] return result - def demix(self, mix: np.ndarray) -> dict: + @staticmethod + def _roformer_chunk_starts(audio_length: int, chunk_size: int, step: int) -> list[int]: + """Return a chunk schedule that covers the input without repeating the tail chunk.""" + if audio_length < 0: + raise ValueError("audio_length must be greater than or equal to 0.") + if chunk_size <= 0: + raise ValueError("chunk_size must be greater than 0.") + if step <= 0 or step > chunk_size: + raise ValueError("step must be greater than 0 and less than or equal to chunk_size.") + + starts = [] + for offset in range(0, audio_length, step): + if offset + chunk_size >= audio_length: + tail_start = max(audio_length - chunk_size, 0) + if not starts or starts[-1] != tail_start: + starts.append(tail_start) + break + starts.append(offset) + return starts + + def demix(self, mix: np.ndarray, override_model_segment_size: bool | None = None) -> dict: """ Demixes the input mix into primary and secondary sources using the model and model data. Args: mix (np.ndarray): The mix to be demixed. + override_model_segment_size (bool | None): Segment-size override for this input. Returns: dict: A dictionary containing the demixed sources. """ orig_mix = mix + if override_model_segment_size is None: + override_model_segment_size = self.override_model_segment_size if self.pitch_shift != 0: self.logger.debug(f"Shifting pitch by -{self.pitch_shift} semitones...") @@ -328,7 +473,7 @@ def demix(self, mix: np.ndarray) -> dict: mix = torch.tensor(mix, dtype=torch.float32) - if self.override_model_segment_size: + if override_model_segment_size: mdx_segment_size = self.segment_size self.logger.debug(f"Using configured segment size: {mdx_segment_size}") else: @@ -360,53 +505,58 @@ def demix(self, mix: np.ndarray) -> dict: step = chunk_size if desired_step <= 0 else min(desired_step, chunk_size) self.logger.debug(f"Step: {step} (desired={desired_step})") - # Create a weighting table and convert it to a PyTorch tensor - window = torch.tensor(signal.windows.hamming(chunk_size), dtype=torch.float32) - device = next(self.model_run.parameters()).device + req_shape = (len(self.model_data_cfgdict.training.instruments),) + tuple(mix.shape) + estimated_buffer_bytes = _estimate_roformer_full_track_buffer_bytes( + num_instruments=req_shape[0], + channels=req_shape[1], + samples=req_shape[2], + chunk_size=chunk_size, + ) + accumulate_on_device = should_accumulate_on_device(device, estimated_buffer_bytes) + accumulation_device = device if accumulate_on_device else torch.device("cpu") + if device.type == "mps" and not accumulate_on_device: + self.logger.info( + "Keeping the overlap-add result/counter buffers on CPU for this input to limit MPS memory use; " + "model inference still runs on MPS " + f"(estimated full-track buffers: {estimated_buffer_bytes / 1024**3:.2f} GiB, " + f"budget: {mps_accumulation_budget_bytes() / 1024**3:.2f} GiB)." + ) + # Keep overlap-add buffers next to the model on unified-memory MPS devices. + window = torch.tensor(signal.windows.hamming(chunk_size), dtype=torch.float32, device=accumulation_device) with torch.no_grad(): - req_shape = (len(self.model_data_cfgdict.training.instruments),) + tuple(mix.shape) - result = torch.zeros(req_shape, dtype=torch.float32) - counter = torch.zeros(req_shape, dtype=torch.float32) + result = torch.zeros(req_shape, dtype=torch.float32, device=accumulation_device) + counter = torch.zeros(req_shape, dtype=torch.float32, device=accumulation_device) - for i in tqdm(range(0, mix.shape[1], step)): - part = mix[:, i : i + chunk_size] + chunk_starts = self._roformer_chunk_starts(mix.shape[1], chunk_size, step) + for start_idx in tqdm(chunk_starts): + part = mix[:, start_idx : start_idx + chunk_size] length = part.shape[-1] - if i + chunk_size > mix.shape[1]: - part = mix[:, -chunk_size:] - length = chunk_size part = part.to(device) - x = self.model_run(part.unsqueeze(0))[0] - x = x.cpu() + x = self._run_roformer_model(part) + if x.device != accumulation_device: + x = x.to(accumulation_device) _release_dml_memory_if_needed(device) - # Perform overlap_add on CPU - if i + chunk_size > mix.shape[1]: - # Fixed to correctly add to the end of the tensor - start_idx = result.shape[-1] - chunk_size - result = self.overlap_add(result, x, window, start_idx, length) - safe_len = min(length, x.shape[-1], window.shape[0]) - if safe_len > 0: - counter[..., start_idx : start_idx + safe_len] += window[:safe_len] - else: - result = self.overlap_add(result, x, window, i, length) - safe_len = min(length, x.shape[-1], window.shape[0]) - if safe_len > 0: - counter[..., i : i + safe_len] += window[:safe_len] - - inferenced_outputs = result / counter.clamp(min=1e-10) + result = self.overlap_add(result, x, window, start_idx, length) + safe_len = min(length, x.shape[-1], window.shape[0]) + if safe_len > 0: + counter[..., start_idx : start_idx + safe_len] += window[:safe_len] - else: - mix = torch.tensor(mix, dtype=torch.float32) + counter.clamp_(min=1e-10) + result.div_(counter) + inferenced_outputs = result + del counter + else: try: num_stems = self.model_run.num_target_instruments except AttributeError: num_stems = self.model_run.module.num_target_instruments self.logger.debug(f"Number of stems: {num_stems}") - if self.override_model_segment_size: + if override_model_segment_size: mdx_segment_size = self.segment_size self.logger.debug(f"Using configured segment size: {mdx_segment_size}") else: @@ -423,7 +573,32 @@ def demix(self, mix: np.ndarray) -> dict: pad_size = hop_size - (mix_shape - chunk_size) % hop_size self.logger.debug(f"Pad size: {pad_size}") - mix = torch.cat([torch.zeros(2, chunk_size - hop_size), mix, torch.zeros(2, pad_size + chunk_size - hop_size)], 1) + padded_length = mix_shape + pad_size + 2 * (chunk_size - hop_size) + estimated_buffer_bytes = _estimate_mdxc_full_track_buffer_bytes( + num_stems=num_stems, + channels=mix.shape[0], + padded_length=padded_length, + ) + accumulate_on_device = should_accumulate_on_device(self.torch_device, estimated_buffer_bytes) + accumulation_device = self.torch_device if accumulate_on_device else torch.device("cpu") + if self.torch_device.type == "mps" and not accumulate_on_device: + self.logger.info( + "Keeping the padded mix and accumulated_outputs on CPU for this input to limit MPS memory use; " + "model inference still runs on MPS " + f"(estimated full-track buffers: {estimated_buffer_bytes / 1024**3:.2f} GiB, " + f"budget: {mps_accumulation_budget_bytes() / 1024**3:.2f} GiB)." + ) + + mix = torch.tensor(mix, dtype=torch.float32, device=accumulation_device) + + mix = torch.cat( + [ + torch.zeros(2, chunk_size - hop_size, device=accumulation_device), + mix, + torch.zeros(2, pad_size + chunk_size - hop_size, device=accumulation_device), + ], + 1, + ) self.logger.debug(f"Mix shape: {mix.shape}") chunks = mix.unfold(1, chunk_size, hop_size).transpose(0, 1) @@ -435,7 +610,9 @@ def demix(self, mix: np.ndarray) -> dict: # accumulated_outputs is used to accumulate the output from processing each batch of chunks through the model. # It starts as a tensor of zeros and is updated in-place as the model processes each batch. # The variable holds the combined result of all processed batches, which, after post-processing, represents the separated audio sources. - accumulated_outputs = torch.zeros(num_stems, *mix.shape) if num_stems > 1 else torch.zeros_like(mix) + accumulated_outputs = ( + torch.zeros(num_stems, *mix.shape, device=accumulation_device) if num_stems > 1 else torch.zeros_like(mix) + ) with torch.no_grad(): count = 0 @@ -448,17 +625,18 @@ def demix(self, mix: np.ndarray) -> dict: # Since single_batch_result can contain multiple output tensors (one for each piece of audio in the batch), # individual_output is used to iterate through these tensors and accumulate them into accumulated_outputs. for individual_output in single_batch_result: - individual_output_cpu = individual_output.cpu() - # Accumulate outputs on CPU - accumulated_outputs[..., count * hop_size : count * hop_size + chunk_size] += individual_output_cpu + if individual_output.device != accumulation_device: + individual_output = individual_output.to(accumulation_device) + accumulated_outputs[..., count * hop_size : count * hop_size + chunk_size] += individual_output count += 1 del single_batch_result _release_dml_memory_if_needed(self.torch_device) self.logger.debug("Calculating inferenced outputs based on accumulated outputs and overlap") - inferenced_outputs = accumulated_outputs[..., chunk_size - hop_size : -(pad_size + chunk_size - hop_size)] / self.overlap - self.logger.debug("Deleting accumulated outputs to free up memory") + accumulated_outputs.div_(self.overlap) + inferenced_outputs = accumulated_outputs[..., chunk_size - hop_size : -(pad_size + chunk_size - hop_size)] + self.logger.debug("Releasing the local accumulator reference after selecting the output view") del accumulated_outputs if num_stems > 1: diff --git a/audio_separator/separator/architectures/vr_separator.py b/audio_separator/separator/architectures/vr_separator.py index d00887c..3c5c83e 100644 --- a/audio_separator/separator/architectures/vr_separator.py +++ b/audio_separator/separator/architectures/vr_separator.py @@ -110,6 +110,34 @@ def __init__(self, common_config, arch_config: dict): self.logger.info("VR Separator initialisation complete") + def _ensure_model_loaded(self, nn_arch_size): + """Load VR weights once and retain them for subsequent separations.""" + if isinstance(self.model_run, torch.nn.Module): + self.logger.debug("Reusing the loaded VR model.") + return + + vr_5_1_models = [56817, 218409] + is_vr_51_model = nn_arch_size in vr_5_1_models or self.is_vr_51_model + if is_vr_51_model: + self.logger.debug("Using CascadedNet for VR 5.1 model...") + model_run = nets_new.CascadedNet( + self.model_params.param["bins"] * 2, + nn_arch_size, + nout=self.model_capacity[0], + nout_lstm=self.model_capacity[1], + ) + else: + self.logger.debug("Determining model capacity...") + model_run = nets.determine_model_capacity(self.model_params.param["bins"] * 2, nn_arch_size) + + # Publish the reusable module only after every loading step succeeds. + # Otherwise a retry could mistake a partial model for a ready one. + model_run.load_state_dict(torch.load(self.model_path, map_location="cpu")) + model_run.to(self.torch_device) + self.model_run = model_run + self.is_vr_51_model = is_vr_51_model + self.logger.debug("Model loaded and moved to device.") + 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,22 +185,11 @@ def separate(self, audio_file_path, custom_output_names=None): self.logger.debug(f"Starting separation for input audio file {self.audio_file_path}...") nn_arch_sizes = [31191, 33966, 56817, 123821, 123812, 129605, 218409, 537238, 537227] # default - vr_5_1_models = [56817, 218409] model_size = math.ceil(os.stat(self.model_path).st_size / 1024) nn_arch_size = min(nn_arch_sizes, key=lambda x: abs(x - model_size)) self.logger.debug(f"Model size determined: {model_size}, NN architecture size: {nn_arch_size}") - if nn_arch_size in vr_5_1_models or self.is_vr_51_model: - self.logger.debug("Using CascadedNet for VR 5.1 model...") - self.model_run = nets_new.CascadedNet(self.model_params.param["bins"] * 2, nn_arch_size, nout=self.model_capacity[0], nout_lstm=self.model_capacity[1]) - self.is_vr_51_model = True - else: - self.logger.debug("Determining model capacity...") - self.model_run = nets.determine_model_capacity(self.model_params.param["bins"] * 2, nn_arch_size) - - self.model_run.load_state_dict(torch.load(self.model_path, map_location="cpu")) - self.model_run.to(self.torch_device) - self.logger.debug("Model loaded and moved to device.") + self._ensure_model_loaded(nn_arch_size) y_spec, v_spec = self.inference_vr(self.loading_mix(), self.torch_device, self.aggressiveness) self.logger.debug("Inference completed.") diff --git a/audio_separator/separator/common_separator.py b/audio_separator/separator/common_separator.py index 34435ea..408d36e 100644 --- a/audio_separator/separator/common_separator.py +++ b/audio_separator/separator/common_separator.py @@ -10,6 +10,7 @@ from pydub import AudioSegment import soundfile as sf from audio_separator.separator.uvr_lib_v5 import spec_utils +from audio_separator.separator.execution_policy import FP32, resolve_execution_policy class CommonSeparator: @@ -61,6 +62,7 @@ def __init__(self, config): # Inferencing device / acceleration config self.torch_device = config.get("torch_device") + self.requested_torch_device = self.torch_device self.torch_device_cpu = config.get("torch_device_cpu") self.torch_device_mps = config.get("torch_device_mps") self.onnx_execution_provider = config.get("onnx_execution_provider") @@ -83,6 +85,15 @@ def __init__(self, config): self.invert_using_spec = config.get("invert_using_spec") self.sample_rate = config.get("sample_rate") self.use_soundfile = config.get("use_soundfile") + self.use_autocast = config.get("use_autocast", False) + self.use_native_fp16 = config.get("use_native_fp16", False) + self.use_torch_compile = config.get("use_torch_compile", False) + self.uses_pytorch_inference = True + self.is_native_fp16 = False + self.effective_precision = FP32 + self.effective_torch_compile = False + self._should_torch_compile = False + self._execution_policy_resolved = False # Roformer-specific loading support self.roformer_loader = None @@ -145,6 +156,24 @@ def __init__(self, config): self.cached_sources_map = {} + def resolve_execution_policy(self, model_family): + """Resolve the requested precision and compilation settings for this model.""" + policy = resolve_execution_policy( + device=self.torch_device, + requested_device=self.requested_torch_device, + model_family=model_family, + use_autocast=self.use_autocast, + use_native_fp16=self.use_native_fp16, + use_torch_compile=self.use_torch_compile, + uses_pytorch_inference=getattr(self, "uses_pytorch_inference", True), + logger=self.logger, + ) + self.effective_precision = policy.precision + self.effective_torch_compile = False + self._should_torch_compile = policy.use_torch_compile + self._execution_policy_resolved = True + return policy + def secondary_stem(self, primary_stem: str): """Determines secondary stem name based on the primary stem name.""" primary_stem = primary_stem if primary_stem else self.NO_STEM diff --git a/audio_separator/separator/execution_policy.py b/audio_separator/separator/execution_policy.py new file mode 100644 index 0000000..7e15ce5 --- /dev/null +++ b/audio_separator/separator/execution_policy.py @@ -0,0 +1,117 @@ +"""Resolve requested inference options against verified execution capabilities.""" + +from dataclasses import dataclass + +from packaging import version +import torch + +from audio_separator.separator.uvr_lib_v5.device_utils import supports_autocast + +FP32 = "fp32" +AUTOCAST = "autocast" +NATIVE_FP16 = "native_fp16" +MIN_TORCH_COMPILE_VERSION = version.parse("2.6") + +# Keep these tables intentionally conservative. A cell should only be added after +# its correctness, numerical quality, and fallback behavior have been verified. +# Speed remains workload-dependent, especially when compilation has a cold start. +NATIVE_FP16_CAPABILITIES = frozenset( + { + ("mps", "mel_band_roformer"), + ("mps", "bs_roformer"), + ("cuda", "mel_band_roformer"), + ("cuda", "bs_roformer"), + } +) + +TORCH_COMPILE_CAPABILITIES = frozenset( + { + (device, model_family, precision) + for device in ("mps", "cuda") + for model_family in ("mel_band_roformer", "bs_roformer") + for precision in (FP32, AUTOCAST, NATIVE_FP16) + } + | { + ("cpu", model_family, precision) + for model_family in ("mel_band_roformer", "bs_roformer") + for precision in (FP32, AUTOCAST) + } +) + + +@dataclass(frozen=True) +class ExecutionPolicy: + """The effective execution choices for one loaded model.""" + + precision: str = FP32 + use_torch_compile: bool = False + + +def _regional_compile_runtime_supported() -> bool: + """Return whether Dynamo can trace the SDPA context used by RoFormer.""" + torch_version = version.parse(torch.__version__.split("+")[0]) + return hasattr(torch, "compile") and torch_version >= MIN_TORCH_COMPILE_VERSION + + +def resolve_execution_policy( + *, + device, + requested_device=None, + model_family: str, + use_autocast: bool, + use_native_fp16: bool, + use_torch_compile: bool, + logger, + uses_pytorch_inference: bool = True, +) -> ExecutionPolicy: + """Resolve requested options, warning when an unverified path is skipped.""" + if use_autocast and use_native_fp16: + raise ValueError("Autocast and native float16 are mutually exclusive precision modes.") + + device_type = getattr(device, "type", str(device)) + requested_device_type = getattr(requested_device, "type", str(requested_device)) if requested_device is not None else device_type + capability_device_type = requested_device_type if requested_device_type == "privateuseone" else device_type + normalized_family = (model_family or "unknown").lower() + precision = FP32 + + if use_native_fp16: + capability = (capability_device_type, normalized_family) + if capability in NATIVE_FP16_CAPABILITIES: + precision = NATIVE_FP16 + else: + logger.warning( + "Native float16 is not supported for device=%s, model=%s; continuing with float32 inference.", + device_type, + normalized_family, + ) + elif use_autocast: + if not uses_pytorch_inference: + logger.warning("Autocast only applies to PyTorch inference; continuing with the model's native precision.") + # torch-directml exposes its device as privateuseone. PyTorch's generic + # autocast context does not support that backend, so never enter it. + elif requested_device_type == "privateuseone": + logger.warning("Autocast is not supported on DirectML; continuing with float32 inference.") + elif supports_autocast(device): + precision = AUTOCAST + else: + logger.warning("Autocast is not available for device=%s; continuing with float32 inference.", device_type) + + compile_enabled = False + if use_torch_compile: + capability = (capability_device_type, normalized_family, precision) + if capability not in TORCH_COMPILE_CAPABILITIES: + logger.warning( + "Regional torch.compile is not supported for device=%s, model=%s, precision=%s; continuing with eager inference.", + device_type, + normalized_family, + precision, + ) + elif not _regional_compile_runtime_supported(): + logger.warning( + "Regional torch.compile requires PyTorch 2.6 or newer; found %s. Continuing with eager inference.", + torch.__version__, + ) + else: + compile_enabled = True + + return ExecutionPolicy(precision=precision, use_torch_compile=compile_enabled) diff --git a/audio_separator/separator/roformer/configuration_normalizer.py b/audio_separator/separator/roformer/configuration_normalizer.py index 6177f25..af733c8 100644 --- a/audio_separator/separator/roformer/configuration_normalizer.py +++ b/audio_separator/separator/roformer/configuration_normalizer.py @@ -164,7 +164,7 @@ def _normalize_single_value(self, key: str, value: Any, model_type: str) -> Any: # Integer normalization elif key in ['dim', 'depth', 'num_stems', 'time_transformer_depth', - 'freq_transformer_depth', 'dim_head', 'heads', + 'freq_transformer_depth', 'linear_transformer_depth', 'dim_head', 'heads', 'mlp_expansion_factor', 'num_bands', 'sample_rate', 'stft_n_fft', 'stft_hop_length', 'stft_win_length', 'mask_estimator_depth']: diff --git a/audio_separator/separator/roformer/roformer_loader.py b/audio_separator/separator/roformer/roformer_loader.py index 95c62dd..6624ee6 100644 --- a/audio_separator/separator/roformer/roformer_loader.py +++ b/audio_separator/separator/roformer/roformer_loader.py @@ -149,6 +149,7 @@ def _create_bs_roformer(self, config: Dict[str, Any]): 'num_stems': config.get('num_stems', 2), 'time_transformer_depth': config.get('time_transformer_depth', 2), 'freq_transformer_depth': config.get('freq_transformer_depth', 2), + 'linear_transformer_depth': config.get('linear_transformer_depth', 0), 'freqs_per_bands': config['freqs_per_bands'], 'dim_head': config.get('dim_head', 64), 'heads': config.get('heads', 8), @@ -283,6 +284,7 @@ def get_default_configuration(self, model_type: str) -> Dict[str, Any]: 'num_stems': 2, 'time_transformer_depth': 2, 'freq_transformer_depth': 2, + 'linear_transformer_depth': 0, 'freqs_per_bands': (2, 4, 8, 16, 32, 64), 'dim_head': 64, 'heads': 8, diff --git a/audio_separator/separator/separator.py b/audio_separator/separator/separator.py index 4561e0c..c3e2694 100644 --- a/audio_separator/separator/separator.py +++ b/audio_separator/separator/separator.py @@ -24,6 +24,7 @@ import onnxruntime as ort from tqdm import tqdm from audio_separator.separator.ensembler import Ensembler +from audio_separator.separator.execution_policy import AUTOCAST, FP32, NATIVE_FP16 # Mapping of common stem name variations to canonical names for ensemble grouping. STEM_NAME_MAP = { @@ -73,7 +74,9 @@ class Separator: invert_using_spec (bool): Flag to invert using spectrogram. sample_rate (int): The sample rate of the audio. use_soundfile (bool): Use soundfile for audio writing, can solve OOM issues. - use_autocast (bool): Flag to use PyTorch autocast for faster inference. + use_autocast (bool): Use PyTorch autocast when the loaded model and device support it. + use_torch_compile (bool): Compile verified repeated model blocks when supported. + use_native_fp16 (bool): Convert a verified model to native float16 inference when supported. MDX Architecture Specific Attributes: hop_length (int): The hop length for STFT. @@ -130,8 +133,13 @@ def __init__( ensemble_weights=None, ensemble_preset=None, info_only=False, + use_torch_compile=False, + use_native_fp16=False, ): """Initialize the separator.""" + if use_autocast and use_native_fp16: + raise ValueError("use_autocast and use_native_fp16 are mutually exclusive precision modes.") + self.logger = logging.getLogger(__name__) self.logger.setLevel(log_level) self.log_level = log_level @@ -211,6 +219,8 @@ def __init__( self.use_soundfile = use_soundfile self.use_autocast = use_autocast + self.use_native_fp16 = use_native_fp16 + self.use_torch_compile = use_torch_compile self.use_directml = use_directml self.chunk_duration = chunk_duration @@ -249,6 +259,9 @@ def __init__( self.model_instance = None self.model_filename = None self.model_filenames = [] + self._loaded_model_filename = None + self._loaded_model_friendly_name = None + self._loaded_model_is_uvr_vip = False self.model_is_uvr_vip = False self.model_friendly_name = None @@ -256,6 +269,20 @@ def __init__( if not info_only: self.setup_accelerated_inferencing_device() + @property + def effective_precision(self): + """Return the precision mode selected for the currently loaded model.""" + if len(getattr(self, "model_filenames", ())) > 1: + return FP32 + return getattr(self.model_instance, "effective_precision", FP32) + + @property + def effective_torch_compile(self): + """Return whether regional compilation is active for the loaded model.""" + if len(getattr(self, "model_filenames", ())) > 1: + return False + return bool(getattr(self.model_instance, "effective_torch_compile", False)) + VALID_ENSEMBLE_ALGORITHMS = [ "avg_wave", "median_wave", "min_wave", "max_wave", "avg_fft", "median_fft", "min_fft", "max_fft", @@ -835,10 +862,17 @@ def load_model_data_using_hash(self, model_path): return model_data - def load_model(self, model_filename="model_bs_roformer_ep_317_sdr_12.9755.ckpt"): + def load_model(self, model_filename="model_bs_roformer_ep_317_sdr_12.9755.ckpt", force_reload=False): """ This method instantiates the architecture-specific separation class, loading the separation model into memory, downloading it first if necessary. + + Consecutive calls with the same single model reuse the loaded instance. Set + ``force_reload`` to ``True`` after changing settings captured at load time. + + Args: + model_filename (str or list): The model filename, or filenames for an ensemble. + force_reload (bool): Reload a matching single model instead of reusing it. """ # If an ensemble preset was loaded and no explicit model list was provided, use preset models if self._ensemble_preset_models is not None and model_filename == "model_bs_roformer_ep_317_sdr_12.9755.ckpt": @@ -852,89 +886,121 @@ def load_model(self, model_filename="model_bs_roformer_ep_317_sdr_12.9755.ckpt") return model_filename = model_filename[0] - self.model_filename = model_filename - self.model_filenames = [model_filename] + if not force_reload and self.model_instance is not None and self._loaded_model_filename == model_filename: + self.model_filename = model_filename + self.model_filenames = [model_filename] + self.model_friendly_name = self._loaded_model_friendly_name + self.model_is_uvr_vip = self._loaded_model_is_uvr_vip + self.logger.info(f"Model {model_filename} is already loaded; reusing the existing instance.") + return self.logger.info(f"Loading model {model_filename}...") load_model_start_time = time.perf_counter() + selected_model_filename = model_filename + previous_model_friendly_name = self.model_friendly_name + previous_model_is_uvr_vip = self.model_is_uvr_vip - # Setting up the model path - model_filename, model_type, model_friendly_name, model_path, yaml_config_filename = self.download_model_files(model_filename) - model_name = model_filename.split(".")[0] - self.logger.debug(f"Model downloaded, friendly name: {model_friendly_name}, model_path: {model_path}") - - if model_path.lower().endswith(".yaml"): - yaml_config_filename = model_path - - if yaml_config_filename is not None: - model_data = self.load_model_data_from_yaml(yaml_config_filename) - else: - model_data = self.load_model_data_using_hash(model_path) + try: + # Setting up the model path + model_filename, model_type, model_friendly_name, model_path, yaml_config_filename = self.download_model_files(model_filename) + model_is_uvr_vip = self.model_is_uvr_vip + model_name = model_filename.split(".")[0] + self.logger.debug(f"Model downloaded, friendly name: {model_friendly_name}, model_path: {model_path}") - common_params = { - "logger": self.logger, - "log_level": self.log_level, - "torch_device": self.torch_device, - "torch_device_cpu": self.torch_device_cpu, - "torch_device_mps": self.torch_device_mps, - "onnx_execution_provider": self.onnx_execution_provider, - "model_name": model_name, - "model_path": model_path, - "model_data": model_data, - "output_format": self.output_format, - "output_bitrate": self.output_bitrate, - "output_dir": self.output_dir, - "normalization_threshold": self.normalization_threshold, - "amplification_threshold": self.amplification_threshold, - "output_single_stem": self.output_single_stem, - "invert_using_spec": self.invert_using_spec, - "sample_rate": self.sample_rate, - "use_soundfile": self.use_soundfile, - } + if model_path.lower().endswith(".yaml"): + yaml_config_filename = model_path - # Instantiate the appropriate separator class depending on the model type - separator_classes = {"MDX": "mdx_separator.MDXSeparator", "VR": "vr_separator.VRSeparator", "Demucs": "demucs_separator.DemucsSeparator", "MDXC": "mdxc_separator.MDXCSeparator"} - - if model_type not in self.arch_specific_params or model_type not in separator_classes: - # Enhanced error message for Roformer models - if "roformer" in model_filename.lower() or (model_data and model_data.get("is_roformer", False)): - error_msg = (f"Roformer model type not properly configured: {model_type}. " - f"This may indicate a configuration validation failure. " - f"Please check the model file and YAML configuration.") - self.logger.error(error_msg) - raise ValueError(error_msg) + if yaml_config_filename is not None: + model_data = self.load_model_data_from_yaml(yaml_config_filename) else: - raise ValueError(f"Model type not supported (yet): {model_type}") + model_data = self.load_model_data_using_hash(model_path) + + common_params = { + "logger": self.logger, + "log_level": self.log_level, + "torch_device": self.torch_device, + "torch_device_cpu": self.torch_device_cpu, + "torch_device_mps": self.torch_device_mps, + "onnx_execution_provider": self.onnx_execution_provider, + "model_name": model_name, + "model_path": model_path, + "model_data": model_data, + "output_format": self.output_format, + "output_bitrate": self.output_bitrate, + "output_dir": self.output_dir, + "normalization_threshold": self.normalization_threshold, + "amplification_threshold": self.amplification_threshold, + "output_single_stem": self.output_single_stem, + "invert_using_spec": self.invert_using_spec, + "sample_rate": self.sample_rate, + "use_soundfile": self.use_soundfile, + "use_autocast": self.use_autocast, + "use_native_fp16": self.use_native_fp16, + "use_torch_compile": self.use_torch_compile, + } - if model_type == "Demucs" and sys.version_info < (3, 10): - raise Exception("Demucs models require Python version 3.10 or newer.") + # Instantiate the appropriate separator class depending on the model type + separator_classes = {"MDX": "mdx_separator.MDXSeparator", "VR": "vr_separator.VRSeparator", "Demucs": "demucs_separator.DemucsSeparator", "MDXC": "mdxc_separator.MDXCSeparator"} + + if model_type not in self.arch_specific_params or model_type not in separator_classes: + # Enhanced error message for Roformer models + if "roformer" in model_filename.lower() or (model_data and model_data.get("is_roformer", False)): + error_msg = (f"Roformer model type not properly configured: {model_type}. " + f"This may indicate a configuration validation failure. " + f"Please check the model file and YAML configuration.") + self.logger.error(error_msg) + raise ValueError(error_msg) + else: + raise ValueError(f"Model type not supported (yet): {model_type}") - self.logger.debug(f"Importing module for model type {model_type}: {separator_classes[model_type]}") + if model_type == "Demucs" and sys.version_info < (3, 10): + raise Exception("Demucs models require Python version 3.10 or newer.") - module_name, class_name = separator_classes[model_type].split(".") - module = importlib.import_module(f"audio_separator.separator.architectures.{module_name}") - separator_class = getattr(module, class_name) + self.logger.debug(f"Importing module for model type {model_type}: {separator_classes[model_type]}") - self.logger.debug(f"Instantiating separator class for model type {model_type}: {separator_class}") + module_name, class_name = separator_classes[model_type].split(".") + module = importlib.import_module(f"audio_separator.separator.architectures.{module_name}") + separator_class = getattr(module, class_name) - try: - self.model_instance = separator_class(common_config=common_params, arch_config=self.arch_specific_params[model_type]) - except Exception as e: - # Enhanced error handling for Roformer models - if "roformer" in model_filename.lower() or (model_data and model_data.get("is_roformer", False)): - error_msg = (f"Failed to instantiate Roformer model: {e}. " - f"This may be due to missing parameters or configuration validation failures.") - self.logger.error(error_msg) - raise RuntimeError(error_msg) from e - else: - raise + self.logger.debug(f"Instantiating separator class for model type {model_type}: {separator_class}") - # Log Roformer implementation version if applicable - if hasattr(self.model_instance, 'is_roformer_model') and self.model_instance.is_roformer_model: - roformer_stats = self.model_instance.get_roformer_loading_stats() - if roformer_stats: - self.logger.info(f"Roformer loading stats: {roformer_stats}") + try: + model_instance = separator_class(common_config=common_params, arch_config=self.arch_specific_params[model_type]) + except Exception as e: + # Enhanced error handling for Roformer models + if "roformer" in model_filename.lower() or (model_data and model_data.get("is_roformer", False)): + error_msg = (f"Failed to instantiate Roformer model: {e}. " + f"This may be due to missing parameters or configuration validation failures.") + self.logger.error(error_msg) + raise RuntimeError(error_msg) from e + else: + raise + + resolve_policy = getattr(model_instance, "resolve_execution_policy", None) + if callable(resolve_policy) and not getattr(model_instance, "_execution_policy_resolved", False): + resolve_policy(model_type.lower()) + + # Log Roformer implementation version if applicable + if hasattr(model_instance, 'is_roformer_model') and model_instance.is_roformer_model: + roformer_stats = model_instance.get_roformer_loading_stats() + if roformer_stats: + self.logger.info(f"Roformer loading stats: {roformer_stats}") + except BaseException: + # Model discovery updates these fields; restore them with the prior + # selection so callers never observe a partially loaded model. + self.model_friendly_name = previous_model_friendly_name + self.model_is_uvr_vip = previous_model_is_uvr_vip + raise + + self.model_instance = model_instance + self.model_filename = selected_model_filename + self.model_filenames = [selected_model_filename] + self._loaded_model_filename = model_filename + self.model_friendly_name = model_friendly_name + self.model_is_uvr_vip = model_is_uvr_vip + self._loaded_model_friendly_name = model_friendly_name + self._loaded_model_is_uvr_vip = model_is_uvr_vip # Log the completion of the model load process self.logger.debug("Loading model completed.") @@ -1027,21 +1093,45 @@ def _separate_file(self, audio_file_path, custom_output_names=None): self.logger.debug(f"Normalization threshold set to {self.normalization_threshold}, waveform will be lowered to this max amplitude to avoid clipping.") self.logger.debug(f"Amplification threshold set to {self.amplification_threshold}, waveform will be scaled up to this max amplitude if below it.") - # Run separation method for the loaded model with autocast enabled if supported by the device + # Run separation using the policy resolved for the loaded model's actual + # inference device. This matters when an architecture falls back to CPU. output_files = None - if self.use_autocast and autocast_mode.is_autocast_available(self.torch_device.type): - self.logger.debug("Autocast available.") - with autocast_mode.autocast(self.torch_device.type): + effective_precision = self.effective_precision + inference_device = getattr(self.model_instance, "torch_device", self.torch_device) + inference_device_type = getattr(inference_device, "type", str(inference_device)) + separation_failed = False + try: + if effective_precision == NATIVE_FP16: + self.logger.debug("Using native float16 inference.") output_files = self.model_instance.separate(audio_file_path, custom_output_names) - else: - self.logger.debug("Autocast unavailable.") - output_files = self.model_instance.separate(audio_file_path, custom_output_names) - - # Clear GPU cache to free up memory - self.model_instance.clear_gpu_cache() - - # Unset separation parameters to prevent accidentally re-using the wrong source files or output paths - self.model_instance.clear_file_specific_paths() + elif effective_precision == AUTOCAST and inference_device_type != "privateuseone": + self.logger.debug("Using autocast inference on %s.", inference_device_type) + with autocast_mode.autocast(inference_device_type): + output_files = self.model_instance.separate(audio_file_path, custom_output_names) + else: + self.logger.debug("Using float32 inference.") + output_files = self.model_instance.separate(audio_file_path, custom_output_names) + except BaseException: + separation_failed = True + raise + finally: + # Reused instances must not retain per-file state after a failed inference. + cleanup_error = None + for cleanup_name in ("clear_gpu_cache", "clear_file_specific_paths"): + try: + getattr(self.model_instance, cleanup_name)() + except Exception as exc: + if separation_failed: + self.logger.warning( + "Cleanup %s failed after separation raised an error: %s", + cleanup_name, + exc, + exc_info=True, + ) + elif cleanup_error is None: + cleanup_error = exc + if cleanup_error is not None: + raise cleanup_error # Remind the user one more time if they used a VIP model, so the message doesn't get lost in the logs self.print_uvr_vip_message() diff --git a/audio_separator/separator/uvr_lib_v5/demucs/hdemucs.py b/audio_separator/separator/uvr_lib_v5/demucs/hdemucs.py index 9e299d4..be528f9 100644 --- a/audio_separator/separator/uvr_lib_v5/demucs/hdemucs.py +++ b/audio_separator/separator/uvr_lib_v5/demucs/hdemucs.py @@ -16,6 +16,7 @@ from .demucs import DConv, rescale_module from .states import capture_init from .spec import spectro, ispectro +from audio_separator.separator.uvr_lib_v5.device_utils import should_fallback_to_cpu_for_demucs_mask def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = "constant", value: float = 0.0): @@ -757,24 +758,17 @@ def forward(self, mix): x = x.view(B, S, -1, Fq, T) x = x * std[:, None] + mean[:, None] - # to cpu as non-cuda GPUs don't support complex numbers - # demucs issue #435 ##432 - # NOTE: in this case z already is on cpu - # TODO: remove this when mps supports complex numbers - - device_type = x.device.type - device_load = f"{device_type}:{x.device.index}" if not device_type == "mps" else device_type - x_is_other_gpu = not device_type in ["cuda", "cpu"] - - if x_is_other_gpu: + original_device = x.device + should_fallback = should_fallback_to_cpu_for_demucs_mask(original_device, self.cac) + if should_fallback: + z = z.cpu() x = x.cpu() zout = self._mask(z, x) x = self._ispec(zout, length) - # back to other device - if x_is_other_gpu: - x = x.to(device_load) + if should_fallback: + x = x.to(original_device) if self.hybrid: xt = xt.view(B, S, -1, length) diff --git a/audio_separator/separator/uvr_lib_v5/demucs/htdemucs.py b/audio_separator/separator/uvr_lib_v5/demucs/htdemucs.py index f3d7a27..7311469 100644 --- a/audio_separator/separator/uvr_lib_v5/demucs/htdemucs.py +++ b/audio_separator/separator/uvr_lib_v5/demucs/htdemucs.py @@ -22,6 +22,7 @@ from .states import capture_init from .spec import spectro, ispectro from .hdemucs import pad1d, ScaledEmbedding, HEncLayer, MultiWrap, HDecLayer +from audio_separator.separator.uvr_lib_v5.device_utils import should_fallback_to_cpu_for_demucs_mask class HTDemucs(nn.Module): @@ -581,16 +582,10 @@ def forward(self, mix): x = x.view(B, S, -1, Fq, T) x = x * std[:, None] + mean[:, None] - # to cpu as non-cuda GPUs don't support complex numbers - # demucs issue #435 ##432 - # NOTE: in this case z already is on cpu - # TODO: remove this when mps supports complex numbers - - device_type = x.device.type - device_load = f"{device_type}:{x.device.index}" if not device_type == "mps" else device_type - x_is_other_gpu = not device_type in ["cuda", "cpu"] - - if x_is_other_gpu: + original_device = x.device + should_fallback = should_fallback_to_cpu_for_demucs_mask(original_device, self.cac) + if should_fallback: + z = z.cpu() x = x.cpu() zout = self._mask(z, x) @@ -602,9 +597,8 @@ def forward(self, mix): else: x = self._ispec(zout, length) - # back to other device - if x_is_other_gpu: - x = x.to(device_load) + if should_fallback: + x = x.to(original_device) if self.use_train_segment: if self.training: diff --git a/audio_separator/separator/uvr_lib_v5/demucs/spec.py b/audio_separator/separator/uvr_lib_v5/demucs/spec.py index 36a4a2e..a653491 100644 --- a/audio_separator/separator/uvr_lib_v5/demucs/spec.py +++ b/audio_separator/separator/uvr_lib_v5/demucs/spec.py @@ -7,15 +7,14 @@ import torch as th +from audio_separator.separator.uvr_lib_v5.device_utils import should_fallback_to_cpu_for_complex_ops + def spectro(x, n_fft=512, hop_length=None, pad=0): *other, length = x.shape x = x.reshape(-1, length) - device_type = x.device.type - is_other_gpu = not device_type in ["cuda", "cpu"] - - if is_other_gpu: + if should_fallback_to_cpu_for_complex_ops(x.device): x = x.cpu() z = th.stft(x, n_fft * (1 + pad), hop_length or n_fft // 4, window=th.hann_window(n_fft).to(x), win_length=n_fft, normalized=True, center=True, return_complex=True, pad_mode="reflect") _, freqs, frame = z.shape @@ -28,10 +27,7 @@ def ispectro(z, hop_length=None, length=None, pad=0): z = z.view(-1, freqs, frames) win_length = n_fft // (1 + pad) - device_type = z.device.type - is_other_gpu = not device_type in ["cuda", "cpu"] - - if is_other_gpu: + if should_fallback_to_cpu_for_complex_ops(z.device): z = z.cpu() x = th.istft(z, n_fft, hop_length, window=th.hann_window(win_length).to(z.real), win_length=win_length, normalized=True, length=length, center=True) _, length = x.shape diff --git a/audio_separator/separator/uvr_lib_v5/device_utils.py b/audio_separator/separator/uvr_lib_v5/device_utils.py new file mode 100644 index 0000000..03da866 --- /dev/null +++ b/audio_separator/separator/uvr_lib_v5/device_utils.py @@ -0,0 +1,159 @@ +"""Device capability helpers for hardware-specific inference paths.""" + +import os +from contextlib import nullcontext +from functools import lru_cache + +import torch + + +_AUTOCAST_SUPPORT_CACHE = {} + +# Full-track buffers grow linearly with audio duration. Keep short and medium +# inputs on MPS, but leave oversized accumulators on CPU so they cannot consume +# most of the MPS working-set budget before model activations are allocated. +# This value is the floor of the budget, and the whole budget when Metal cannot +# report a working-set size. +MAX_MPS_FULL_TRACK_BUFFER_BYTES = 1024**3 + +# Share of the *free* Metal working set that full-track buffers may occupy. +# Half means the buffers can never take more room than they leave behind for +# model activations, which are what actually fail when the working set runs out. +MPS_BUFFER_HEADROOM_SHARE = 0.5 + +# Overrides the computed budget, in GiB. Intended for diagnosis and for callers +# who know their own headroom better than the heuristic does. +MPS_BUFFER_BUDGET_ENV = "AUDIO_SEPARATOR_MPS_BUFFER_BUDGET_GIB" + + +def _supports_autocast(device_type: str) -> bool: + """Return whether PyTorch can safely enter autocast for a device type.""" + if device_type in _AUTOCAST_SUPPORT_CACHE: + return _AUTOCAST_SUPPORT_CACHE[device_type] + + # PyTorch reports privateuseone as autocast-capable even though + # torch-directml does not register the AMP hooks required by the context. + if device_type == "privateuseone": + _AUTOCAST_SUPPORT_CACHE[device_type] = False + return False + + try: + is_available = getattr(torch.amp.autocast_mode, "is_autocast_available", None) + if is_available is not None and not is_available(device_type): + _AUTOCAST_SUPPORT_CACHE[device_type] = False + return False + with torch.autocast(device_type=device_type, enabled=False): + pass + supported = True + except (AssertionError, RuntimeError, TypeError, ValueError): + supported = False + + _AUTOCAST_SUPPORT_CACHE[device_type] = supported + return supported + + +def supports_autocast(device: torch.device) -> bool: + """Return whether a device can use PyTorch's generic autocast context.""" + return _supports_autocast(device.type) + + +def autocast_disabled(device: torch.device): + """Disable autocast when supported, otherwise return a no-op context.""" + if not supports_autocast(device): + return nullcontext() + return torch.autocast(device_type=device.type, enabled=False) + + +def _probe_complex_scatter_add(spectrum: torch.Tensor) -> None: + """Exercise the complex scatter operation used by MelBand RoFormer.""" + source = spectrum[:, :2, :2] + indices = torch.zeros(source.shape, dtype=torch.long, device=source.device) + torch.zeros_like(source).scatter_add_(1, indices, source) + + +@lru_cache(maxsize=32) +def _supports_complex_spectral_ops(device_type: str, device_index: int) -> bool: + """Return whether a device can execute the complex operations used by the models.""" + if device_type in {"cpu", "cuda"}: + return True + + # DirectML cannot represent complex tensors. Avoid probing unsupported + # operations on its out-of-tree backend slot. + if device_type == "privateuseone": + return False + + try: + device = torch.device(f"{device_type}:{device_index}") if device_index >= 0 else torch.device(device_type) + sample_length = 1024 + n_fft = 256 + hop_length = 64 + sample = torch.randn(1, sample_length, device=device) + window = torch.hann_window(n_fft, device=device) + spectrum = torch.stft(sample, n_fft=n_fft, hop_length=hop_length, window=window, center=True, return_complex=True) + spectrum = torch.view_as_complex(torch.view_as_real(spectrum).contiguous()) * torch.ones_like(spectrum) + _probe_complex_scatter_add(spectrum) + torch.istft(spectrum, n_fft=n_fft, hop_length=hop_length, window=window, center=True, length=sample_length) + return True + except Exception: + return False + + +def should_fallback_to_cpu_for_complex_ops(device: torch.device) -> bool: + """Return whether complex spectral operations should use the legacy CPU path.""" + if os.environ.get("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX") == "1": + return True + + device_index = -1 if device.index is None else int(device.index) + return not _supports_complex_spectral_ops(device.type, device_index) + + +def should_fallback_to_cpu_for_demucs_mask(device: torch.device, cac: bool) -> bool: + """Keep non-CaC Demucs Wiener masking on CPU because the spectral probe does not cover it.""" + return (device.type == "mps" and not cac) or should_fallback_to_cpu_for_complex_ops(device) + + +def _mps_memory_reading(counter: str) -> int: + """Return a torch.mps memory counter in bytes, or 0 when it is unavailable.""" + try: + if not torch.backends.mps.is_available(): + return 0 + value = getattr(torch.mps, counter)() + except (AttributeError, RuntimeError, OSError, ValueError): + return 0 + + return int(value) if value and value > 0 else 0 + + +def mps_accumulation_budget_bytes() -> int: + """Return how many bytes of duration-scaled buffers may stay on MPS. + + Model weights are already resident by the time this is called, so the budget + is measured against what is still free rather than against the whole working + set: buffers may take at most MPS_BUFFER_HEADROOM_SHARE of the remaining + room, which leaves at least as much again for activations. It never drops + below MAX_MPS_FULL_TRACK_BUFFER_BYTES. + """ + override = os.environ.get(MPS_BUFFER_BUDGET_ENV) + if override: + try: + override_gib = float(override) + except ValueError: + override_gib = 0.0 + if override_gib > 0: + return int(override_gib * 1024**3) + + recommended = _mps_memory_reading("recommended_max_memory") + if recommended <= 0: + return MAX_MPS_FULL_TRACK_BUFFER_BYTES + + # driver_allocated_memory counts the allocator's cached blocks as well as + # live tensors, so free room is understated while blocks are being reused. + # Erring small here is the safe direction for a working-set guard. + free = max(recommended - _mps_memory_reading("driver_allocated_memory"), 0) + + return max(int(free * MPS_BUFFER_HEADROOM_SHARE), MAX_MPS_FULL_TRACK_BUFFER_BYTES) + + +def should_accumulate_on_device(device: torch.device, estimated_bytes: int) -> bool: + """Return whether duration-scaled buffers fit the bounded MPS fast path.""" + return device.type == "mps" and estimated_bytes <= mps_accumulation_budget_bytes() diff --git a/audio_separator/separator/uvr_lib_v5/roformer/attend.py b/audio_separator/separator/uvr_lib_v5/roformer/attend.py index 5bf7b8f..ad3d03c 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/attend.py +++ b/audio_separator/separator/uvr_lib_v5/roformer/attend.py @@ -4,6 +4,7 @@ import torch from torch import nn, einsum +from torch.nn.attention import SDPBackend, sdpa_kernel import torch.nn.functional as F from einops import rearrange, reduce @@ -29,6 +30,22 @@ def exists(val): return val is not None +def _sdpa_backends(config): + """Translate the legacy SDPA flags to the Dynamo-compatible API.""" + backends = [] + if config.enable_flash: + backends.append(SDPBackend.FLASH_ATTENTION) + if config.enable_mem_efficient: + backends.append(SDPBackend.EFFICIENT_ATTENTION) + if config.enable_math: + backends.append(SDPBackend.MATH) + + # torch.backends.cuda.sdp_kernel enabled cuDNN by default even though the + # local config predates that fourth flag. Preserve that behavior. + backends.append(SDPBackend.CUDNN_ATTENTION) + return backends + + def once(fn): called = False @@ -49,8 +66,9 @@ def inner(x): class Attend(nn.Module): - def __init__(self, dropout=0.0, flash=False): + def __init__(self, dropout=0.0, flash=False, scale=None): super().__init__() + self.scale = scale self.dropout = dropout self.attn_dropout = nn.Dropout(dropout) @@ -84,9 +102,15 @@ def flash_attn(self, q, k, v): if is_cuda and q.dtype != torch.float16: config = FlashAttentionConfig(False, True, True) - # pytorch 2.0 flash attn: q, k, v, mask, dropout, softmax_scale - with torch.backends.cuda.sdp_kernel(**config._asdict()): - out = F.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout if self.training else 0.0) + # Keep SDPA backend selection inside the graphable PyTorch API. + with sdpa_kernel(_sdpa_backends(config)): + out = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.dropout if self.training else 0.0, + scale=self.scale, + ) return out @@ -101,7 +125,7 @@ def forward(self, q, k, v): q_len, k_len, device = q.shape[-2], k.shape[-2], q.device - scale = q.shape[-1] ** -0.5 + scale = self.scale if exists(self.scale) else q.shape[-1] ** -0.5 # DML has no SDPA โ€” fall through to the einsum path. Gated so every # other device keeps its exact existing behavior. (Issue #292) 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 4f4543c..c8918cd 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py +++ b/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py @@ -11,11 +11,13 @@ from beartype import beartype from rotary_embedding_torch import RotaryEmbedding -from rotary_embedding_torch.rotary_embedding_torch import rotate_half as _rotate_half_no_cat from einops import rearrange, pack, unpack from einops.layers.torch import Rearrange +from audio_separator.separator.uvr_lib_v5.device_utils import should_fallback_to_cpu_for_complex_ops +from .rotary import rotate_queries_or_keys as _rotate_queries_or_keys + # helper functions @@ -29,26 +31,6 @@ def _is_dml_device(device) -> bool: return device.type == "privateuseone" -def _rotate_queries_or_keys(rotary_embed, t): - """Apply rotary position embedding, avoiding zero-width tensor ops on DML. - - rotary_embedding_torch's apply_rotary_emb concatenates (possibly empty) - unrotated edge slices around the rotated block; torch-directml rejects - zero-sized tensor ops with 'The parameter is incorrect.'. These models - always rotate the full head dimension, so the edge slices are empty and - the concat is a no-op โ€” compute the rotation directly instead. Verified - equivalent to the library implementation by unit test. (Issue #292) - """ - if not _is_dml_device(t.device): - return rotary_embed.rotate_queries_or_keys(t) - seq_len = t.shape[-2] - freqs = rotary_embed.forward(rotary_embed.get_seq_pos(seq_len, device=t.device, dtype=t.dtype), seq_len=seq_len) - if freqs.shape[-1] != t.shape[-1]: - # Partial-dim rotation would need the edge concat โ€” unreachable here - # (RotaryEmbedding(dim=dim_head) rotates the full head dim). - return rotary_embed.rotate_queries_or_keys(t) - return t * freqs.cos() + _rotate_half_no_cat(t) * freqs.sin() - def exists(val): return val is not None @@ -69,6 +51,8 @@ def unpack_one(t, ps, pattern): def l2norm(t): + if t.dtype in (torch.float16, torch.bfloat16): + return F.normalize(t.float(), dim=-1, p=2).to(t.dtype) return F.normalize(t, dim=-1, p=2) @@ -80,7 +64,11 @@ def __init__(self, dim): def forward(self, x): x = x.to(self.gamma.device) - return F.normalize(x, dim=-1) * self.scale * self.gamma + if x.dtype in (torch.float16, torch.bfloat16): + normalized = F.normalize(x.float(), dim=-1).to(x.dtype) + else: + normalized = F.normalize(x, dim=-1) + return normalized * self.scale * self.gamma # attention @@ -460,7 +448,8 @@ 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 + # Use the legacy CPU hop unless the current MPS runtime supports every required complex operation. + x_is_mps = original_device.type == "mps" and should_fallback_to_cpu_for_complex_ops(original_device) # 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 @@ -486,7 +475,7 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False): stft_window = self.stft_window_fn().to(device) - if x_is_dml: + if x_is_mps or 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: @@ -498,6 +487,10 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False): x = rearrange(stft_repr, "b f t c -> b t (f c)") + band_split_dtype = next(self.band_split.parameters()).dtype + if x.dtype != band_split_dtype: + x = x.to(band_split_dtype) + x = self.band_split(x) # axial / hierarchical attention @@ -540,10 +533,13 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False): # complex number multiplication - if x_is_dml: + if x_is_mps or x_is_dml: stft_repr = stft_repr.cpu() mask = mask.cpu() + if mask.dtype != stft_repr.dtype: + mask = mask.to(stft_repr.dtype) + stft_repr = torch.view_as_complex(stft_repr) mask = torch.view_as_complex(mask) @@ -553,7 +549,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 or x_is_dml) else stft_window, return_complex=False).to(device) + recon_audio = torch.istft(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 5fa0ce3..157eb5a 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 @@ -11,12 +11,14 @@ from beartype import beartype from rotary_embedding_torch import RotaryEmbedding -from rotary_embedding_torch.rotary_embedding_torch import rotate_half as _rotate_half_no_cat from einops import rearrange, pack, unpack, reduce, repeat from librosa import filters +from audio_separator.separator.uvr_lib_v5.device_utils import should_fallback_to_cpu_for_complex_ops +from .rotary import rotate_queries_or_keys as _rotate_queries_or_keys + def _is_dml_device(device) -> bool: @@ -28,26 +30,6 @@ def _is_dml_device(device) -> bool: return device.type == "privateuseone" -def _rotate_queries_or_keys(rotary_embed, t): - """Apply rotary position embedding, avoiding zero-width tensor ops on DML. - - rotary_embedding_torch's apply_rotary_emb concatenates (possibly empty) - unrotated edge slices around the rotated block; torch-directml rejects - zero-sized tensor ops with 'The parameter is incorrect.'. These models - always rotate the full head dimension, so the edge slices are empty and - the concat is a no-op โ€” compute the rotation directly instead. Verified - equivalent to the library implementation by unit test. (Issue #292) - """ - if not _is_dml_device(t.device): - return rotary_embed.rotate_queries_or_keys(t) - seq_len = t.shape[-2] - freqs = rotary_embed.forward(rotary_embed.get_seq_pos(seq_len, device=t.device, dtype=t.dtype), seq_len=seq_len) - if freqs.shape[-1] != t.shape[-1]: - # Partial-dim rotation would need the edge concat โ€” unreachable here - # (RotaryEmbedding(dim=dim_head) rotates the full head dim). - return rotary_embed.rotate_queries_or_keys(t) - return t * freqs.cos() + _rotate_half_no_cat(t) * freqs.sin() - def exists(val): return val is not None @@ -78,7 +60,11 @@ def __init__(self, dim): def forward(self, x): x = x.to(self.gamma.device) - return F.normalize(x, dim=-1) * self.scale * self.gamma + if x.dtype in (torch.float16, torch.bfloat16): + normalized = F.normalize(x.float(), dim=-1).to(x.dtype) + else: + normalized = F.normalize(x, dim=-1) + return normalized * self.scale * self.gamma class FeedForward(Module): @@ -369,7 +355,8 @@ 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 + # Use the legacy CPU hop unless the current MPS runtime supports every required complex operation. + x_is_mps = original_device.type == "mps" and should_fallback_to_cpu_for_complex_ops(original_device) # 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 @@ -412,6 +399,9 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False): x = stft_repr[batch_arange, self.freq_indices.cpu()] if x_is_mps else stft_repr[batch_arange, self.freq_indices] x = rearrange(x, "b f t c -> b t (f c)") + band_split_dtype = next(self.band_split.parameters()).dtype + if x.dtype != band_split_dtype: + x = x.to(band_split_dtype) x = self.band_split(x) @@ -443,6 +433,8 @@ def forward(self, raw_audio, target=None, return_loss_breakdown=False): stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c") + if masks.dtype != stft_repr.dtype: + masks = masks.to(stft_repr.dtype) stft_repr = torch.view_as_complex(stft_repr) masks = torch.view_as_complex(masks) diff --git a/audio_separator/separator/uvr_lib_v5/roformer/rotary.py b/audio_separator/separator/uvr_lib_v5/roformer/rotary.py new file mode 100644 index 0000000..0f4445b --- /dev/null +++ b/audio_separator/separator/uvr_lib_v5/roformer/rotary.py @@ -0,0 +1,80 @@ +import torch +from rotary_embedding_torch.rotary_embedding_torch import apply_rotary_emb, rotate_half + +from audio_separator.separator.uvr_lib_v5.device_utils import autocast_disabled + + +def _is_dml_device(device: torch.device) -> bool: + """Return whether a device uses torch-directml's private backend slot.""" + return device.type == "privateuseone" + + +def _float32_frequencies(rotary_embed, *, seq_len: int, device: torch.device) -> torch.Tensor: + """Build or retrieve rotary angles without allowing autocast to reduce precision.""" + # A compiled graph is shared by many regional Transformer instances. Their + # time and frequency embeddings have different cache shapes, so reading or + # mutating cached_freqs here creates instance-state guards and repeated + # Dynamo recompilations. Angle construction is cheap relative to attention + # and becomes part of the compiled graph, so skip the mutable cache while + # Dynamo is tracing. + is_compiling = torch.compiler.is_compiling() + should_cache = ( + not is_compiling + and rotary_embed.cache_if_possible + and not rotary_embed.learned_freq + and rotary_embed.freqs_for != "pixel" + ) + cached_freqs = rotary_embed.cached_freqs if should_cache else None + + if ( + should_cache + and cached_freqs is not None + and cached_freqs.dtype == torch.float32 + and cached_freqs.device == device + and seq_len <= cached_freqs.shape[0] + ): + return cached_freqs[:seq_len].detach() + + positions = rotary_embed.get_seq_pos(seq_len, device=device, dtype=torch.float32) + base_frequencies = rotary_embed.freqs.to(dtype=torch.float32) + frequencies = torch.einsum("..., f -> ... f", positions, base_frequencies) + frequencies = torch.repeat_interleave(frequencies, 2, dim=-1) + + if should_cache: + # Replace an older low-precision cache instead of allowing it to be + # reused after leaving an autocast region. + rotary_embed.cached_freqs = frequencies.detach() + + return frequencies + + +def rotate_queries_or_keys(rotary_embed, tensor: torch.Tensor) -> torch.Tensor: + """Apply full-head rotary embeddings with float32 angles on every backend. + + rotary-embedding-torch 0.6.5 disables CUDA autocast only, so CPU and MPS + autocast can otherwise lower the precision of its position/angle einsum. + DirectML also rejects the empty edge tensors concatenated by the upstream + helper when the full head dimension is rotated, so that case keeps the + existing concat-free implementation. + + This module reads rotary-embedding-torch internals (freqs, cached_freqs, + cache_if_possible, learned_freq, freqs_for, default_seq_dim, get_seq_pos) + rather than a public API, which is why pyproject pins the dependency to + the 0.6.x series these helpers are validated against. + """ + input_dtype = tensor.dtype + seq_dim = rotary_embed.default_seq_dim + seq_len = tensor.shape[seq_dim] + + with autocast_disabled(tensor.device): + frequencies = _float32_frequencies(rotary_embed, seq_len=seq_len, device=tensor.device) + + if seq_dim == -3: + frequencies = frequencies[:, None, :] + + if _is_dml_device(tensor.device) and frequencies.shape[-1] == tensor.shape[-1]: + rotated = tensor * frequencies.cos() + rotate_half(tensor) * frequencies.sin() + else: + rotated = apply_rotary_emb(frequencies, tensor, seq_dim=seq_dim) + + return rotated.to(dtype=input_dtype) diff --git a/audio_separator/separator/uvr_lib_v5/stft.py b/audio_separator/separator/uvr_lib_v5/stft.py index f440395..5ff9f86 100644 --- a/audio_separator/separator/uvr_lib_v5/stft.py +++ b/audio_separator/separator/uvr_lib_v5/stft.py @@ -1,5 +1,7 @@ import torch +from audio_separator.separator.uvr_lib_v5.device_utils import should_fallback_to_cpu_for_complex_ops + class STFT: """ @@ -18,11 +20,11 @@ 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 __call__(self, input_tensor): - # Determine if the input tensor's device is not a standard computing device (i.e., not CPU or CUDA). - is_non_standard_device = not input_tensor.device.type in ["cuda", "cpu"] + original_device = input_tensor.device + should_fallback = should_fallback_to_cpu_for_complex_ops(original_device) - # If on a non-standard device, temporarily move the tensor to CPU for processing. - if is_non_standard_device: + # Use CPU only when the current backend lacks the required complex operations. + if should_fallback: input_tensor = input_tensor.cpu() # Transfer the pre-defined window tensor to the same device as the input tensor. @@ -48,9 +50,9 @@ def __call__(self, input_tensor): [*batch_dimensions, channel_dim * 2, -1, permuted_stft_output.shape[-1]] ) - # If the original tensor was on a non-standard device, move the processed tensor back to that device. - if is_non_standard_device: - final_output = final_output.to(self.device) + # Restore the input device after using the CPU fallback. + if should_fallback: + final_output = final_output.to(original_device) # Return the transformed tensor, sliced to retain only the required frequency dimension (`dim_f`). return final_output[..., : self.dim_f, :] @@ -97,11 +99,11 @@ def prepare_for_istft(self, padded_tensor, batch_dimensions, channel_dim, num_fr return complex_tensor def inverse(self, input_tensor): - # Determine if the input tensor's device is not a standard computing device (i.e., not CPU or CUDA). - is_non_standard_device = not input_tensor.device.type in ["cuda", "cpu"] + original_device = input_tensor.device + should_fallback = should_fallback_to_cpu_for_complex_ops(original_device) - # If on a non-standard device, temporarily move the tensor to CPU for processing. - if is_non_standard_device: + # Use CPU only when the current backend lacks the required complex operations. + if should_fallback: input_tensor = input_tensor.cpu() # Transfer the pre-defined Hann window tensor to the same device as the input tensor. @@ -119,8 +121,8 @@ def inverse(self, input_tensor): # Reshape ISTFT result to restore original batch and channel dimensions. final_output = istft_result.reshape([*batch_dimensions, 2, -1]) - # If the original tensor was on a non-standard device, move the processed tensor back to that device. - if is_non_standard_device: - final_output = final_output.to(self.device) + # Restore the input device after using the CPU fallback. + if should_fallback: + final_output = final_output.to(original_device) return final_output diff --git a/audio_separator/separator/uvr_lib_v5/tfc_tdf_v3.py b/audio_separator/separator/uvr_lib_v5/tfc_tdf_v3.py index 4d3356f..59b8974 100644 --- a/audio_separator/separator/uvr_lib_v5/tfc_tdf_v3.py +++ b/audio_separator/separator/uvr_lib_v5/tfc_tdf_v3.py @@ -2,6 +2,8 @@ import torch.nn as nn from functools import partial +from audio_separator.separator.uvr_lib_v5.device_utils import should_fallback_to_cpu_for_complex_ops + class STFT: def __init__(self, n_fft, hop_length, dim_f, device): self.n_fft = n_fft @@ -12,8 +14,9 @@ def __init__(self, n_fft, hop_length, dim_f, device): def __call__(self, x): - x_is_mps = not x.device.type in ["cuda", "cpu"] - if x_is_mps: + original_device = x.device + should_fallback = should_fallback_to_cpu_for_complex_ops(original_device) + if should_fallback: x = x.cpu() window = self.window.to(x.device) @@ -24,15 +27,16 @@ def __call__(self, x): x = x.permute([0, 3, 1, 2]) x = x.reshape([*batch_dims, c, 2, -1, x.shape[-1]]).reshape([*batch_dims, c * 2, -1, x.shape[-1]]) - if x_is_mps: - x = x.to(self.device) + if should_fallback: + x = x.to(original_device) return x[..., :self.dim_f, :] def inverse(self, x): - x_is_mps = not x.device.type in ["cuda", "cpu"] - if x_is_mps: + original_device = x.device + should_fallback = should_fallback_to_cpu_for_complex_ops(original_device) + if should_fallback: x = x.cpu() window = self.window.to(x.device) @@ -47,8 +51,8 @@ def inverse(self, x): x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop_length, window=window, center=True) x = x.reshape([*batch_dims, 2, -1]) - if x_is_mps: - x = x.to(self.device) + if should_fallback: + x = x.to(original_device) return x @@ -266,4 +270,3 @@ def forward(self, x): return x - diff --git a/audio_separator/utils/cli.py b/audio_separator/utils/cli.py index 23bbc90..62e3369 100755 --- a/audio_separator/utils/cli.py +++ b/audio_separator/utils/cli.py @@ -60,7 +60,18 @@ def main(): single_stem_help = "Output only single stem, e.g. Instrumental, Vocals, Drums, Bass, Guitar, Piano, Other. Example: --single_stem=Instrumental" sample_rate_help = "Modify the sample rate of the output audio (default: %(default)s). Example: --sample_rate=44100" use_soundfile_help = "Use soundfile to write audio output (default: %(default)s). Example: --use_soundfile" - use_autocast_help = "Use PyTorch autocast for faster inference (default: %(default)s). Do not use for CPU inference. Example: --use_autocast" + use_autocast_help = ( + "Use PyTorch autocast when supported (default: %(default)s). Example: --use_autocast" + ) + use_native_fp16_help = ( + "Use native float16 for verified model/device combinations (default: %(default)s). " + "Mutually exclusive with --use_autocast. Example: --use_native_fp16" + ) + use_torch_compile_help = ( + "Compile verified repeated model blocks when supported (default: %(default)s). " + "Best for long inputs or repeated same-shape runs; a fresh compiler cache can make " + "the first run slower. Example: --use_torch_compile" + ) use_directml_help = "Use DirectML for hardware-accelerated inference on Windows AMD/Intel GPUs (experimental; requires the 'dml' extra). Example: --use_directml" chunk_duration_help = "Split audio into chunks of this duration in seconds (default: %(default)s = no chunking). Useful for processing very long audio files on systems with limited memory. Recommended: 600 (10 minutes) for files >1 hour. Chunks are concatenated without overlap/crossfade. Example: --chunk_duration=600" ensemble_algorithm_help = "Algorithm to use for ensembling multiple models (default: avg_wave). Choices: avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft, uvr_max_spec, uvr_min_spec, ensemble_wav. Example: --ensemble_algorithm=uvr_max_spec" @@ -76,7 +87,10 @@ def main(): common_params.add_argument("--single_stem", default=None, help=single_stem_help) common_params.add_argument("--sample_rate", type=int, default=44100, help=sample_rate_help) common_params.add_argument("--use_soundfile", action="store_true", help=use_soundfile_help) - common_params.add_argument("--use_autocast", action="store_true", help=use_autocast_help) + precision_params = common_params.add_mutually_exclusive_group() + precision_params.add_argument("--use_autocast", action="store_true", help=use_autocast_help) + precision_params.add_argument("--use_native_fp16", action="store_true", help=use_native_fp16_help) + common_params.add_argument("--use_torch_compile", action="store_true", help=use_torch_compile_help) common_params.add_argument("--use_directml", action="store_true", help=use_directml_help) common_params.add_argument("--chunk_duration", type=float, default=None, help=chunk_duration_help) common_params.add_argument( @@ -249,6 +263,8 @@ def main(): sample_rate=args.sample_rate, use_soundfile=args.use_soundfile, use_autocast=args.use_autocast, + use_native_fp16=args.use_native_fp16, + use_torch_compile=args.use_torch_compile, use_directml=args.use_directml, chunk_duration=args.chunk_duration, ensemble_algorithm=args.ensemble_algorithm, diff --git a/poetry.lock b/poetry.lock index 7ffdb76..fc2ed8d 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. [[package]] name = "absl-py" @@ -764,6 +764,106 @@ ssh = ["bcrypt (>=3.1.5)"] test = ["certifi (>=2024)", "cryptography-vectors (==46.0.6)", "pretend (>=0.7)", "pytest (>=7.4.0)", "pytest-benchmark (>=4.0)", "pytest-cov (>=2.10.1)", "pytest-xdist (>=3.5.0)"] test-randomorder = ["pytest-randomly"] +[[package]] +name = "cuda-bindings" +version = "13.3.1" +description = "Python bindings for CUDA" +optional = false +python-versions = ">=3.10" +groups = ["main"] +markers = "python_version < \"3.15\" and (platform_machine == \"arm64\" or python_version == \"3.14\" or sys_platform == \"linux\") and platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\" or python_version == \"3.14\")" +files = [ + {file = "cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86"}, + {file = "cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9851b0caa8bfd3bc6fa054eaf57bea7c8e9c3a62db2d2621224677f49f3c53d0"}, + {file = "cuda_bindings-13.3.1-cp310-cp310-win_amd64.whl", hash = "sha256:c0c4b1a995098c46695c24257a342dc97d6e6d3f3050b944c9f43bd26d734051"}, + {file = "cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5879712accf6e14bb01aa5e67440eb84998b8d104b509cc7a6dc0b8f656a474"}, + {file = "cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:04436a9364059c84b8f9636f359eccda1cf814341f5b670c71d80d2f79dbc708"}, + {file = "cuda_bindings-13.3.1-cp311-cp311-win_amd64.whl", hash = "sha256:507b0e19e7f934c5e30f30f0244ad70a75812619a7d3a0d742543caae1bd50f1"}, + {file = "cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49"}, + {file = "cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e32d08f71ebcdf00f0f41eab2eb37e8da94c8ed411cc9f7f7a019ce6b34abe3a"}, + {file = "cuda_bindings-13.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:b134dd8c5c66ae4c4ad814f7aee88fd215353c077010cbc47e3b55ed35ec9eff"}, + {file = "cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9efb21c1ee64981e184b9e0ba5eb3179e5ba3d4b51665a6cb52b8ef3d01a7cbf"}, + {file = "cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2732904099e0a4d4db774a5fc6d91ee95fae065b4d2ecabb4968c5fe2406c9d7"}, + {file = "cuda_bindings-13.3.1-cp313-cp313-win_amd64.whl", hash = "sha256:18c8c167c8907b8f02531ca810534315c458dabef31f7965095619bf647b9202"}, + {file = "cuda_bindings-13.3.1-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ab2f74ed65bfef4163ba07a8db16f1085e0729291db12a2423aff84ee8278b8"}, + {file = "cuda_bindings-13.3.1-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:efd4c814d311ec08c981f6dded1dbe7d4b371067ee4f6c14cccec4bde9590f80"}, + {file = "cuda_bindings-13.3.1-cp314-cp314-win_amd64.whl", hash = "sha256:8de12ef60bf40756852cb62bbb40460609269f6ece522903d1cc93d73a3ececb"}, + {file = "cuda_bindings-13.3.1-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c3c772dfff49681541d59630c90f858e173ac926b9c593a2b7123f2a1043cc76"}, + {file = "cuda_bindings-13.3.1-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:36febb7c1079d68a981dbbd8d5a67235b399802b82075c9388624719607e52b9"}, + {file = "cuda_bindings-13.3.1-cp314-cp314t-win_amd64.whl", hash = "sha256:61120b5e4f4a63f67efd7e7396914cb9ef871bb1f0021e990fb70277be240a4d"}, +] + +[package.dependencies] +cuda-pathfinder = ">=1.4.2" + +[package.extras] +all = ["cuda-toolkit (==13.*)", "cuda-toolkit[cufile] (==13.*) ; sys_platform == \"linux\"", "cuda-toolkit[nvfatbin,nvjitlink,nvrtc,nvvm] (==13.*)", "nvidia-cudla (==13.*) ; platform_system == \"Linux\" and platform_machine == \"aarch64\""] + +[[package]] +name = "cuda-pathfinder" +version = "1.6.0" +description = "Pathfinder for CUDA components" +optional = false +python-versions = ">=3.10" +groups = ["main"] +markers = "python_version < \"3.15\" and (platform_machine == \"arm64\" or python_version == \"3.14\" or sys_platform == \"linux\") and platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\" or python_version == \"3.14\")" +files = [ + {file = "cuda_pathfinder-1.6.0-py3-none-any.whl", hash = "sha256:1503af579d8379c24bdd65528379bc57039b0455be9f5f9686cf8e473a1fce51"}, +] + +[[package]] +name = "cuda-toolkit" +version = "13.0.3" +description = "CUDA Toolkit meta-package" +optional = false +python-versions = "*" +groups = ["main"] +markers = "(platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\" or python_version >= \"3.14\")" +files = [ + {file = "cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f"}, +] + +[package.dependencies] +nvidia-cublas = {version = "==13.1.1.3.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (extra == \"cublas\" or extra == \"cusolver\")"} +nvidia-cuda-cupti = {version = "==13.0.85.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cupti\""} +nvidia-cuda-nvrtc = {version = "==13.0.88.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (extra == \"cublas\" or extra == \"nvrtc\")"} +nvidia-cuda-runtime = {version = "==13.0.96.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cudart\""} +nvidia-cufft = {version = "==12.0.0.61.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cufft\""} +nvidia-cufile = {version = "==1.15.1.6.*", optional = true, markers = "sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cufile\""} +nvidia-curand = {version = "==10.4.0.35.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"curand\""} +nvidia-cusolver = {version = "==12.0.4.66.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cusolver\""} +nvidia-cusparse = {version = "==12.6.3.3.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (extra == \"cusolver\" or extra == \"cusparse\")"} +nvidia-nvjitlink = {version = ">=13.0.88,<14", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (extra == \"cufft\" or extra == \"cusolver\" or extra == \"cusparse\" or extra == \"nvjitlink\")"} +nvidia-nvtx = {version = "==13.0.85.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"nvtx\""} + +[package.extras] +all = ["nvidia-cublas (==13.1.1.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-cccl (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-crt (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-culibos (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")", "nvidia-cuda-cupti (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-cuxxfilt (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-nvcc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-nvrtc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-opencl (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-profiler-api (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-runtime (==13.0.96.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-sanitizer-api (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cufft (==12.0.0.61.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cufile (==1.15.1.6.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")", "nvidia-curand (==10.4.0.35.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cusolver (==12.0.4.66.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cusparse (==12.6.3.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-npp (==13.0.1.2.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvfatbin (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjpeg (==13.0.1.86.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvml-dev (==13.0.87.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvptxcompiler (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvtx (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvvm (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cccl = ["nvidia-cuda-cccl (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +crt = ["nvidia-cuda-crt (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cublas = ["nvidia-cublas (==13.1.1.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-nvrtc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cudart = ["nvidia-cuda-runtime (==13.0.96.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cufft = ["nvidia-cufft (==12.0.0.61.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cufile = ["nvidia-cufile (==1.15.1.6.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")"] +culibos = ["nvidia-cuda-culibos (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")"] +cupti = ["nvidia-cuda-cupti (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +curand = ["nvidia-curand (==10.4.0.35.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cusolver = ["nvidia-cublas (==13.1.1.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cusolver (==12.0.4.66.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cusparse (==12.6.3.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cusparse = ["nvidia-cusparse (==12.6.3.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cuxxfilt = ["nvidia-cuda-cuxxfilt (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +npp = ["nvidia-npp (==13.0.1.2.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvcc = ["nvidia-cuda-crt (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-nvcc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-runtime (==13.0.96.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvvm (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvfatbin = ["nvidia-nvfatbin (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvjitlink = ["nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvjpeg = ["nvidia-nvjpeg (==13.0.1.86.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvml = ["nvidia-nvml-dev (==13.0.87.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvptxcompiler = ["nvidia-nvptxcompiler (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvrtc = ["nvidia-cuda-nvrtc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvtx = ["nvidia-nvtx (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvvm = ["nvidia-nvvm (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +opencl = ["nvidia-cuda-opencl (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +profiler = ["nvidia-cuda-profiler-api (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +sanitizer = ["nvidia-cuda-sanitizer-api (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] + [[package]] name = "cycler" version = "0.12.1" @@ -1152,18 +1252,18 @@ files = [ google-auth = ">=2.14.1,<3.0.0" googleapis-common-protos = ">=1.56.3,<2.0.0" grpcio = [ - {version = ">=1.75.1,<2.0.0", optional = true, markers = "python_version >= \"3.14\" and extra == \"grpc\""}, + {version = ">=1.33.2,<2.0.0", optional = true, markers = "python_version < \"3.11\" and extra == \"grpc\""}, {version = ">=1.49.1,<2.0.0", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""}, - {version = ">=1.33.2,<2.0.0", optional = true, markers = "extra == \"grpc\""}, + {version = ">=1.75.1,<2.0.0", optional = true, markers = "python_version >= \"3.14\" and extra == \"grpc\""}, ] grpcio-status = [ - {version = ">=1.75.1,<2.0.0", optional = true, markers = "python_version >= \"3.14\" and extra == \"grpc\""}, - {version = ">=1.49.1,<2.0.0", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""}, {version = ">=1.33.2,<2.0.0", optional = true, markers = "extra == \"grpc\""}, + {version = ">=1.49.1,<2.0.0", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""}, + {version = ">=1.75.1,<2.0.0", optional = true, markers = "python_version >= \"3.14\" and extra == \"grpc\""}, ] proto-plus = [ - {version = ">=1.25.0,<2.0.0", markers = "python_version >= \"3.13\""}, {version = ">=1.22.3,<2.0.0"}, + {version = ">=1.25.0,<2.0.0", markers = "python_version >= \"3.13\""}, ] protobuf = ">=4.25.8,<7.0.0" requests = ">=2.20.0,<3.0.0" @@ -1236,12 +1336,12 @@ google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0", extras google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0" google-cloud-core = ">=1.4.1,<3.0.0" grpcio = [ - {version = ">=1.75.1,<2.0.0", markers = "python_version >= \"3.14\""}, {version = ">=1.33.2,<2.0.0", markers = "python_version < \"3.14\""}, + {version = ">=1.75.1,<2.0.0", markers = "python_version >= \"3.14\""}, ] proto-plus = [ + {version = ">=1.22.3,<2.0.0", markers = "python_version < \"3.13\""}, {version = ">=1.25.0,<2.0.0", markers = "python_version >= \"3.13\""}, - {version = ">=1.22.3,<2.0.0"}, ] protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<7.0.0" @@ -2013,10 +2113,10 @@ files = [ [package.dependencies] numpy = [ - {version = ">=2.1.0", markers = "python_version >= \"3.13\""}, - {version = ">=1.26.0", markers = "python_version == \"3.12\""}, {version = ">=1.23.3", markers = "python_version == \"3.11\""}, - {version = ">=1.21.2", markers = "python_version == \"3.10\""}, + {version = ">=1.21.2", markers = "python_version >= \"3.10\""}, + {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, + {version = ">=2.1.0", markers = "python_version >= \"3.13\""}, ] [package.extras] @@ -2351,19 +2451,39 @@ files = [ ] [[package]] -name = "nvidia-cublas-cu12" -version = "12.6.4.1" +name = "nvidia-cublas" +version = "13.1.1.3" description = "CUBLAS native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08ed2686e9875d01b58e3cb379c6896df8e76c75e0d4a7f7dace3d7b6d9ef8eb"}, - {file = "nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:235f728d6e2a409eddf1df58d5b0921cf80cfa9e72b9f2775ccb7b4a87984668"}, - {file = "nvidia_cublas_cu12-12.6.4.1-py3-none-win_amd64.whl", hash = "sha256:9e4fa264f4d8a4eb0cdbd34beadc029f453b3bafae02401e999cf3d5a5af75f8"}, + {file = "nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5"}, + {file = "nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436"}, + {file = "nvidia_cublas-13.1.1.3-py3-none-win_amd64.whl", hash = "sha256:b6cdce694e47ff6aadf0a69df1cab6628d696f5ff56e8d16af50309d855fa20f"}, ] +[package.dependencies] +nvidia-cuda-nvrtc = "*" + +[[package]] +name = "nvidia-cublas" +version = "13.6.0.2" +description = "CUBLAS native runtime libraries" +optional = false +python-versions = ">=3" +groups = ["main"] +markers = "(python_version >= \"3.14\" or platform_machine != \"aarch64\" and platform_machine != \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and (platform_machine != \"aarch64\" and platform_machine != \"x86_64\" or sys_platform != \"linux\") and (platform_machine != \"x86_64\" or sys_platform != \"linux\" and sys_platform != \"win32\") and (python_version >= \"3.14\" or sys_platform == \"darwin\" or sys_platform == \"linux\")" +files = [ + {file = "nvidia_cublas-13.6.0.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:946f6a252b1cc72d8de912c75975fd6d8ba44f67d4e5044fe764ddb909f4a688"}, + {file = "nvidia_cublas-13.6.0.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b82c80c886cea6da6e149a5c3bdba274f12b7e4ec4b00a050b916b0446fb4153"}, + {file = "nvidia_cublas-13.6.0.2-py3-none-win_amd64.whl", hash = "sha256:3b5bcd6bfb6f65010ebf195851bcb9b2aa34b9fe08479432002991c1fe84b67d"}, +] + +[package.dependencies] +nvidia-cuda-nvrtc = "*" + [[package]] name = "nvidia-cublas-cu12" version = "12.8.4.1" @@ -2371,7 +2491,7 @@ description = "CUBLAS native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b86f6dd8935884615a0683b663891d43781b819ac4f2ba2b0c9604676af346d0"}, {file = "nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142"}, @@ -2379,19 +2499,17 @@ files = [ ] [[package]] -name = "nvidia-cuda-cupti-cu12" -version = "12.6.80" +name = "nvidia-cuda-cupti" +version = "13.0.85" description = "CUDA profiling tools runtime libs." optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:166ee35a3ff1587f2490364f90eeeb8da06cd867bd5b701bf7f9a02b78bc63fc"}, - {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_aarch64.whl", hash = "sha256:358b4a1d35370353d52e12f0a7d1769fc01ff74a191689d3870b2123156184c4"}, - {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6768bad6cab4f19e8292125e5f1ac8aa7d1718704012a0e3272a6f61c4bce132"}, - {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73"}, - {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-win_amd64.whl", hash = "sha256:bbe6ae76e83ce5251b56e8c8e61a964f757175682bbad058b170b136266ab00a"}, + {file = "nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151"}, + {file = "nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8"}, + {file = "nvidia_cuda_cupti-13.0.85-py3-none-win_amd64.whl", hash = "sha256:683f58d301548deeefcb8f6fac1b8d907691b9d8b18eccab417f51e362102f00"}, ] [[package]] @@ -2401,7 +2519,7 @@ description = "CUDA profiling tools runtime libs." optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4412396548808ddfed3f17a467b104ba7751e6b58678a4b840675c56d21cf7ed"}, {file = "nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182"}, @@ -2409,17 +2527,31 @@ files = [ ] [[package]] -name = "nvidia-cuda-nvrtc-cu12" -version = "12.6.77" +name = "nvidia-cuda-nvrtc" +version = "13.0.88" +description = "NVRTC native runtime libraries" +optional = false +python-versions = ">=3" +groups = ["main"] +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" +files = [ + {file = "nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575"}, + {file = "nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b"}, + {file = "nvidia_cuda_nvrtc-13.0.88-py3-none-win_amd64.whl", hash = "sha256:6bcd4e7f8e205cbe644f5a98f2f799bef9556fefc89dd786e79a16312ce49872"}, +] + +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.3.33" description = "NVRTC native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "(python_version >= \"3.14\" or platform_machine != \"aarch64\" and platform_machine != \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and (platform_machine != \"aarch64\" and platform_machine != \"x86_64\" or sys_platform != \"linux\") and (platform_machine != \"x86_64\" or sys_platform != \"linux\" and sys_platform != \"win32\") and (python_version >= \"3.14\" or sys_platform == \"darwin\" or sys_platform == \"linux\")" files = [ - {file = "nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5847f1d6e5b757f1d2b3991a01082a44aad6f10ab3c5c0213fa3e25bddc25a13"}, - {file = "nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:35b0cc6ee3a9636d5409133e79273ce1f3fd087abb0532d2d2e8fff1fe9efc53"}, - {file = "nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-win_amd64.whl", hash = "sha256:f7007dbd914c56bd80ea31bc43e8e149da38f68158f423ba845fc3292684e45a"}, + {file = "nvidia_cuda_nvrtc-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:82530788b8c6164a54d3fd9ae8bcca8893d397c4aeb998861982a03bbe41e204"}, + {file = "nvidia_cuda_nvrtc-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7b05ecda494c6dabc44231a608b060a71008a730d9dfda932cc508e6d29159e0"}, + {file = "nvidia_cuda_nvrtc-13.3.33-py3-none-win_amd64.whl", hash = "sha256:7d2af818851c0c224d5f92221e9226e51ee23c236df4b51f9194563979c888be"}, ] [[package]] @@ -2429,7 +2561,7 @@ description = "NVRTC native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994"}, {file = "nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fc1fec1e1637854b4c0a65fb9a8346b51dd9ee69e61ebaccc82058441f15bce8"}, @@ -2437,19 +2569,17 @@ files = [ ] [[package]] -name = "nvidia-cuda-runtime-cu12" -version = "12.6.77" +name = "nvidia-cuda-runtime" +version = "13.0.96" description = "CUDA Runtime native Libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6116fad3e049e04791c0256a9778c16237837c08b27ed8c8401e2e45de8d60cd"}, - {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d461264ecb429c84c8879a7153499ddc7b19b5f8d84c204307491989a365588e"}, - {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ba3b56a4f896141e25e19ab287cd71e52a6a0f4b29d0d31609f60e3b4d5219b7"}, - {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8"}, - {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-win_amd64.whl", hash = "sha256:86c58044c824bf3c173c49a2dbc7a6c8b53cb4e4dca50068be0bf64e9dab3f7f"}, + {file = "nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55"}, + {file = "nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548"}, + {file = "nvidia_cuda_runtime-13.0.96-py3-none-win_amd64.whl", hash = "sha256:f79298c8a098cec150a597c8eba58ecdab96e3bdc4b9bc4f9983635031740492"}, ] [[package]] @@ -2459,7 +2589,7 @@ description = "CUDA Runtime native Libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:52bf7bbee900262ffefe5e9d5a2a69a30d97e2bc5bb6cc866688caa976966e3d"}, {file = "nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90"}, @@ -2468,56 +2598,54 @@ files = [ [[package]] name = "nvidia-cudnn-cu12" -version = "9.5.1.17" +version = "9.10.2.21" description = "cuDNN runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ - {file = "nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:9fd4584468533c61873e5fda8ca41bac3a38bcb2d12350830c69b0a96a7e4def"}, - {file = "nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:30ac3869f6db17d170e0e556dd6cc5eee02647abc31ca856634d5a40f82c15b2"}, - {file = "nvidia_cudnn_cu12-9.5.1.17-py3-none-win_amd64.whl", hash = "sha256:d7af0f8a4f3b4b9dbb3122f2ef553b45694ed9c384d5a75bab197b8eefb79ab8"}, + {file = "nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9132cc3f8958447b4910a1720036d9eff5928cc3179b0a51fb6d167c6cc87d8"}, + {file = "nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8"}, + {file = "nvidia_cudnn_cu12-9.10.2.21-py3-none-win_amd64.whl", hash = "sha256:c6288de7d63e6cf62988f0923f96dc339cea362decb1bf5b3141883392a7d65e"}, ] [package.dependencies] nvidia-cublas-cu12 = "*" [[package]] -name = "nvidia-cudnn-cu12" -version = "9.10.2.21" +name = "nvidia-cudnn-cu13" +version = "9.20.0.48" description = "cuDNN runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "(platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\" or python_version >= \"3.14\")" files = [ - {file = "nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9132cc3f8958447b4910a1720036d9eff5928cc3179b0a51fb6d167c6cc87d8"}, - {file = "nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8"}, - {file = "nvidia_cudnn_cu12-9.10.2.21-py3-none-win_amd64.whl", hash = "sha256:c6288de7d63e6cf62988f0923f96dc339cea362decb1bf5b3141883392a7d65e"}, + {file = "nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1"}, + {file = "nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0c45dd8eeb50b603f07995b1b300c62ffe6a1980482b82b3bcf94a4ca9d49304"}, + {file = "nvidia_cudnn_cu13-9.20.0.48-py3-none-win_amd64.whl", hash = "sha256:af8139732b99c0118be65ea5aac97f0d46018f8c552889e49d2fb0c6261a4a24"}, ] [package.dependencies] -nvidia-cublas-cu12 = "*" +nvidia-cublas = "*" [[package]] -name = "nvidia-cufft-cu12" -version = "11.3.0.4" +name = "nvidia-cufft" +version = "12.0.0.61" description = "CUFFT native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d16079550df460376455cba121db6564089176d9bac9e4f360493ca4741b22a6"}, - {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8510990de9f96c803a051822618d42bf6cb8f069ff3f48d93a8486efdacb48fb"}, - {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ccba62eb9cef5559abd5e0d54ceed2d9934030f51163df018532142a8ec533e5"}, - {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.whl", hash = "sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca"}, - {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-win_amd64.whl", hash = "sha256:6048ebddfb90d09d2707efb1fd78d4e3a77cb3ae4dc60e19aab6be0ece2ae464"}, + {file = "nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5"}, + {file = "nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3"}, + {file = "nvidia_cufft-12.0.0.61-py3-none-win_amd64.whl", hash = "sha256:2abce5b39d2f5ae12730fb7e5db6696533e36c26e2d3e8fd1750bdd2853364eb"}, ] [package.dependencies] -nvidia-nvjitlink-cu12 = "*" +nvidia-nvjitlink = "*" [[package]] name = "nvidia-cufft-cu12" @@ -2526,7 +2654,7 @@ description = "CUFFT native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a"}, {file = "nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74"}, @@ -2537,16 +2665,16 @@ files = [ nvidia-nvjitlink-cu12 = "*" [[package]] -name = "nvidia-cufile-cu12" -version = "1.11.1.6" +name = "nvidia-cufile" +version = "1.15.1.6" description = "cuFile GPUDirect libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ - {file = "nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc23469d1c7e52ce6c1d55253273d32c565dd22068647f3aa59b3c6b005bf159"}, - {file = "nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:8f57a0051dcf2543f6dc2b98a98cb2719c37d3cee1baba8965d57f3bbc90d4db"}, + {file = "nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44"}, + {file = "nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1"}, ] [[package]] @@ -2556,26 +2684,24 @@ description = "cuFile GPUDirect libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc"}, {file = "nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:4beb6d4cce47c1a0f1013d72e02b0994730359e17801d395bdcbf20cfb3bb00a"}, ] [[package]] -name = "nvidia-curand-cu12" -version = "10.3.7.77" +name = "nvidia-curand" +version = "10.4.0.35" description = "CURAND native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:6e82df077060ea28e37f48a3ec442a8f47690c7499bff392a5938614b56c98d8"}, - {file = "nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a42cd1344297f70b9e39a1e4f467a4e1c10f1da54ff7a85c12197f6c652c8bdf"}, - {file = "nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:99f1a32f1ac2bd134897fc7a203f779303261268a65762a623bf30cc9fe79117"}, - {file = "nvidia_curand_cu12-10.3.7.77-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:7b2ed8e95595c3591d984ea3603dd66fe6ce6812b886d59049988a712ed06b6e"}, - {file = "nvidia_curand_cu12-10.3.7.77-py3-none-win_amd64.whl", hash = "sha256:6d6d935ffba0f3d439b7cd968192ff068fafd9018dbf1b85b37261b13cfc9905"}, + {file = "nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a"}, + {file = "nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc"}, + {file = "nvidia_curand-10.4.0.35-py3-none-win_amd64.whl", hash = "sha256:65b1710aa6961d326b411e314b374290904c5ddf41dc3f766ebc3f1d7d4ca69f"}, ] [[package]] @@ -2585,7 +2711,7 @@ description = "CURAND native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dfab99248034673b779bc6decafdc3404a8a6f502462201f2f31f11354204acd"}, {file = "nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9"}, @@ -2593,25 +2719,23 @@ files = [ ] [[package]] -name = "nvidia-cusolver-cu12" -version = "11.7.1.2" +name = "nvidia-cusolver" +version = "12.0.4.66" description = "CUDA solver native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0"}, - {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9e49843a7707e42022babb9bcfa33c29857a93b88020c4e4434656a655b698c"}, - {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6"}, - {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dbbe4fc38ec1289c7e5230e16248365e375c3673c9c8bac5796e2e20db07f56e"}, - {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-win_amd64.whl", hash = "sha256:6813f9d8073f555444a8705f3ab0296d3e1cb37a16d694c5fc8b862a0d8706d7"}, + {file = "nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2"}, + {file = "nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112"}, + {file = "nvidia_cusolver-12.0.4.66-py3-none-win_amd64.whl", hash = "sha256:16515bd33a8e76bb54d024cfa068fa68d30e80fc34b9e1090813ea9362e0cb65"}, ] [package.dependencies] -nvidia-cublas-cu12 = "*" -nvidia-cusparse-cu12 = "*" -nvidia-nvjitlink-cu12 = "*" +nvidia-cublas = "*" +nvidia-cusparse = "*" +nvidia-nvjitlink = "*" [[package]] name = "nvidia-cusolver-cu12" @@ -2620,7 +2744,7 @@ description = "CUDA solver native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0"}, {file = "nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450"}, @@ -2633,23 +2757,38 @@ nvidia-cusparse-cu12 = "*" nvidia-nvjitlink-cu12 = "*" [[package]] -name = "nvidia-cusparse-cu12" -version = "12.5.4.2" +name = "nvidia-cusparse" +version = "12.6.3.3" description = "CUSPARSE native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d25b62fb18751758fe3c93a4a08eff08effedfe4edf1c6bb5afd0890fe88f887"}, - {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7aa32fa5470cf754f72d1116c7cbc300b4e638d3ae5304cfa4a638a5b87161b1"}, - {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7556d9eca156e18184b94947ade0fba5bb47d69cec46bf8660fd2c71a4b48b73"}, - {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f"}, - {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-win_amd64.whl", hash = "sha256:4acb8c08855a26d737398cba8fb6f8f5045d93f82612b4cfd84645a2332ccf20"}, + {file = "nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c"}, + {file = "nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b"}, + {file = "nvidia_cusparse-12.6.3.3-py3-none-win_amd64.whl", hash = "sha256:cbcf42feb737bd7ec15b4c0a63e62351886bd3f975027b8815d7f720a2b5ea79"}, ] [package.dependencies] -nvidia-nvjitlink-cu12 = "*" +nvidia-nvjitlink = "*" + +[[package]] +name = "nvidia-cusparse" +version = "12.8.2.51" +description = "CUSPARSE native runtime libraries" +optional = false +python-versions = ">=3" +groups = ["main"] +markers = "" +files = [ + {file = "nvidia_cusparse-12.8.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:00469fcf62c4d464a1225abd9b20864ecff35e3fbc9fb992572e83d358927755"}, + {file = "nvidia_cusparse-12.8.2.51-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:65cbcc4e37a34fca4ee7df2fd57da103593842cda1bbb4a144664ecfe59873a5"}, + {file = "nvidia_cusparse-12.8.2.51-py3-none-win_amd64.whl", hash = "sha256:2ee59291cd362038f3d40d57c7cd09b26d689f3873ae5c94b31c3270772d41b8"}, +] + +[package.dependencies] +nvidia-nvjitlink = "*" [[package]] name = "nvidia-cusparse-cu12" @@ -2658,7 +2797,7 @@ description = "CUSPARSE native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc"}, {file = "nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b"}, @@ -2670,70 +2809,70 @@ nvidia-nvjitlink-cu12 = "*" [[package]] name = "nvidia-cusparselt-cu12" -version = "0.6.3" +version = "0.7.1" description = "NVIDIA cuSPARSELt" optional = false python-versions = "*" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ - {file = "nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8371549623ba601a06322af2133c4a44350575f5a3108fb75f3ef20b822ad5f1"}, - {file = "nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:e5c8a26c36445dd2e6812f1177978a24e2d37cacce7e090f297a688d1ec44f46"}, - {file = "nvidia_cusparselt_cu12-0.6.3-py3-none-win_amd64.whl", hash = "sha256:3b325bcbd9b754ba43df5a311488fca11a6b5dc3d11df4d190c000cf1a0765c7"}, + {file = "nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8878dce784d0fac90131b6817b607e803c36e629ba34dc5b433471382196b6a5"}, + {file = "nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623"}, + {file = "nvidia_cusparselt_cu12-0.7.1-py3-none-win_amd64.whl", hash = "sha256:f67fbb5831940ec829c9117b7f33807db9f9678dc2a617fbe781cac17b4e1075"}, ] [[package]] -name = "nvidia-cusparselt-cu12" -version = "0.7.1" +name = "nvidia-cusparselt-cu13" +version = "0.8.1" description = "NVIDIA cuSPARSELt" optional = false python-versions = "*" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "(platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\" or python_version >= \"3.14\")" files = [ - {file = "nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8878dce784d0fac90131b6817b607e803c36e629ba34dc5b433471382196b6a5"}, - {file = "nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623"}, - {file = "nvidia_cusparselt_cu12-0.7.1-py3-none-win_amd64.whl", hash = "sha256:f67fbb5831940ec829c9117b7f33807db9f9678dc2a617fbe781cac17b4e1075"}, + {file = "nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f"}, + {file = "nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0"}, + {file = "nvidia_cusparselt_cu13-0.8.1-py3-none-win_amd64.whl", hash = "sha256:dccbd362f91a7b9024d1f55ee9f548ac065027ff15d8c8b0db889ab3a8f31215"}, ] [[package]] name = "nvidia-nccl-cu12" -version = "2.26.2" +version = "2.27.3" description = "NVIDIA Collective Communication Library (NCCL) Runtime" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ - {file = "nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5c196e95e832ad30fbbb50381eb3cbd1fadd5675e587a548563993609af19522"}, - {file = "nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:694cf3879a206553cc9d7dbda76b13efaf610fdb70a50cba303de1b0d1530ac6"}, + {file = "nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9ddf1a245abc36c550870f26d537a9b6087fb2e2e3d6e0ef03374c6fd19d984f"}, + {file = "nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adf27ccf4238253e0b826bce3ff5fa532d65fc42322c8bfdfaf28024c0fbe039"}, ] [[package]] -name = "nvidia-nccl-cu12" -version = "2.27.3" +name = "nvidia-nccl-cu13" +version = "2.29.7" description = "NVIDIA Collective Communication Library (NCCL) Runtime" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "(platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\" or python_version >= \"3.14\")" files = [ - {file = "nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9ddf1a245abc36c550870f26d537a9b6087fb2e2e3d6e0ef03374c6fd19d984f"}, - {file = "nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adf27ccf4238253e0b826bce3ff5fa532d65fc42322c8bfdfaf28024c0fbe039"}, + {file = "nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5"}, + {file = "nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d"}, ] [[package]] -name = "nvidia-nvjitlink-cu12" -version = "12.6.85" +name = "nvidia-nvjitlink" +version = "13.3.33" description = "Nvidia JIT LTO Library" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_nvjitlink_cu12-12.6.85-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:eedc36df9e88b682efe4309aa16b5b4e78c2407eac59e8c10a6a47535164369a"}, - {file = "nvidia_nvjitlink_cu12-12.6.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cf4eaa7d4b6b543ffd69d6abfb11efdeb2db48270d94dfd3a452c24150829e41"}, - {file = "nvidia_nvjitlink_cu12-12.6.85-py3-none-win_amd64.whl", hash = "sha256:e61120e52ed675747825cdd16febc6a0730537451d867ee58bee3853b1b13d1c"}, + {file = "nvidia_nvjitlink-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:26a6de7fb4c8fdaa7703d3dad720d6d427ddfea5c48a528fd97c11733ad830e5"}, + {file = "nvidia_nvjitlink-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ce48b37dfeb3cb1eae4cf85adacb47d7a6539ea2272870c9a3628ce275c2037e"}, + {file = "nvidia_nvjitlink-13.3.33-py3-none-win_amd64.whl", hash = "sha256:4297ee49639b4f2e07255a1d69b3acc7ab2d011bb892b403e91ac98368962e3b"}, ] [[package]] @@ -2743,7 +2882,7 @@ description = "Nvidia JIT LTO Library" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88"}, {file = "nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:adccd7161ace7261e01bb91e44e88da350895c270d23f744f0820c818b7229e7"}, @@ -2751,19 +2890,30 @@ files = [ ] [[package]] -name = "nvidia-nvtx-cu12" -version = "12.6.77" +name = "nvidia-nvshmem-cu13" +version = "3.4.5" +description = "NVSHMEM creates a global address space that provides efficient and scalable communication for NVIDIA GPU clusters." +optional = false +python-versions = ">=3" +groups = ["main"] +markers = "(platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\" or python_version >= \"3.14\")" +files = [ + {file = "nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9"}, + {file = "nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80"}, +] + +[[package]] +name = "nvidia-nvtx" +version = "13.0.85" description = "NVIDIA Tools Extension" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or python_version >= \"3.14\" and platform_machine == \"x86_64\" and platform_system == \"Linux\" and (sys_platform == \"linux\" or sys_platform == \"win32\")" files = [ - {file = "nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f44f8d86bb7d5629988d61c8d3ae61dddb2015dee142740536bc7481b022fe4b"}, - {file = "nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:adcaabb9d436c9761fca2b13959a2d237c5f9fd406c8e4b723c695409ff88059"}, - {file = "nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b90bed3df379fa79afbd21be8e04a0314336b8ae16768b58f2d34cb1d04cd7d2"}, - {file = "nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6574241a3ec5fdc9334353ab8c479fe75841dbe8f4532a8fc97ce63503330ba1"}, - {file = "nvidia_nvtx_cu12-12.6.77-py3-none-win_amd64.whl", hash = "sha256:2fb11a4af04a5e6c84073e6404d26588a34afd35379f0855a99797897efa75c0"}, + {file = "nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4"}, + {file = "nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6"}, + {file = "nvidia_nvtx-13.0.85-py3-none-win_amd64.whl", hash = "sha256:d66ea44254dd3c6eacc300047af6e1288d2269dd072b417e0adffbf479e18519"}, ] [[package]] @@ -2773,7 +2923,7 @@ description = "NVIDIA Tools Extension" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ {file = "nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7ad891da111ebafbf7e015d34879f7112832fc239ff0d7d776b6cb685274615"}, {file = "nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f"}, @@ -3769,7 +3919,7 @@ description = "Easily download, build, install, upgrade, and uninstall Python pa optional = false python-versions = ">=3.9" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" or python_version >= \"3.12\"" +markers = "platform_machine == \"arm64\" and (sys_platform == \"darwin\" or sys_platform == \"linux\") or python_version >= \"3.12\" or sys_platform == \"linux\" or platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\") and (platform_machine == \"arm64\" or platform_machine == \"x86_64\") or platform_machine == \"x86_64\" and platform_system == \"Linux\"" files = [ {file = "setuptools-80.9.0-py3-none-any.whl", hash = "sha256:062d34222ad13e0cc312a4c02d73f059e86a4acbfbdea8f8f76b28c99f306922"}, {file = "setuptools-80.9.0.tar.gz", hash = "sha256:f36b47402ecde768dbfafc46e8e4207b4360c654f1f3bb84475f0a28628fb19c"}, @@ -4028,69 +4178,6 @@ files = [ {file = "tomli-2.2.1.tar.gz", hash = "sha256:cd45e1dc79c835ce60f7404ec8119f2eb06d38b1deba146f07ced3bbc44505ff"}, ] -[[package]] -name = "torch" -version = "2.7.1" -description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration" -optional = false -python-versions = ">=3.9.0" -groups = ["main"] -markers = "python_version >= \"3.14\"" -files = [ - {file = "torch-2.7.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:a103b5d782af5bd119b81dbcc7ffc6fa09904c423ff8db397a1e6ea8fd71508f"}, - {file = "torch-2.7.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:fe955951bdf32d182ee8ead6c3186ad54781492bf03d547d31771a01b3d6fb7d"}, - {file = "torch-2.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:885453d6fba67d9991132143bf7fa06b79b24352f4506fd4d10b309f53454162"}, - {file = "torch-2.7.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:d72acfdb86cee2a32c0ce0101606f3758f0d8bb5f8f31e7920dc2809e963aa7c"}, - {file = "torch-2.7.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:236f501f2e383f1cb861337bdf057712182f910f10aeaf509065d54d339e49b2"}, - {file = "torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:06eea61f859436622e78dd0cdd51dbc8f8c6d76917a9cf0555a333f9eac31ec1"}, - {file = "torch-2.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:8273145a2e0a3c6f9fd2ac36762d6ee89c26d430e612b95a99885df083b04e52"}, - {file = "torch-2.7.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:aea4fc1bf433d12843eb2c6b2204861f43d8364597697074c8d38ae2507f8730"}, - {file = "torch-2.7.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:27ea1e518df4c9de73af7e8a720770f3628e7f667280bce2be7a16292697e3fa"}, - {file = "torch-2.7.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:c33360cfc2edd976c2633b3b66c769bdcbbf0e0b6550606d188431c81e7dd1fc"}, - {file = "torch-2.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:d8bf6e1856ddd1807e79dc57e54d3335f2b62e6f316ed13ed3ecfe1fc1df3d8b"}, - {file = "torch-2.7.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:787687087412c4bd68d315e39bc1223f08aae1d16a9e9771d95eabbb04ae98fb"}, - {file = "torch-2.7.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:03563603d931e70722dce0e11999d53aa80a375a3d78e6b39b9f6805ea0a8d28"}, - {file = "torch-2.7.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:d632f5417b6980f61404a125b999ca6ebd0b8b4bbdbb5fbbba44374ab619a412"}, - {file = "torch-2.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:23660443e13995ee93e3d844786701ea4ca69f337027b05182f5ba053ce43b38"}, - {file = "torch-2.7.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:0da4f4dba9f65d0d203794e619fe7ca3247a55ffdcbd17ae8fb83c8b2dc9b585"}, - {file = "torch-2.7.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e08d7e6f21a617fe38eeb46dd2213ded43f27c072e9165dc27300c9ef9570934"}, - {file = "torch-2.7.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:30207f672328a42df4f2174b8f426f354b2baa0b7cca3a0adb3d6ab5daf00dc8"}, - {file = "torch-2.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:79042feca1c634aaf6603fe6feea8c6b30dfa140a6bbc0b973e2260c7e79a22e"}, - {file = "torch-2.7.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:988b0cbc4333618a1056d2ebad9eb10089637b659eb645434d0809d8d937b946"}, - {file = "torch-2.7.1-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:e0d81e9a12764b6f3879a866607c8ae93113cbcad57ce01ebde63eb48a576369"}, - {file = "torch-2.7.1-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:8394833c44484547ed4a47162318337b88c97acdb3273d85ea06e03ffff44998"}, - {file = "torch-2.7.1-cp39-cp39-win_amd64.whl", hash = "sha256:df41989d9300e6e3c19ec9f56f856187a6ef060c3662fe54f4b6baf1fc90bd19"}, - {file = "torch-2.7.1-cp39-none-macosx_11_0_arm64.whl", hash = "sha256:a737b5edd1c44a5c1ece2e9f3d00df9d1b3fb9541138bee56d83d38293fb6c9d"}, -] - -[package.dependencies] -filelock = "*" -fsspec = "*" -jinja2 = "*" -networkx = "*" -nvidia-cublas-cu12 = {version = "12.6.4.1", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cuda-cupti-cu12 = {version = "12.6.80", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cuda-nvrtc-cu12 = {version = "12.6.77", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cuda-runtime-cu12 = {version = "12.6.77", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cudnn-cu12 = {version = "9.5.1.17", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cufft-cu12 = {version = "11.3.0.4", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cufile-cu12 = {version = "1.11.1.6", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-curand-cu12 = {version = "10.3.7.77", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cusolver-cu12 = {version = "11.7.1.2", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cusparse-cu12 = {version = "12.5.4.2", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-cusparselt-cu12 = {version = "0.6.3", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-nccl-cu12 = {version = "2.26.2", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-nvjitlink-cu12 = {version = "12.6.85", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -nvidia-nvtx-cu12 = {version = "12.6.77", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -setuptools = {version = "*", markers = "python_version >= \"3.12\""} -sympy = ">=1.13.3" -triton = {version = "3.3.1", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} -typing-extensions = ">=4.10.0" - -[package.extras] -opt-einsum = ["opt-einsum (>=3.3)"] -optree = ["optree (>=0.13.0)"] - [[package]] name = "torch" version = "2.8.0" @@ -4098,7 +4185,7 @@ description = "Tensors and Dynamic neural networks in Python with strong GPU acc optional = false python-versions = ">=3.9.0" groups = ["main"] -markers = "python_version < \"3.14\"" +markers = "(sys_platform != \"darwin\" and sys_platform != \"linux\" or platform_machine != \"arm64\") and sys_platform != \"linux\" and python_version <= \"3.13\"" files = [ {file = "torch-2.8.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0be92c08b44009d4131d1ff7a8060d10bafdb7ddcb7359ef8d8c5169007ea905"}, {file = "torch-2.8.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:89aa9ee820bb39d4d72b794345cccef106b574508dd17dbec457949678c76011"}, @@ -4155,6 +4242,62 @@ opt-einsum = ["opt-einsum (>=3.3)"] optree = ["optree (>=0.13.0)"] pyyaml = ["pyyaml"] +[[package]] +name = "torch" +version = "2.13.0" +description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration" +optional = false +python-versions = ">=3.10" +groups = ["main"] +markers = "platform_machine == \"arm64\" and (sys_platform == \"darwin\" or sys_platform == \"linux\") or python_version >= \"3.14\" or sys_platform == \"linux\"" +files = [ + {file = "torch-2.13.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:94f0de129916f77b8dc2c7a8eff644cfeddfe59e39c9f55e9f6e17543410281d"}, + {file = "torch-2.13.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0ab4b69f3ee03a62a002cfbf77b1ca5e88aceb4ea64cb4388bb28f638ddbb045"}, + {file = "torch-2.13.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:c78b7b4d04461855a764cf01bae9a462bb88bc93defcfa11235cbc8fdf3e12c4"}, + {file = "torch-2.13.0-cp310-cp310-win_amd64.whl", hash = "sha256:2bd30b6b730d987fa386ce3898933762c5cb8cc82eb0535211d787cc3ce2dfeb"}, + {file = "torch-2.13.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:e76f9bcecc52b8ff711239a2f7547d5353df95878ab232f0773c1d95928b92f8"}, + {file = "torch-2.13.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:092790c696a760c729fd5722835f50b9d81fd7c8f141571f3f3cf4081a8f664c"}, + {file = "torch-2.13.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:60fcdcb2f3876e21146cb4524ef06397d727ca9ad5f020818547e25075fe3cb7"}, + {file = "torch-2.13.0-cp311-cp311-win_amd64.whl", hash = "sha256:a0d8b11f16a48d60e2015d8213aa0390744cbebb98e58b62b3514dddc656e330"}, + {file = "torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027"}, + {file = "torch-2.13.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:572df8be8ffb4599c88cbd6a0726f1f854f4da65d2e3c09f0e2c2283333cd6d4"}, + {file = "torch-2.13.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:796633c4cdf0fe2cdced72d8f88f22e73dbcfce83132763162f6d4bff13b820b"}, + {file = "torch-2.13.0-cp312-cp312-win_amd64.whl", hash = "sha256:024c6cc0c1b085f2f91f20a3dc27b0471d021c31ce84b81be3afdc39f791fd9d"}, + {file = "torch-2.13.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:33449899ce5496c1b84b4853179d94fd102028ae1407314d9fb956bb79e70d09"}, + {file = "torch-2.13.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:1e09d6a722504957c694faceca843acde562786df1144ebcc5a74075ec7f6005"}, + {file = "torch-2.13.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:a3a9a21312872af8a26950b2c15680335a386a1f56ed03e780653d78b9607e9e"}, + {file = "torch-2.13.0-cp313-cp313-win_amd64.whl", hash = "sha256:49b58f1e2c52440abb6f17c28f0335fe6c6d01ad1a7f55b0183b81e4b34d64e6"}, + {file = "torch-2.13.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d849b390e07d8d333ce8ecaf91b273c656c598379a19c9acf1318a883f6b391c"}, + {file = "torch-2.13.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:a3893dc2da0a972a8ca5d698c85a9f967559ac5f8ee1797b77408aa8734d073c"}, + {file = "torch-2.13.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:49f1ea385c754e54919408a9bb3b5a72b0b755bbe2c916c1d6f70afbec4908a2"}, + {file = "torch-2.13.0-cp314-cp314-win_amd64.whl", hash = "sha256:4f8573e3ce9ebcd53fe922f01077a6085ccdfbe5f12fd215883a9d87d7a744fd"}, + {file = "torch-2.13.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:c28def70706c2f9ecc752574766e8ae4da9b810ab6676b611166761a78a9f1e1"}, + {file = "torch-2.13.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:31061ff56ed8fbf26c749806905aeb749ebeb819810fd5d52508aa5afd90dddc"}, + {file = "torch-2.13.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:cc26eead4cf51d0b544e31e364dcf000846549c273bd148936fe9d24d29acb92"}, + {file = "torch-2.13.0-cp314-cp314t-win_amd64.whl", hash = "sha256:a7de8a313090dc5c7d7ba4bfe5c3be222528f9a4dba1acc83bddb1157360c4b8"}, +] + +[package.dependencies] +cuda-bindings = {version = ">=13.0.3,<14", markers = "platform_system == \"Linux\" and python_version < \"3.15\""} +cuda-toolkit = {version = "13.0.3", extras = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], markers = "platform_system == \"Linux\""} +filelock = "*" +fsspec = ">=0.8.5" +jinja2 = "*" +networkx = ">=2.5.1" +nvidia-cudnn-cu13 = {version = "9.20.0.48", markers = "platform_system == \"Linux\""} +nvidia-cusparselt-cu13 = {version = "0.8.1", markers = "platform_system == \"Linux\""} +nvidia-nccl-cu13 = {version = "2.29.7", markers = "platform_system == \"Linux\""} +nvidia-nvshmem-cu13 = {version = "3.4.5", markers = "platform_system == \"Linux\""} +setuptools = ">=77.0.3" +sympy = ">=1.13.3" +triton = {version = "3.7.1", markers = "platform_system == \"Linux\" and python_version < \"3.15\""} +typing-extensions = ">=4.10.0" + +[package.extras] +opt-einsum = ["opt-einsum (>=3.3)"] +optree = ["optree (>=0.13.0)"] +pyyaml = ["pyyaml"] + [[package]] name = "torch-directml" version = "0.1.13.dev221216" @@ -4173,50 +4316,6 @@ files = [ {file = "torch_directml-0.1.13.dev221216-cp39-cp39-win_amd64.whl", hash = "sha256:c140f0170a864d53f6a7bbbc7ef6e831ba9c05bef42ec31e5f69f056e96ba0f1"}, ] -[[package]] -name = "torchvision" -version = "0.22.1" -description = "image and video datasets and models for torch deep learning" -optional = false -python-versions = ">=3.9" -groups = ["main"] -markers = "python_version >= \"3.14\"" -files = [ - {file = "torchvision-0.22.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:3b47d8369ee568c067795c0da0b4078f39a9dfea6f3bc1f3ac87530dfda1dd56"}, - {file = "torchvision-0.22.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:990de4d657a41ed71680cd8be2e98ebcab55371f30993dc9bd2e676441f7180e"}, - {file = "torchvision-0.22.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:3347f690c2eed6d02aa0edfb9b01d321e7f7cf1051992d96d8d196c39b881d49"}, - {file = "torchvision-0.22.1-cp310-cp310-win_amd64.whl", hash = "sha256:86ad938f5a6ca645f0d5fb19484b1762492c2188c0ffb05c602e9e9945b7b371"}, - {file = "torchvision-0.22.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4addf626e2b57fc22fd6d329cf1346d474497672e6af8383b7b5b636fba94a53"}, - {file = "torchvision-0.22.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:8b4a53a6067d63adba0c52f2b8dd2290db649d642021674ee43c0c922f0c6a69"}, - {file = "torchvision-0.22.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:b7866a3b326413e67724ac46f1ee594996735e10521ba9e6cdbe0fa3cd98c2f2"}, - {file = "torchvision-0.22.1-cp311-cp311-win_amd64.whl", hash = "sha256:bb3f6df6f8fd415ce38ec4fd338376ad40c62e86052d7fc706a0dd51efac1718"}, - {file = "torchvision-0.22.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:153f1790e505bd6da123e21eee6e83e2e155df05c0fe7d56347303067d8543c5"}, - {file = "torchvision-0.22.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:964414eef19459d55a10e886e2fca50677550e243586d1678f65e3f6f6bac47a"}, - {file = "torchvision-0.22.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:699c2d70d33951187f6ed910ea05720b9b4aaac1dcc1135f53162ce7d42481d3"}, - {file = "torchvision-0.22.1-cp312-cp312-win_amd64.whl", hash = "sha256:75e0897da7a8e43d78632f66f2bdc4f6e26da8d3f021a7c0fa83746073c2597b"}, - {file = "torchvision-0.22.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9c3ae3319624c43cc8127020f46c14aa878406781f0899bb6283ae474afeafbf"}, - {file = "torchvision-0.22.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:4a614a6a408d2ed74208d0ea6c28a2fbb68290e9a7df206c5fef3f0b6865d307"}, - {file = "torchvision-0.22.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:7ee682be589bb1a002b7704f06b8ec0b89e4b9068f48e79307d2c6e937a9fdf4"}, - {file = "torchvision-0.22.1-cp313-cp313-win_amd64.whl", hash = "sha256:2566cafcfa47ecfdbeed04bab8cef1307c8d4ef75046f7624b9e55f384880dfe"}, - {file = "torchvision-0.22.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:043d9e35ed69c2e586aff6eb9e2887382e7863707115668ac9d140da58f42cba"}, - {file = "torchvision-0.22.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:27142bcc8a984227a6dcf560985e83f52b82a7d3f5fe9051af586a2ccc46ef26"}, - {file = "torchvision-0.22.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:ef46e065502f7300ad6abc98554131c35dc4c837b978d91306658f1a65c00baa"}, - {file = "torchvision-0.22.1-cp313-cp313t-win_amd64.whl", hash = "sha256:7414eeacfb941fa21acddcd725f1617da5630ec822e498660a4b864d7d998075"}, - {file = "torchvision-0.22.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:8be941b4d35c0aba819be70fdbbbed8ceb60401ce6996b8cfaaba1300ce62263"}, - {file = "torchvision-0.22.1-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:154a2bdc37a16122c2024f2f77e65f5986020b40c013515c694b5d357fac99a1"}, - {file = "torchvision-0.22.1-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:ef7dee376f42900c0e7b0e34624f391d9ece70ab90ee74b42de0c1fffe371284"}, - {file = "torchvision-0.22.1-cp39-cp39-win_amd64.whl", hash = "sha256:e01631046fda25a1eca2f58d5fdc9a152b93740eb82435cdb27c5151b8d20c02"}, -] - -[package.dependencies] -numpy = "*" -pillow = ">=5.3.0,<8.3.dev0 || >=8.4.dev0" -torch = "2.7.1" - -[package.extras] -gdown = ["gdown (>=4.7.3)"] -scipy = ["scipy"] - [[package]] name = "torchvision" version = "0.23.0" @@ -4224,7 +4323,7 @@ description = "image and video datasets and models for torch deep learning" optional = false python-versions = ">=3.9" groups = ["main"] -markers = "python_version < \"3.14\"" +markers = "(sys_platform != \"darwin\" and sys_platform != \"linux\" or platform_machine != \"arm64\") and sys_platform != \"linux\" and python_version <= \"3.13\"" files = [ {file = "torchvision-0.23.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:7266871daca00ad46d1c073e55d972179d12a58fa5c9adec9a3db9bbed71284a"}, {file = "torchvision-0.23.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:31c583ba27426a3a04eca8c05450524105c1564db41be6632f7536ef405a6de2"}, @@ -4261,6 +4360,50 @@ torch = "2.8.0" gdown = ["gdown (>=4.7.3)"] scipy = ["scipy"] +[[package]] +name = "torchvision" +version = "0.28.0" +description = "image and video datasets and models for torch deep learning" +optional = false +python-versions = "!=3.14.1,>=3.10" +groups = ["main"] +markers = "platform_machine == \"arm64\" and (sys_platform == \"darwin\" or sys_platform == \"linux\") or python_version >= \"3.14\" or sys_platform == \"linux\"" +files = [ + {file = "torchvision-0.28.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:2a1ef4b6f4bf5828b48cfad97372c8982db906830884b2868ba5c3df937a7d81"}, + {file = "torchvision-0.28.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:546fd85345cf8652f6cd099d4f9884b0ca5c2f3fae78689a21dd2f35ea6b622f"}, + {file = "torchvision-0.28.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:6dfb0f45e2b4ceb4e76f158c3fbb5f44387099f3c466e3423a09ab665a194aba"}, + {file = "torchvision-0.28.0-cp310-cp310-win_amd64.whl", hash = "sha256:7fad44dc9582570c7d92c4487d36ac46998f40cc39b438e8b8f5111a935ce4e8"}, + {file = "torchvision-0.28.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:83fe6c020866a85acd7d97deccc45ff11d66daf42916d04396a4309c66c0ccb8"}, + {file = "torchvision-0.28.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:5a38bc6da3d72621be003400b66f66a2b4c6d644fde05f680c2cb7ca8cf8dd6c"}, + {file = "torchvision-0.28.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:7e80f543b22503d9415e126db5f0ff3917036925e38560ee6b9ae38c571a4002"}, + {file = "torchvision-0.28.0-cp311-cp311-win_amd64.whl", hash = "sha256:9a45ea67235d965ef52187130d20002a4de20c54ea3d927a24286961d268dc37"}, + {file = "torchvision-0.28.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:e9f54c30cd52e3ef7fd034cc69b7bb7e0964e1c8f8743e018ab92e95b40f9eee"}, + {file = "torchvision-0.28.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5cf78ebc401ce64ae19b8c55de866bb836797d559a4de9c25ccbe74cfa642d3a"}, + {file = "torchvision-0.28.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:028a3d481b37d785605620d7cdad897064c5a55bae2aa1f2658766333e291940"}, + {file = "torchvision-0.28.0-cp312-cp312-win_amd64.whl", hash = "sha256:87dc16b2df427c1318ad335f1e2be2b3b15b2cf20f7934c83b0505a48425ee5d"}, + {file = "torchvision-0.28.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:d483b4aa3f5237569053f749cd1a2b5bb548ca456e40461a5dd087f21149d123"}, + {file = "torchvision-0.28.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:bb6dd6918460ed89cc7644adcc2402991474d6933cf1ce92b390641cb233fddf"}, + {file = "torchvision-0.28.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:ad7b3a439265cc3739a4ab5b4c998c0e38ea99c0ee7ca4dea35c5d0b099ec237"}, + {file = "torchvision-0.28.0-cp313-cp313-win_amd64.whl", hash = "sha256:7e9dd6f60d6e15f8dc27d4f877fdb6002fc70d70272412135f1c2ff9cfa08d3b"}, + {file = "torchvision-0.28.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:3bd9dba55224a9db4a2d77f6feaa5651770d8c8e86d3d0ddb0fa6bec54c8712b"}, + {file = "torchvision-0.28.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:89f90e29b0966352811b12589f3a3c61943bf2bb9487b9d7bbec10efb1096bb5"}, + {file = "torchvision-0.28.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:36beb0782976906069ca03d4c9aacaf4b6b838b06ed6c20960ea9c51cce7acdd"}, + {file = "torchvision-0.28.0-cp314-cp314-win_amd64.whl", hash = "sha256:3557cc7b539f46dabcda2b6f2b14017ccbeef024de466d4fc5835fc3f287f769"}, + {file = "torchvision-0.28.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:09ce8f56e81f19b9c378ae7bb109f83f6659fd8bc3cd14241a48e4af46e9ed49"}, + {file = "torchvision-0.28.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:62c7d110f86a039245b587e4fae60278c649f3bd42ff79cfbc1178eca4e72542"}, + {file = "torchvision-0.28.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:904cf89af220f8c6b2ed0296bb5065b474ce43b77558e48b2bf9de8b0ba17204"}, + {file = "torchvision-0.28.0-cp314-cp314t-win_amd64.whl", hash = "sha256:46f581979c010ad6da6bd85ee602aa707e1ff44312670223b7a0ee517ad06d47"}, +] + +[package.dependencies] +numpy = "*" +pillow = ">=5.3.0,<8.3.dev0 || >=8.4.dev0" +torch = "2.13.0" + +[package.extras] +gdown = ["gdown (>=4.7.3)"] +scipy = ["scipy"] + [[package]] name = "tqdm" version = "4.67.1" @@ -4285,49 +4428,52 @@ telegram = ["requests"] [[package]] name = "triton" -version = "3.3.1" +version = "3.4.0" description = "A language and compiler for custom Deep Learning operations" optional = false -python-versions = "*" +python-versions = "<3.14,>=3.9" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version >= \"3.14\"" +markers = "python_version <= \"3.13\" and platform_system == \"Linux\" and platform_machine == \"x86_64\" and sys_platform != \"linux\"" files = [ - {file = "triton-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b74db445b1c562844d3cfad6e9679c72e93fdfb1a90a24052b03bb5c49d1242e"}, - {file = "triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b"}, - {file = "triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9999e83aba21e1a78c1f36f21bce621b77bcaa530277a50484a7cb4a822f6e43"}, - {file = "triton-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b89d846b5a4198317fec27a5d3a609ea96b6d557ff44b56c23176546023c4240"}, - {file = "triton-3.3.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3198adb9d78b77818a5388bff89fa72ff36f9da0bc689db2f0a651a67ce6a42"}, - {file = "triton-3.3.1-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f6139aeb04a146b0b8e0fbbd89ad1e65861c57cfed881f21d62d3cb94a36bab7"}, + {file = "triton-3.4.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7ff2785de9bc02f500e085420273bb5cc9c9bb767584a4aa28d6e360cec70128"}, + {file = "triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b70f5e6a41e52e48cfc087436c8a28c17ff98db369447bcaff3b887a3ab4467"}, + {file = "triton-3.4.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:31c1d84a5c0ec2c0f8e8a072d7fd150cab84a9c239eaddc6706c081bfae4eb04"}, + {file = "triton-3.4.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00be2964616f4c619193cb0d1b29a99bd4b001d7dc333816073f92cf2a8ccdeb"}, + {file = "triton-3.4.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7936b18a3499ed62059414d7df563e6c163c5e16c3773678a3ee3d417865035d"}, + {file = "triton-3.4.0-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:98e5c1442eaeabae2e2452ae765801bd53cd4ce873cab0d1bdd59a32ab2d9397"}, ] [package.dependencies] setuptools = ">=40.8.0" [package.extras] -build = ["cmake (>=3.20)", "lit"] +build = ["cmake (>=3.20,<4.0)", "lit"] tests = ["autopep8", "isort", "llnl-hatchet", "numpy", "pytest", "pytest-forked", "pytest-xdist", "scipy (>=1.7.1)"] tutorials = ["matplotlib", "pandas", "tabulate"] [[package]] name = "triton" -version = "3.4.0" +version = "3.7.1" description = "A language and compiler for custom Deep Learning operations" optional = false -python-versions = "<3.14,>=3.9" +python-versions = "<3.15,>=3.10" groups = ["main"] -markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.14\"" -files = [ - {file = "triton-3.4.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7ff2785de9bc02f500e085420273bb5cc9c9bb767584a4aa28d6e360cec70128"}, - {file = "triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b70f5e6a41e52e48cfc087436c8a28c17ff98db369447bcaff3b887a3ab4467"}, - {file = "triton-3.4.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:31c1d84a5c0ec2c0f8e8a072d7fd150cab84a9c239eaddc6706c081bfae4eb04"}, - {file = "triton-3.4.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00be2964616f4c619193cb0d1b29a99bd4b001d7dc333816073f92cf2a8ccdeb"}, - {file = "triton-3.4.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7936b18a3499ed62059414d7df563e6c163c5e16c3773678a3ee3d417865035d"}, - {file = "triton-3.4.0-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:98e5c1442eaeabae2e2452ae765801bd53cd4ce873cab0d1bdd59a32ab2d9397"}, +markers = "python_version < \"3.15\" and (platform_machine == \"arm64\" or python_version == \"3.14\" or sys_platform == \"linux\") and platform_system == \"Linux\" and (sys_platform == \"darwin\" or sys_platform == \"linux\" or python_version == \"3.14\")" +files = [ + {file = "triton-3.7.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3daf64305d6cea88d3334c65ebc9bcd0c64c9564a977084366aa768d57cbcf64"}, + {file = "triton-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ee89fbf782ec2ad50391dd1cf26cbea4f4467154c37f4773026da8fc31c0f58e"}, + {file = "triton-3.7.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4a0e1cd4c4a76370ed74a8432a53cea28716827d19e40ffc732233e35ceb3f6"}, + {file = "triton-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6744957e9fd610a29680ec2346057d0c86948ed3812468670719f391e94b44a5"}, + {file = "triton-3.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9497f2e696ee368862a181a90b2dcc03ca978cc4f602abd67c7d81022a6988e1"}, + {file = "triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e40869937a68206ec70d7f25bb7ec6433cb083f9135e1f36dbd318dc449a728"}, + {file = "triton-3.7.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdbfc09d9ec58bc5e68321525653220de7515c199e7a8097a97c85e62b52cd0a"}, + {file = "triton-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:58c0e131da05134a2a4788ccbcc0c1105cf0f54c8e98f19e34cd465396dc15eb"}, + {file = "triton-3.7.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe4ea396a06171f1f1f58cbd39c70b09294398f7dd7c620939bab54ad6f934fa"}, + {file = "triton-3.7.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2020153b08280415ec0da6607834e79166442147e78e144df06b508c75b186d2"}, + {file = "triton-3.7.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c58e4c61f0c73b5dba3b5d19b4a7093c32f90dc18b2a7f121a7c16ccd31107b7"}, + {file = "triton-3.7.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10ba85fa2cca4a2fbdeb36bf1cb082f2c252bda55bf9fccd74f65ec5bc647e68"}, ] -[package.dependencies] -setuptools = ">=40.8.0" - [package.extras] build = ["cmake (>=3.20,<4.0)", "lit"] tests = ["autopep8", "isort", "llnl-hatchet", "numpy", "pytest", "pytest-forked", "pytest-xdist", "scipy (>=1.7.1)"] @@ -4370,5 +4516,5 @@ gpu = ["onnxruntime-gpu"] [metadata] lock-version = "2.1" -python-versions = ">=3.10" -content-hash = "bcd472b09d517e01e462d3e1d09045637f71bec5bfd4bcd2e4f87d47d6d4cd44" +python-versions = ">=3.10,!=3.14.1" +content-hash = "191745b98ea3a5a70729fa29797229a680d1c9dbb1d878029bbc7080a66ad7ae" diff --git a/pyproject.toml b/pyproject.toml index c6f3ad9..57fb409 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,9 +1,57 @@ +# Poetry 2's hybrid PEP 621 layout publishes broad runtime dependency markers +# while [tool.poetry] keeps the narrower cross-platform lock used by contributors. [build-system] -requires = ["poetry-core"] +requires = ["poetry-core>=2.0.0"] build-backend = "poetry.core.masonry.api" -[tool.poetry] +[project] name = "audio-separator" +requires-python = ">=3.10,!=3.14.1" +dependencies = [ + "audioop-lts>=0.2.1; python_version >= '3.13' and python_version < '4.0'", + "beartype>=0.18.5,<0.19.0", + "diffq>=0.2; sys_platform != 'win32'", + "diffq-fixed>=0.2; sys_platform == 'win32'", + "einops>=0.7", + "julius>=0.2", + "librosa>=0.10", + "ml_collections", + "numpy>=2", + "onnx-weekly", + "onnx2torch-py313>=1.6", + "packaging", + "pydub>=0.25", + "pyyaml", + "requests>=2", + "resampy>=0.4", + "rotary-embedding-torch>=0.6.1,<0.7.0", + "samplerate==0.1.0", + "scipy>=1.13.0,<2.0.0", + "six>=1.16", + "soundfile>=0.12", + "torch>=2.13,<3; sys_platform == 'darwin' and platform_machine == 'arm64'", + "torch>=2.3,<3; sys_platform != 'darwin' or platform_machine != 'arm64'", + "tqdm", +] +dynamic = [ + "authors", + "classifiers", + "description", + "keywords", + "license", + "readme", + "scripts", + "urls", + "version", +] + +[project.optional-dependencies] +cpu = ["onnxruntime>=1.17"] +gpu = ["onnxruntime-gpu>=1.17"] +dml = ["onnxruntime-directml>=1.17", "torch_directml"] + +[tool.poetry] +requires-poetry = ">=2.0.0" version = "0.44.5" description = "Easy to use audio stem separation, using various models from UVR trained primarily by @Anjok07" authors = ["Andrew Beveridge "] @@ -30,38 +78,17 @@ classifiers = [ ] [tool.poetry.dependencies] -python = ">=3.10" -requests = ">=2" -numpy = ">=2" -librosa = ">=0.10" -samplerate = "0.1.0" -six = ">=1.16" -torch = ">=2.3" -torch_directml = {version = "*", optional = true} -tqdm = "*" -pydub = ">=0.25" -audioop-lts = { version = ">=0.2.1", python = "^3.13" } -onnx-weekly = { version = "*" } -onnx2torch-py313 = ">=1.6" -onnxruntime = { version = ">=1.17", optional = true } -onnxruntime-gpu = { version = ">=1.17", optional = true } -onnxruntime-directml = { version = ">=1.17", optional = true } -julius = ">=0.2" -diffq-fixed = { version = ">=0.2", platform = "win32" } -diffq = { version = ">=0.2", platform = "!=win32" } -einops = ">=0.7" -pyyaml = "*" -ml_collections = "*" -resampy = ">=0.4" -beartype = "^0.18.5" -rotary-embedding-torch = "^0.6.1" -scipy = "^1.13.0" -soundfile = ">=0.12" - -[tool.poetry.extras] -cpu = ["onnxruntime"] -gpu = ["onnxruntime-gpu"] -dml = ["onnxruntime-directml", "torch_directml"] +python = ">=3.10,!=3.14.1" +# PyTorch 2.13 is the measured baseline on Apple Silicon and Linux/CUDA; its +# Apple Silicon wheels target macOS 14 or newer. Preserve upstream's PyTorch +# 2.8 / torchvision 0.23 lock on unvalidated Python <3.14 platforms. +# Python 3.14 needs the first pair with cp314 wheels; torchvision 0.28 excludes +# Python 3.14.1, which is mirrored by the project Python constraint above. +torch = [ + { version = ">=2.13,<3", markers = "sys_platform == 'linux' or (sys_platform == 'darwin' and platform_machine == 'arm64')" }, + { version = ">=2.3,<2.9", markers = "sys_platform != 'linux' and (sys_platform != 'darwin' or platform_machine != 'arm64') and python_version < '3.14'" }, + { version = ">=2.13,<3", markers = "sys_platform != 'linux' and (sys_platform != 'darwin' or platform_machine != 'arm64') and python_version >= '3.14'" }, +] [tool.poetry.scripts] audio-separator = 'audio_separator.utils.cli:main' diff --git a/tests/unit/test_bs_roformer_fp16.py b/tests/unit/test_bs_roformer_fp16.py new file mode 100644 index 0000000..6ef6cae --- /dev/null +++ b/tests/unit/test_bs_roformer_fp16.py @@ -0,0 +1,245 @@ +import copy +import platform +import subprocess +from unittest.mock import Mock, patch + +import pytest +import torch +from rotary_embedding_torch import RotaryEmbedding + +from audio_separator.separator.architectures.mdxc_separator import MDXCSeparator +from audio_separator.separator.execution_policy import NATIVE_FP16 +from audio_separator.separator.roformer.roformer_loader import RoformerLoader +from audio_separator.separator.uvr_lib_v5.roformer import bs_roformer as bs_module + + +def _apple_gpu_is_virtualized() -> bool: + """Detect a paravirtualized Metal device (hosted CI Macs report VirtualMac*).""" + if platform.system() != "Darwin": + return False + try: + result = subprocess.run(["/usr/sbin/sysctl", "-n", "hw.model"], capture_output=True, text=True, timeout=5, check=False) + except (OSError, subprocess.TimeoutExpired): + return False + if result.returncode != 0: + return False + return result.stdout.strip().startswith("VirtualMac") + + +def _tiny_bs_roformer(*, linear_transformer_depth=0): + torch.manual_seed(0) + return bs_module.BSRoformer( + dim=16, + depth=1, + stereo=False, + num_stems=1, + time_transformer_depth=1, + freq_transformer_depth=1, + linear_transformer_depth=linear_transformer_depth, + freqs_per_bands=(17, 16), # sums to 64 // 2 + 1 + dim_head=8, + heads=2, + flash_attn=False, + stft_n_fft=64, + stft_hop_length=16, + stft_win_length=64, + mask_estimator_depth=1, + ).eval() + + +def _half_preserving_rotary_frequencies(model): + rotary_frequencies = [module.freqs.detach().clone() for module in model.modules() if isinstance(module, RotaryEmbedding)] + + model.half() + + rotary_modules = [module for module in model.modules() if isinstance(module, RotaryEmbedding)] + for rotary, frequencies in zip(rotary_modules, rotary_frequencies): + rotary.freqs.data = frequencies.to(rotary.freqs.device) + rotary.cached_freqs = None + + return model + + +def _non_silent_audio(): + sample_indices = torch.arange(256, dtype=torch.float32) + return ( + 0.35 * torch.sin(2 * torch.pi * 440 * sample_indices / 44100) + 0.15 * torch.sin(2 * torch.pi * 880 * sample_indices / 44100) + ).unsqueeze(0) + + +def test_bs_rms_norm_keeps_half_silence_finite_and_zero_on_cpu(): + norm = bs_module.RMSNorm(16).half() + + output = norm(torch.zeros(2, 4, 16, dtype=torch.float16)) + + assert output.dtype == torch.float16 + assert torch.isfinite(output).all() + assert torch.count_nonzero(output) == 0 + + +def test_bs_linear_attention_normalization_keeps_half_silence_finite_on_cpu(): + output = bs_module.l2norm(torch.zeros(2, 4, 16, dtype=torch.float16)) + + assert output.dtype == torch.float16 + assert torch.isfinite(output).all() + assert torch.count_nonzero(output) == 0 + + +@pytest.mark.parametrize("flash", [False, True]) +def test_bs_linear_attention_uses_its_configured_similarity_scale(flash): + torch.manual_seed(0) + q = torch.randn(1, 2, 4, 8) + k = torch.randn(1, 2, 4, 8) + v = torch.randn(1, 2, 4, 8) + scale = 8.0 + attend = bs_module.Attend(scale=scale, flash=flash) + + similarity = torch.einsum("b h i d, b h j d -> b h i j", q, k) * scale + expected = torch.einsum("b h i j, b h j d -> b h i d", similarity.softmax(dim=-1), v) + + torch.testing.assert_close(attend(q, k, v), expected) + + +def test_bs_half_forward_keeps_silence_finite_and_zero_on_cpu(): + model = _half_preserving_rotary_frequencies(_tiny_bs_roformer()) + + with torch.no_grad(): + output = model(torch.zeros(1, 256)) + + assert output.shape == (1, 1, 256) + assert output.dtype == torch.float32 + assert torch.isfinite(output).all() + assert torch.count_nonzero(output) == 0 + + +def test_bs_native_fp16_is_selected_and_applied_on_mps(): + separator = object.__new__(MDXCSeparator) + separator.logger = Mock() + separator.model_run = _tiny_bs_roformer() + separator.roformer_model_type = "bs_roformer" + separator.torch_device = torch.device("mps") + separator.requested_torch_device = separator.torch_device + separator.use_autocast = False + separator.use_native_fp16 = True + separator.use_torch_compile = False + + separator._configure_model_precision() + + assert separator.effective_precision == NATIVE_FP16 + assert separator.is_native_fp16 is True + assert next(separator.model_run.band_split.parameters()).dtype == torch.float16 + + +def test_bs_half_forward_keeps_silence_finite_in_forced_dml_complex_fallback(): + model = _half_preserving_rotary_frequencies(_tiny_bs_roformer()) + + with patch.object(bs_module, "_is_dml_device", return_value=True), torch.no_grad(): + output = model(torch.zeros(1, 256)) + + assert output.shape == (1, 1, 256) + assert output.dtype == torch.float32 + assert torch.isfinite(output).all() + assert torch.count_nonzero(output) == 0 + + +def test_regional_compile_target_flattening_handles_two_and_three_transformer_bs_blocks(): + separator = object.__new__(MDXCSeparator) + two_transformer_block = torch.nn.ModuleList([torch.nn.Identity(), torch.nn.Identity()]) + three_transformer_block = torch.nn.ModuleList([torch.nn.Identity(), torch.nn.Identity(), torch.nn.Identity()]) + separator.model_run = torch.nn.Module() + separator.model_run.layers = torch.nn.ModuleList([two_transformer_block, three_transformer_block]) + + targets = separator._regional_compile_targets() + + assert targets == [*two_transformer_block, *three_transformer_block] + + +def test_roformer_loader_forwards_linear_transformer_depth_to_bs_model(): + loader = RoformerLoader() + model = Mock() + config = { + "dim": 16, + "depth": 1, + "linear_transformer_depth": 2, + "freqs_per_bands": (17, 16), + } + + with patch.object(bs_module, "BSRoformer", return_value=model) as bs_roformer: + loaded = loader._create_bs_roformer(config) + + assert loaded is model + assert bs_roformer.call_args.kwargs["linear_transformer_depth"] == 2 + + +def test_bs_linear_attention_block_runs_in_half_and_is_included_in_regional_compile(): + model = _half_preserving_rotary_frequencies(_tiny_bs_roformer(linear_transformer_depth=1)) + separator = object.__new__(MDXCSeparator) + separator.logger = Mock() + separator.model_run = model + separator.roformer_model_type = "bs_roformer" + separator.torch_device = torch.device("cpu") + separator._should_torch_compile = True + + with patch.object(torch.nn.Module, "compile", autospec=True) as compile_module: + separator._configure_model_compilation() + + assert len(model.layers[0]) == 3 + assert compile_module.call_count == 3 + assert separator.effective_torch_compile is True + + with torch.no_grad(): + output = model(torch.zeros(1, 256)) + + assert output.shape == (1, 1, 256) + assert output.dtype == torch.float32 + assert torch.isfinite(output).all() + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available") +@pytest.mark.parametrize("force_cpu_complex", [False, True]) +def test_bs_half_mps_forward_keeps_silence_finite_and_zero(force_cpu_complex): + model = _half_preserving_rotary_frequencies(_tiny_bs_roformer()).to("mps") + audio = torch.zeros(1, 256, device="mps") + original_stft = torch.stft + + with ( + patch.object(bs_module, "should_fallback_to_cpu_for_complex_ops", return_value=force_cpu_complex), + patch.object(bs_module.torch, "stft", wraps=original_stft) as stft, + torch.no_grad(), + ): + output = model(audio) + + assert output.shape == (1, 1, 256) + assert output.device.type == "mps" + assert torch.isfinite(output).all() + assert torch.count_nonzero(output) == 0 + assert stft.call_args.args[0].device.type == ("cpu" if force_cpu_complex else "mps") + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available") +@pytest.mark.skipif( + _apple_gpu_is_virtualized(), + reason="fp16 SNR gates need a real Apple GPU: virtualized Metal degrades half-precision accumulation", +) +@pytest.mark.parametrize("force_cpu_complex", [False, True]) +def test_bs_half_mps_forward_matches_cpu_float32(force_cpu_complex): + cpu_model = _tiny_bs_roformer() + mps_model = _half_preserving_rotary_frequencies(copy.deepcopy(cpu_model)).to("mps") + audio = _non_silent_audio() + + with torch.no_grad(): + reference = cpu_model(audio) + with patch.object(bs_module, "should_fallback_to_cpu_for_complex_ops", return_value=force_cpu_complex): + output = mps_model(audio.to("mps")) + + output = output.cpu().float() + error = output - reference + reference_rms = reference.square().mean().sqrt() + error_rms = error.square().mean().sqrt() + snr = 20 * torch.log10(reference_rms / error_rms) + + assert output.shape == reference.shape + assert torch.isfinite(output).all() + assert reference_rms.item() > 1e-4 + assert snr.item() > 30 + torch.testing.assert_close(output, reference, rtol=0.1, atol=1e-3) diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index ce49eef..7280a1b 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -41,6 +41,8 @@ def common_expected_args(): "sample_rate": 44100, "use_soundfile": False, "use_autocast": False, + "use_native_fp16": False, + "use_torch_compile": False, "use_directml": False, "chunk_duration": None, "ensemble_algorithm": None, @@ -258,6 +260,78 @@ def test_cli_use_autocast_argument(common_expected_args): mock_separator.assert_called_once_with(**expected_args) +def test_cli_use_native_fp16_argument(common_expected_args): + test_args = ["cli.py", "test_audio.mp3", "--use_native_fp16"] + with patch("sys.argv", test_args): + with patch("audio_separator.separator.Separator") as mock_separator: + mock_separator_instance = mock_separator.return_value + mock_separator_instance.separate.return_value = ["output_file.mp3"] + main() + + expected_args = common_expected_args.copy() + expected_args["use_native_fp16"] = True + + mock_separator.assert_called_once_with(**expected_args) + + +def test_cli_rejects_multiple_precision_modes(): + test_args = ["cli.py", "test_audio.mp3", "--use_autocast", "--use_native_fp16"] + + with patch("sys.argv", test_args), pytest.raises(SystemExit, match="2"): + main() + + +@pytest.mark.parametrize( + ("precision_flag", "expected_key"), + [ + ("--use_autocast", "use_autocast"), + ("--use_native_fp16", "use_native_fp16"), + ], +) +def test_cli_precision_mode_can_be_combined_with_torch_compile( + common_expected_args, + precision_flag, + expected_key, +): + test_args = ["cli.py", "test_audio.mp3", precision_flag, "--use_torch_compile"] + + with patch("sys.argv", test_args), patch("audio_separator.separator.Separator") as mock_separator: + mock_separator.return_value.separate.return_value = ["output_file.mp3"] + main() + + expected_args = common_expected_args.copy() + expected_args[expected_key] = True + expected_args["use_torch_compile"] = True + mock_separator.assert_called_once_with(**expected_args) + + +def test_cli_rejects_multiple_precision_modes_even_with_torch_compile(): + test_args = [ + "cli.py", + "test_audio.mp3", + "--use_autocast", + "--use_native_fp16", + "--use_torch_compile", + ] + + with patch("sys.argv", test_args), pytest.raises(SystemExit, match="2"): + main() + + +def test_cli_use_torch_compile_argument(common_expected_args): + test_args = ["cli.py", "test_audio.mp3", "--use_torch_compile"] + with patch("sys.argv", test_args): + with patch("audio_separator.separator.Separator") as mock_separator: + mock_separator_instance = mock_separator.return_value + mock_separator_instance.separate.return_value = ["output_file.mp3"] + main() + + expected_args = common_expected_args.copy() + expected_args["use_torch_compile"] = True + + mock_separator.assert_called_once_with(**expected_args) + + # Test using use_directml argument def test_cli_use_directml_argument(common_expected_args): test_args = ["cli.py", "test_audio.mp3", "--use_directml"] diff --git a/tests/unit/test_configuration_normalizer.py b/tests/unit/test_configuration_normalizer.py index 86a1404..273a976 100644 --- a/tests/unit/test_configuration_normalizer.py +++ b/tests/unit/test_configuration_normalizer.py @@ -184,6 +184,7 @@ def test_normalize_parameter_values_numbers(self): config = { 'dim': '512', 'depth': '12.0', # Float string to int + 'linear_transformer_depth': '1', 'sample_rate': 44100.0, # Float to int 'attn_dropout': '0.1', 'ff_dropout': 0.2 @@ -193,6 +194,7 @@ def test_normalize_parameter_values_numbers(self): assert result['dim'] == 512 assert result['depth'] == 12 + assert result['linear_transformer_depth'] == 1 assert result['sample_rate'] == 44100 assert result['attn_dropout'] == 0.1 assert result['ff_dropout'] == 0.2 @@ -200,6 +202,7 @@ def test_normalize_parameter_values_numbers(self): # Check types assert isinstance(result['dim'], int) assert isinstance(result['depth'], int) + assert isinstance(result['linear_transformer_depth'], int) assert isinstance(result['sample_rate'], int) assert isinstance(result['attn_dropout'], float) assert isinstance(result['ff_dropout'], float) diff --git a/tests/unit/test_demucs_cleanup.py b/tests/unit/test_demucs_cleanup.py new file mode 100644 index 0000000..fecbf09 --- /dev/null +++ b/tests/unit/test_demucs_cleanup.py @@ -0,0 +1,33 @@ +import logging +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch + +from audio_separator.separator.architectures.demucs_separator import DemucsSeparator + + +def test_demucs_model_is_released_after_inference_failure(): + separator = object.__new__(DemucsSeparator) + separator.logger = logging.getLogger(__name__) + separator.model_path = "/tmp/htdemucs.yaml" + separator.segment_size = "Default" + separator.torch_device = torch.device("cpu") + separator.prepare_mix = MagicMock(return_value=np.zeros((2, 16), dtype=np.float32)) + separator.demix_demucs = MagicMock(side_effect=RuntimeError("demix failed")) + separator.clear_gpu_cache = MagicMock(side_effect=RuntimeError("cleanup failed")) + + model = MagicMock(spec=torch.nn.Module) + model.sources = ["drums", "bass", "other", "vocals"] + + with ( + patch("audio_separator.separator.architectures.demucs_separator.HDemucs"), + patch("audio_separator.separator.architectures.demucs_separator.get_demucs_model", return_value=model), + patch("audio_separator.separator.architectures.demucs_separator.demucs_segments", return_value=model), + ): + with pytest.raises(RuntimeError, match="demix failed"): + separator.separate("input.wav") + + assert not hasattr(separator, "demucs_model_instance") + separator.clear_gpu_cache.assert_called_once_with() diff --git a/tests/unit/test_demucs_import.py b/tests/unit/test_demucs_import.py new file mode 100644 index 0000000..077f76c --- /dev/null +++ b/tests/unit/test_demucs_import.py @@ -0,0 +1,16 @@ +import importlib +from pathlib import Path + + +def test_checkpoint_compatible_top_level_demucs_import(monkeypatch): + """Demucs modules remain importable under checkpoint-compatible top-level names.""" + uvr_lib_path = Path(__file__).resolve().parents[2] / "audio_separator" / "separator" / "uvr_lib_v5" + monkeypatch.syspath_prepend(str(uvr_lib_path)) + + hdemucs = importlib.import_module("demucs.hdemucs") + htdemucs = importlib.import_module("demucs.htdemucs") + spec = importlib.import_module("demucs.spec") + + assert hdemucs.HDemucs is not None + assert htdemucs.HTDemucs is not None + assert spec.spectro is not None diff --git a/tests/unit/test_device_utils.py b/tests/unit/test_device_utils.py new file mode 100644 index 0000000..584b511 --- /dev/null +++ b/tests/unit/test_device_utils.py @@ -0,0 +1,237 @@ +from unittest.mock import patch + +import pytest +import torch + +from audio_separator.separator.uvr_lib_v5 import device_utils + + +@pytest.fixture(autouse=True) +def clear_device_capability_cache(): + device_utils._supports_complex_spectral_ops.cache_clear() + device_utils._AUTOCAST_SUPPORT_CACHE.clear() + yield + device_utils._supports_complex_spectral_ops.cache_clear() + device_utils._AUTOCAST_SUPPORT_CACHE.clear() + + +def test_privateuseone_is_not_treated_as_autocast_capable(): + with ( + patch.object(device_utils.torch.amp.autocast_mode, "is_autocast_available", return_value=True) as available, + patch.object(device_utils.torch, "autocast", side_effect=AssertionError("unsupported")) as autocast, + ): + assert device_utils.supports_autocast(torch.device("privateuseone")) is False + + available.assert_not_called() + autocast.assert_not_called() + + +def test_autocast_support_requires_a_working_context(): + with ( + patch.object(device_utils.torch.amp.autocast_mode, "is_autocast_available", return_value=True), + patch.object(device_utils.torch, "autocast", side_effect=AssertionError("unsupported")), + ): + assert device_utils.supports_autocast(torch.device("mps")) is False + + +def test_autocast_support_handles_torch_without_availability_probe(): + with ( + patch.object(device_utils.torch.amp.autocast_mode, "is_autocast_available", None), + patch.object(device_utils.torch, "autocast") as autocast, + ): + assert device_utils.supports_autocast(torch.device("cpu")) is True + + autocast.assert_called_once_with(device_type="cpu", enabled=False) + + +@pytest.mark.parametrize("device_type", ["cpu", "cuda"]) +def test_standard_devices_are_supported_without_a_runtime_probe(device_type): + with patch.object(device_utils.torch, "device") as device_constructor: + assert device_utils._supports_complex_spectral_ops(device_type, -1) is True + + device_constructor.assert_not_called() + + +def test_directml_short_circuits_without_a_runtime_probe(): + with patch.object(device_utils.torch, "device") as device_constructor: + assert device_utils._supports_complex_spectral_ops("privateuseone", -1) is False + + device_constructor.assert_not_called() + + +def test_probe_failure_preserves_cpu_fallback(): + with patch.object(device_utils.torch, "device", side_effect=RuntimeError("unsupported")): + assert device_utils._supports_complex_spectral_ops("mps", -1) is False + + +def test_force_cpu_environment_flag_overrides_capability_probe(monkeypatch): + monkeypatch.setenv("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX", "1") + with patch.object(device_utils, "_supports_complex_spectral_ops") as probe: + assert device_utils.should_fallback_to_cpu_for_complex_ops(torch.device("cpu")) is True + + probe.assert_not_called() + + +@pytest.mark.parametrize("value", [None, "", "0", "false", "off"]) +def test_disabled_force_cpu_environment_values_use_capability_probe(monkeypatch, value): + if value is None: + monkeypatch.delenv("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX", raising=False) + else: + monkeypatch.setenv("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX", value) + + with patch.object(device_utils, "_supports_complex_spectral_ops", return_value=True) as probe: + assert device_utils.should_fallback_to_cpu_for_complex_ops(torch.device("mps")) is False + + probe.assert_called_once_with("mps", -1) + + +def test_fallback_decision_uses_cached_capability_probe(monkeypatch): + monkeypatch.delenv("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX", raising=False) + with patch.object(device_utils, "_supports_complex_spectral_ops", return_value=True) as probe: + assert device_utils.should_fallback_to_cpu_for_complex_ops(torch.device("mps")) is False + + probe.assert_called_once_with("mps", -1) + + +@pytest.mark.parametrize( + ("device_type", "expected"), + [("cpu", False), ("cuda", False), ("privateuseone", False), ("mps", True)], +) +def test_device_accumulation_is_limited_to_mps(device_type, expected): + assert device_utils.should_accumulate_on_device(torch.device(device_type), estimated_bytes=1024) is expected + + +def _metal(recommended, driver=0): + """Patch the two torch.mps counters the budget is derived from.""" + + def reading(counter): + return recommended if counter == "recommended_max_memory" else driver + + return patch.object(device_utils, "_mps_memory_reading", side_effect=reading) + + +def test_mps_accumulation_falls_back_to_cpu_above_bounded_buffer_size(): + with _metal(recommended=0): + limit = device_utils.mps_accumulation_budget_bytes() + + assert limit == device_utils.MAX_MPS_FULL_TRACK_BUFFER_BYTES + assert device_utils.should_accumulate_on_device(torch.device("mps"), estimated_bytes=limit) is True + assert device_utils.should_accumulate_on_device(torch.device("mps"), estimated_bytes=limit + 1) is False + + +def test_accumulation_budget_is_half_of_the_free_working_set(): + recommended = 16 * 1024**3 + driver = 2 * 1024**3 + + with _metal(recommended, driver): + budget = device_utils.mps_accumulation_budget_bytes() + + assert budget == int((recommended - driver) * device_utils.MPS_BUFFER_HEADROOM_SHARE) + assert budget == 7 * 1024**3 + + +def test_accumulation_budget_shrinks_as_the_working_set_fills(): + recommended = 16 * 1024**3 + + with _metal(recommended, driver=0): + idle = device_utils.mps_accumulation_budget_bytes() + with _metal(recommended, driver=12 * 1024**3): + loaded = device_utils.mps_accumulation_budget_bytes() + + assert idle == 8 * 1024**3 + assert loaded == 2 * 1024**3 + + +def test_accumulation_budget_never_drops_below_the_floor(): + # Free room this small would put half of it under the floor. + with _metal(recommended=16 * 1024**3, driver=15 * 1024**3): + assert device_utils.mps_accumulation_budget_bytes() == device_utils.MAX_MPS_FULL_TRACK_BUFFER_BYTES + + +def test_accumulation_budget_survives_an_overcommitted_working_set(): + # driver_allocated can exceed the recommendation; free room must not go negative. + with _metal(recommended=8 * 1024**3, driver=12 * 1024**3): + assert device_utils.mps_accumulation_budget_bytes() == device_utils.MAX_MPS_FULL_TRACK_BUFFER_BYTES + + +@pytest.mark.parametrize("recommended", [0, 8 * 1024**3]) +def test_accumulation_budget_env_override_wins(monkeypatch, recommended): + monkeypatch.setenv(device_utils.MPS_BUFFER_BUDGET_ENV, "6.5") + + with _metal(recommended): + assert device_utils.mps_accumulation_budget_bytes() == int(6.5 * 1024**3) + + +@pytest.mark.parametrize("value", ["", "0", "-2", "not-a-number"]) +def test_accumulation_budget_ignores_unusable_env_override(monkeypatch, value): + monkeypatch.setenv(device_utils.MPS_BUFFER_BUDGET_ENV, value) + + with _metal(recommended=0): + assert device_utils.mps_accumulation_budget_bytes() == device_utils.MAX_MPS_FULL_TRACK_BUFFER_BYTES + + +@pytest.mark.parametrize("counter", ["recommended_max_memory", "driver_allocated_memory"]) +@pytest.mark.parametrize("failure", [AttributeError, RuntimeError, OSError, ValueError]) +def test_memory_reading_survives_a_failing_metal_query(counter, failure): + with ( + patch.object(device_utils.torch.backends.mps, "is_available", return_value=True), + patch.object(device_utils.torch.mps, counter, side_effect=failure("boom")), + ): + assert device_utils._mps_memory_reading(counter) == 0 + + +@pytest.mark.parametrize("counter", ["recommended_max_memory", "driver_allocated_memory"]) +def test_memory_reading_is_zero_without_mps(counter): + with patch.object(device_utils.torch.backends.mps, "is_available", return_value=False): + assert device_utils._mps_memory_reading(counter) == 0 + + +def test_a_larger_budget_keeps_more_inputs_on_mps(): + estimate = 3 * 1024**3 + + with _metal(recommended=0): + assert device_utils.should_accumulate_on_device(torch.device("mps"), estimate) is False + + with _metal(recommended=16 * 1024**3, driver=0): + assert device_utils.should_accumulate_on_device(torch.device("mps"), estimate) is True + + +@pytest.mark.parametrize( + ("device_type", "cac", "complex_fallback", "expected", "probe_called"), + [ + ("mps", False, False, True, False), + ("mps", True, False, False, True), + ("mps", True, True, True, True), + ("cpu", False, False, False, True), + ("cuda", False, False, False, True), + ], +) +def test_demucs_mask_fallback_covers_non_cac_mps_wiener_path( + device_type, + cac, + complex_fallback, + expected, + probe_called, +): + with patch.object( + device_utils, + "should_fallback_to_cpu_for_complex_ops", + return_value=complex_fallback, + ) as complex_probe: + assert device_utils.should_fallback_to_cpu_for_demucs_mask(torch.device(device_type), cac) is expected + + assert complex_probe.called is probe_called + + +def test_probe_rejects_a_device_without_complex_scatter_support(): + cpu_device = torch.device("cpu") + with ( + patch.object(device_utils.torch, "device", return_value=cpu_device), + patch.object(device_utils, "_probe_complex_scatter_add", side_effect=RuntimeError("unsupported")), + ): + assert device_utils._supports_complex_spectral_ops("mps", -1) is False + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available") +def test_mps_probe_returns_a_boolean(): + assert isinstance(device_utils._supports_complex_spectral_ops("mps", -1), bool) diff --git a/tests/unit/test_execution_policy.py b/tests/unit/test_execution_policy.py new file mode 100644 index 0000000..69fc685 --- /dev/null +++ b/tests/unit/test_execution_policy.py @@ -0,0 +1,207 @@ +from unittest.mock import Mock, patch + +import pytest +import torch + +from audio_separator.separator.execution_policy import AUTOCAST, FP32, NATIVE_FP16, resolve_execution_policy +from audio_separator.separator.separator import Separator + + +def _resolve(*, device="mps", requested_device=None, model="mel_band_roformer", autocast=False, native=False, compile=False, pytorch=True): + logger = Mock() + policy = resolve_execution_policy( + device=torch.device(device), + requested_device=torch.device(requested_device) if requested_device else None, + model_family=model, + use_autocast=autocast, + use_native_fp16=native, + use_torch_compile=compile, + uses_pytorch_inference=pytorch, + logger=logger, + ) + return policy, logger + + +@pytest.mark.parametrize("device", ["mps", "cuda"]) +@pytest.mark.parametrize("model", ["mel_band_roformer", "bs_roformer"]) +def test_verified_roformer_native_fp16_is_enabled_on_accelerators(device, model): + policy, logger = _resolve(device=device, model=model, native=True) + + assert policy.precision == NATIVE_FP16 + assert policy.use_torch_compile is False + logger.warning.assert_not_called() + + +@pytest.mark.parametrize("device", ["mps", "cuda"]) +@pytest.mark.parametrize("model", ["mel_band_roformer", "bs_roformer"]) +@pytest.mark.parametrize( + ("autocast", "native", "expected_precision"), + [ + (False, False, FP32), + (True, False, AUTOCAST), + (False, True, NATIVE_FP16), + ], +) +def test_verified_roformer_precision_modes_can_be_combined_with_compile( + device, + model, + autocast, + native, + expected_precision, +): + with patch("audio_separator.separator.execution_policy.supports_autocast", return_value=True): + policy, logger = _resolve( + device=device, + model=model, + autocast=autocast, + native=native, + compile=True, + ) + + assert policy.precision == expected_precision + assert policy.use_torch_compile is True + logger.warning.assert_not_called() + + +@pytest.mark.parametrize("model", ["mel_band_roformer", "bs_roformer"]) +@pytest.mark.parametrize( + ("autocast", "expected_precision"), + [ + (False, FP32), + (True, AUTOCAST), + ], +) +def test_verified_cpu_roformer_modes_can_be_combined_with_compile(model, autocast, expected_precision): + with patch("audio_separator.separator.execution_policy.supports_autocast", return_value=True): + policy, logger = _resolve( + device="cpu", + model=model, + autocast=autocast, + compile=True, + ) + + assert policy.precision == expected_precision + assert policy.use_torch_compile is True + logger.warning.assert_not_called() + + +def test_unsupported_cpu_native_fp16_falls_back_to_verified_fp32_compile(): + policy, logger = _resolve(device="cpu", native=True, compile=True) + + assert policy.precision == FP32 + assert policy.use_torch_compile is True + logger.warning.assert_called_once() + + +@pytest.mark.parametrize("model", ["vr", "demucs"]) +def test_unverified_models_reject_native_fp16_and_compile(model): + policy, logger = _resolve(model=model, native=True, compile=True) + + assert policy.precision == FP32 + assert policy.use_torch_compile is False + assert logger.warning.call_count == 2 + + +@pytest.mark.parametrize("model", ["vr", "demucs"]) +def test_unverified_models_keep_existing_autocast_but_reject_compile(model): + with patch("audio_separator.separator.execution_policy.supports_autocast", return_value=True): + policy, logger = _resolve(model=model, autocast=True, compile=True) + + assert policy.precision == AUTOCAST + assert policy.use_torch_compile is False + logger.warning.assert_called_once() + + +def test_autocast_and_compile_are_orthogonal_requests(): + with patch("audio_separator.separator.execution_policy.supports_autocast", return_value=True): + policy, logger = _resolve(autocast=True, compile=True) + + assert policy.precision == AUTOCAST + assert policy.use_torch_compile is True + logger.warning.assert_not_called() + + +def test_compile_falls_back_on_torch_versions_without_traceable_sdpa_context(): + with ( + patch("audio_separator.separator.execution_policy.supports_autocast", return_value=True), + patch("audio_separator.separator.execution_policy._regional_compile_runtime_supported", return_value=False), + ): + policy, logger = _resolve(autocast=True, compile=True) + + assert policy.precision == AUTOCAST + assert policy.use_torch_compile is False + logger.warning.assert_called_once() + assert "PyTorch 2.6 or newer" in logger.warning.call_args.args[0] + + +def test_directml_never_enables_autocast(): + with patch("audio_separator.separator.execution_policy.supports_autocast") as available: + policy, logger = _resolve(device="privateuseone", autocast=True) + + assert policy.precision == FP32 + assert policy.use_torch_compile is False + available.assert_not_called() + logger.warning.assert_called_once() + + +def test_onnx_runtime_model_does_not_report_autocast_as_effective(): + with patch("audio_separator.separator.execution_policy.supports_autocast") as available: + policy, logger = _resolve(device="mps", model="mdx", autocast=True, pytorch=False) + + assert policy.precision == FP32 + assert policy.use_torch_compile is False + available.assert_not_called() + logger.warning.assert_called_once() + + +def test_directml_cpu_fallback_still_uses_float32_eager(): + with patch("audio_separator.separator.execution_policy.supports_autocast") as available: + policy, logger = _resolve( + device="cpu", + requested_device="privateuseone", + model="bs_roformer", + autocast=True, + compile=True, + ) + + assert policy.precision == FP32 + assert policy.use_torch_compile is False + available.assert_not_called() + assert logger.warning.call_count == 2 + + +def test_directml_cpu_fallback_rejects_native_fp16_and_compile(): + policy, logger = _resolve( + device="cpu", + requested_device="privateuseone", + model="bs_roformer", + native=True, + compile=True, + ) + + assert policy.precision == FP32 + assert policy.use_torch_compile is False + assert logger.warning.call_count == 2 + + +def test_resolver_rejects_conflicting_precision_modes(): + with pytest.raises(ValueError, match="mutually exclusive"): + _resolve(autocast=True, native=True) + + +def test_constructor_rejects_conflicting_precision_modes(): + with pytest.raises(ValueError, match="mutually exclusive"): + Separator(info_only=True, use_autocast=True, use_native_fp16=True) + + +def test_effective_properties_follow_loaded_model_state(): + separator = object.__new__(Separator) + separator.model_instance = None + + assert separator.effective_precision == FP32 + assert separator.effective_torch_compile is False + + separator.model_instance = Mock(effective_precision=NATIVE_FP16, effective_torch_compile=True) + + assert separator.effective_precision == NATIVE_FP16 + assert separator.effective_torch_compile is True diff --git a/tests/unit/test_mdxc_roformer_chunk_starts.py b/tests/unit/test_mdxc_roformer_chunk_starts.py new file mode 100644 index 0000000..698e9be --- /dev/null +++ b/tests/unit/test_mdxc_roformer_chunk_starts.py @@ -0,0 +1,54 @@ +from unittest.mock import Mock + +import pytest + +from audio_separator.separator.architectures.mdxc_separator import MDXCSeparator + + +@pytest.mark.parametrize( + ("audio_length", "chunk_size", "step", "expected"), + [ + (0, 4, 2, []), + (3, 4, 2, [0]), + (4, 4, 2, [0]), + (10, 4, 3, [0, 3, 6]), + (11, 4, 3, [0, 3, 6, 7]), + (20, 8, 5, [0, 5, 10, 12]), + (20, 8, 2, [0, 2, 4, 6, 8, 10, 12]), + ], +) +def test_roformer_chunk_starts_cover_tail_once(audio_length, chunk_size, step, expected): + assert MDXCSeparator._roformer_chunk_starts(audio_length, chunk_size, step) == expected + + +@pytest.mark.parametrize(("chunk_size", "step"), [(0, 1), (-1, 1), (4, 0), (4, -1), (4, 5)]) +def test_roformer_chunk_starts_reject_invalid_chunk_schedule(chunk_size, step): + with pytest.raises(ValueError): + MDXCSeparator._roformer_chunk_starts(audio_length=10, chunk_size=chunk_size, step=step) + + +def test_roformer_chunk_starts_reject_negative_audio_length(): + with pytest.raises(ValueError): + MDXCSeparator._roformer_chunk_starts(audio_length=-1, chunk_size=4, step=2) + + +def test_short_audio_override_does_not_mutate_separator(): + separator = object.__new__(MDXCSeparator) + separator.override_model_segment_size = False + separator.logger = Mock() + + assert separator._use_model_segment_override(5.0) is True + assert separator.override_model_segment_size is False + assert separator.logger.warning.call_count == 2 + assert separator._use_model_segment_override(20.0) is False + assert separator._use_model_segment_override(10.0) is False + + +def test_explicit_segment_override_applies_to_all_inputs(): + separator = object.__new__(MDXCSeparator) + separator.override_model_segment_size = True + separator.logger = Mock() + + assert separator._use_model_segment_override(5.0) is True + assert separator._use_model_segment_override(20.0) is True + separator.logger.warning.assert_not_called() diff --git a/tests/unit/test_model_reuse.py b/tests/unit/test_model_reuse.py new file mode 100644 index 0000000..89b048b --- /dev/null +++ b/tests/unit/test_model_reuse.py @@ -0,0 +1,365 @@ +import logging +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from audio_separator.separator import Separator +from audio_separator.separator.architectures.vr_separator import VRSeparator + + +@pytest.fixture +def separator(tmp_path): + return Separator(model_file_dir=tmp_path / "models", output_dir=tmp_path / "output", info_only=True) + + +def test_load_model_reuses_matching_instance(separator): + loaded_instance = object() + separator.model_instance = loaded_instance + separator._loaded_model_filename = "model.ckpt" + + with patch.object(separator, "download_model_files") as download_model_files: + separator.load_model("model.ckpt") + + download_model_files.assert_not_called() + assert separator.model_instance is loaded_instance + assert separator.model_filename == "model.ckpt" + assert separator.model_filenames == ["model.ckpt"] + + +def test_load_model_normalizes_single_item_list_before_reuse(separator): + separator.model_instance = object() + separator._loaded_model_filename = "model.ckpt" + + with patch.object(separator, "download_model_files") as download_model_files: + separator.load_model(["model.ckpt"]) + + download_model_files.assert_not_called() + assert separator.model_filename == "model.ckpt" + assert separator.model_filenames == ["model.ckpt"] + + +def test_load_model_reuse_restores_loaded_model_metadata(separator): + loaded_instance = object() + separator.model_instance = loaded_instance + separator._loaded_model_filename = "first.ckpt" + separator._loaded_model_friendly_name = "First model" + separator._loaded_model_is_uvr_vip = False + + def download_second_model(_model_filename): + separator.model_friendly_name = "Second VIP model" + separator.model_is_uvr_vip = True + return "second.ckpt", "MDXC", "Second VIP model", "/tmp/second.ckpt", None + + with ( + patch.object(separator, "download_model_files", side_effect=download_second_model) as download_model_files, + patch.object(separator, "load_model_data_using_hash", return_value={}), + ): + separator.download_model_and_data("second.ckpt") + separator.load_model("first.ckpt") + + download_model_files.assert_called_once_with("second.ckpt") + assert separator.model_instance is loaded_instance + assert separator.model_friendly_name == "First model" + assert separator.model_is_uvr_vip is False + + +def test_load_model_force_reload_bypasses_reuse(separator): + separator.model_instance = object() + separator._loaded_model_filename = "model.ckpt" + + with patch.object(separator, "download_model_files", side_effect=RuntimeError("reload attempted")) as download_model_files: + with pytest.raises(RuntimeError, match="reload attempted"): + separator.load_model("model.ckpt", force_reload=True) + + download_model_files.assert_called_once_with("model.ckpt") + + +def test_load_model_reloads_when_instance_is_missing(separator): + separator._loaded_model_filename = "model.ckpt" + + with patch.object(separator, "download_model_files", side_effect=RuntimeError("reload attempted")): + with pytest.raises(RuntimeError, match="reload attempted"): + separator.load_model("model.ckpt") + + +def test_load_model_reloads_different_model_without_poisoning_cache(separator): + loaded_instance = object() + separator.model_instance = loaded_instance + separator.model_filename = "first.ckpt" + separator.model_filenames = ["first.ckpt"] + separator._loaded_model_filename = "first.ckpt" + separator.model_friendly_name = "First model" + separator.model_is_uvr_vip = False + + def fail_download(_model_filename): + separator.model_friendly_name = "Second model" + separator.model_is_uvr_vip = True + raise RuntimeError("load failed") + + with patch.object(separator, "download_model_files", side_effect=fail_download): + with pytest.raises(RuntimeError, match="load failed"): + separator.load_model("second.ckpt") + + assert separator.model_instance is loaded_instance + assert separator.model_filename == "first.ckpt" + assert separator.model_filenames == ["first.ckpt"] + assert separator._loaded_model_filename == "first.ckpt" + assert separator.model_friendly_name == "First model" + assert separator.model_is_uvr_vip is False + + +@pytest.mark.parametrize( + "failure", + [RuntimeError("construction failed"), SystemExit("construction stopped")], + ids=["runtime-error", "system-exit"], +) +def test_load_model_constructor_failure_preserves_previous_selection(separator, failure): + loaded_instance = object() + separator.model_instance = loaded_instance + separator.model_filename = "first.ckpt" + separator.model_filenames = ["first.ckpt"] + separator._loaded_model_filename = "first.ckpt" + separator.model_friendly_name = "First model" + separator.model_is_uvr_vip = False + separator_class = MagicMock(side_effect=failure) + architecture_module = SimpleNamespace(MDXCSeparator=separator_class) + + def resolve_download(_model_filename): + separator.model_friendly_name = "Second model" + separator.model_is_uvr_vip = True + return "second.ckpt", "MDXC", "Second model", "/tmp/second.ckpt", None + + with ( + patch.object(separator, "download_model_files", side_effect=resolve_download), + patch.object(separator, "load_model_data_using_hash", return_value={}), + patch("audio_separator.separator.separator.importlib.import_module", return_value=architecture_module), + ): + with pytest.raises(type(failure), match=str(failure)): + separator.load_model("second.ckpt") + + assert separator.model_instance is loaded_instance + assert separator.model_filename == "first.ckpt" + assert separator.model_filenames == ["first.ckpt"] + assert separator._loaded_model_filename == "first.ckpt" + assert separator.model_friendly_name == "First model" + assert separator.model_is_uvr_vip is False + + +def test_load_model_policy_failure_preserves_previous_ensemble_selection(separator): + ensemble_models = ["first.ckpt", "third.ckpt"] + loaded_instance = object() + separator.model_instance = loaded_instance + separator.model_filename = list(ensemble_models) + separator.model_filenames = list(ensemble_models) + separator._loaded_model_filename = "first.ckpt" + separator.model_friendly_name = "First model" + separator.model_is_uvr_vip = False + resolve_execution_policy = MagicMock(side_effect=RuntimeError("policy failed")) + candidate_instance = SimpleNamespace( + _execution_policy_resolved=False, + is_roformer_model=False, + resolve_execution_policy=resolve_execution_policy, + ) + separator_class = MagicMock(return_value=candidate_instance) + architecture_module = SimpleNamespace(MDXCSeparator=separator_class) + + with ( + patch.object( + separator, + "download_model_files", + return_value=("second.ckpt", "MDXC", "Second model", "/tmp/second.ckpt", None), + ), + patch.object(separator, "load_model_data_using_hash", return_value={}), + patch("audio_separator.separator.separator.importlib.import_module", return_value=architecture_module), + ): + with pytest.raises(RuntimeError, match="policy failed"): + separator.load_model("second.ckpt") + + resolve_execution_policy.assert_called_once_with("mdxc") + assert separator.model_instance is loaded_instance + assert separator.model_filename == ensemble_models + assert separator.model_filenames == ensemble_models + assert separator.model_filename is not ensemble_models + assert separator._loaded_model_filename == "first.ckpt" + assert separator.model_friendly_name == "First model" + assert separator.model_is_uvr_vip is False + + +def test_successful_load_populates_cache_for_the_next_call(separator): + loaded_instance = object() + separator_class = MagicMock(return_value=loaded_instance) + architecture_module = SimpleNamespace(MDXCSeparator=separator_class) + + with ( + patch.object( + separator, + "download_model_files", + return_value=("model.ckpt", "MDXC", "Model", "/tmp/model.ckpt", None), + ) as download_model_files, + patch.object(separator, "load_model_data_using_hash", return_value={}), + patch("audio_separator.separator.separator.importlib.import_module", return_value=architecture_module), + ): + separator.load_model("model.ckpt") + separator.load_model("model.ckpt") + + download_model_files.assert_called_once_with("model.ckpt") + separator_class.assert_called_once() + assert separator.model_instance is loaded_instance + assert separator._loaded_model_filename == "model.ckpt" + assert separator._loaded_model_friendly_name == "Model" + assert separator._loaded_model_is_uvr_vip is False + + +def test_force_reload_repeats_a_successful_load(separator): + separator_class = MagicMock(side_effect=[object(), object()]) + architecture_module = SimpleNamespace(MDXCSeparator=separator_class) + + with ( + patch.object( + separator, + "download_model_files", + return_value=("model.ckpt", "MDXC", "Model", "/tmp/model.ckpt", None), + ) as download_model_files, + patch.object(separator, "load_model_data_using_hash", return_value={}), + patch("audio_separator.separator.separator.importlib.import_module", return_value=architecture_module), + ): + separator.load_model("model.ckpt") + first_instance = separator.model_instance + separator.load_model("model.ckpt", force_reload=True) + + assert download_model_files.call_count == 2 + assert separator_class.call_count == 2 + assert separator.model_instance is not first_instance + assert separator._loaded_model_filename == "model.ckpt" + + +def test_load_model_preserves_multi_model_ensemble_behavior(separator): + models = ["first.ckpt", "second.ckpt"] + + with patch.object(separator, "download_model_files") as download_model_files: + separator.load_model(models) + + download_model_files.assert_not_called() + assert separator.model_filename == models + assert separator.model_filenames == models + assert separator.model_filename is not models + + +def test_load_model_expands_ensemble_preset_before_reuse(separator): + preset_models = ["first.ckpt", "second.ckpt"] + separator._ensemble_preset_models = preset_models + + with patch.object(separator, "download_model_files") as download_model_files: + separator.load_model() + + download_model_files.assert_not_called() + assert separator.model_filename == preset_models + assert separator.model_filenames == preset_models + + +def _make_vr_separator(model_run): + separator = object.__new__(VRSeparator) + separator.logger = logging.getLogger(__name__) + separator.model_run = model_run + separator.model_params = SimpleNamespace(param={"bins": 128}) + separator.model_capacity = (32, 128) + separator.is_vr_51_model = False + separator.model_path = "/tmp/model.pth" + separator.torch_device = torch.device("cpu") + return separator + + +def test_vr_model_reuses_loaded_torch_module(): + loaded_model = torch.nn.Linear(2, 2) + separator = _make_vr_separator(loaded_model) + + with ( + patch("audio_separator.separator.architectures.vr_separator.nets.determine_model_capacity") as determine_model_capacity, + patch("audio_separator.separator.architectures.vr_separator.torch.load") as load_weights, + ): + separator._ensure_model_loaded(31191) + + determine_model_capacity.assert_not_called() + load_weights.assert_not_called() + assert separator.model_run is loaded_model + + +def test_vr_model_loads_placeholder_once(): + separator = _make_vr_separator(lambda: None) + loaded_model = MagicMock(spec=torch.nn.Module) + state_dict = {"weight": torch.tensor([1.0])} + + with ( + patch( + "audio_separator.separator.architectures.vr_separator.nets.determine_model_capacity", + return_value=loaded_model, + ) as determine_model_capacity, + patch( + "audio_separator.separator.architectures.vr_separator.torch.load", + return_value=state_dict, + ) as load_weights, + ): + separator._ensure_model_loaded(31191) + + determine_model_capacity.assert_called_once_with(256, 31191) + load_weights.assert_called_once_with("/tmp/model.pth", map_location="cpu") + loaded_model.load_state_dict.assert_called_once_with(state_dict) + loaded_model.to.assert_called_once_with(torch.device("cpu")) + assert separator.model_run is loaded_model + + +def test_vr_model_retries_after_weight_loading_failure(): + placeholder = lambda: None + separator = _make_vr_separator(placeholder) + failed_model = MagicMock(spec=torch.nn.Module) + failed_model.load_state_dict.side_effect = RuntimeError("invalid weights") + loaded_model = MagicMock(spec=torch.nn.Module) + + with ( + patch( + "audio_separator.separator.architectures.vr_separator.nets.determine_model_capacity", + side_effect=[failed_model, loaded_model], + ) as determine_model_capacity, + patch( + "audio_separator.separator.architectures.vr_separator.torch.load", + side_effect=[{"bad": torch.tensor([1.0])}, {"weight": torch.tensor([2.0])}], + ) as load_weights, + ): + with pytest.raises(RuntimeError, match="invalid weights"): + separator._ensure_model_loaded(31191) + + assert separator.model_run is placeholder + separator._ensure_model_loaded(31191) + + assert determine_model_capacity.call_count == 2 + assert load_weights.call_count == 2 + loaded_model.load_state_dict.assert_called_once() + loaded_model.to.assert_called_once_with(torch.device("cpu")) + assert separator.model_run is loaded_model + + +def test_effective_mode_is_reset_while_an_ensemble_is_selected(separator): + separator.model_instance = SimpleNamespace(effective_precision="native_fp16", effective_torch_compile=True) + separator._loaded_model_filename = "first.ckpt" + + separator.load_model(["first.ckpt", "second.ckpt"]) + + assert separator.effective_precision == "fp32" + assert separator.effective_torch_compile is False + + +def test_failed_separation_clears_reused_instance_state(separator): + model_instance = MagicMock() + model_instance.effective_precision = "fp32" + model_instance.torch_device = torch.device("cpu") + model_instance.separate.side_effect = RuntimeError("inference failed") + model_instance.clear_gpu_cache.side_effect = RuntimeError("cleanup failed") + separator.model_instance = model_instance + + with pytest.raises(RuntimeError, match="inference failed"): + separator._separate_file("input.wav") + + model_instance.clear_gpu_cache.assert_called_once_with() + model_instance.clear_file_specific_paths.assert_called_once_with() diff --git a/tests/unit/test_mps_device_accumulation.py b/tests/unit/test_mps_device_accumulation.py new file mode 100644 index 0000000..4973680 --- /dev/null +++ b/tests/unit/test_mps_device_accumulation.py @@ -0,0 +1,226 @@ +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import numpy as np +import pytest +import torch + +from audio_separator.separator.architectures.demucs_separator import ( + DemucsSeparator, + _estimate_demucs_full_track_buffer_bytes, +) +from audio_separator.separator.architectures.mdxc_separator import ( + MDXCSeparator, + _estimate_mdxc_full_track_buffer_bytes, + _estimate_roformer_full_track_buffer_bytes, +) +from audio_separator.separator.uvr_lib_v5 import device_utils +from audio_separator.separator.uvr_lib_v5.device_utils import ( + MAX_MPS_FULL_TRACK_BUFFER_BYTES, + should_accumulate_on_device, +) + + +@pytest.fixture +def fixed_floor_budget(): + """Pin the budget to the fixed floor so crossover maths is host-independent.""" + with patch.object(device_utils, "_mps_memory_reading", return_value=0): + yield MAX_MPS_FULL_TRACK_BUFFER_BYTES + + +def _demucs_separator(device: torch.device) -> DemucsSeparator: + separator = object.__new__(DemucsSeparator) + separator.logger = Mock() + separator.torch_device = device + separator.demucs_model_instance = Mock() + separator.demucs_model_instance.sources = ["drums", "bass", "other", "vocals"] + separator.demucs_model_instance.models = [] + separator.shifts = 0 + separator.segments_enabled = True + separator.overlap = 0.25 + return separator + + +def _mix() -> np.ndarray: + rng = np.random.default_rng(0) + return rng.standard_normal((2, 128), dtype=np.float32) + + +@pytest.mark.parametrize( + "estimate_bytes", + [ + pytest.param( + lambda samples: _estimate_roformer_full_track_buffer_bytes(2, 2, samples, chunk_size=485100), + id="roformer", + ), + pytest.param( + lambda samples: _estimate_mdxc_full_track_buffer_bytes(2, 2, padded_length=samples + 485100), + id="mdxc", + ), + pytest.param( + lambda samples: _estimate_demucs_full_track_buffer_bytes( + 2, + samples, + 4, + shifts=2, + num_bag_models=1, + ), + id="demucs", + ), + ], +) +def test_full_track_buffer_estimates_cross_mps_limit_at_adjacent_sample_counts(estimate_bytes, fixed_floor_budget): + below_samples = 0 + above_samples = 1 + while estimate_bytes(above_samples) <= MAX_MPS_FULL_TRACK_BUFFER_BYTES: + below_samples = above_samples + above_samples *= 2 + + while above_samples - below_samples > 1: + midpoint = (below_samples + above_samples) // 2 + if estimate_bytes(midpoint) <= MAX_MPS_FULL_TRACK_BUFFER_BYTES: + below_samples = midpoint + else: + above_samples = midpoint + + below_limit = estimate_bytes(below_samples) + above_limit = estimate_bytes(above_samples) + + assert above_samples == below_samples + 1 + assert below_limit <= MAX_MPS_FULL_TRACK_BUFFER_BYTES < above_limit + assert should_accumulate_on_device(torch.device("mps"), below_limit) is True + assert should_accumulate_on_device(torch.device("mps"), above_limit) is False + + +class _FakeMDXCModel: + num_target_instruments = 2 + + def __init__(self, expected_device: torch.device): + self.expected_device = expected_device + + def __call__(self, batch): + assert batch.device.type == self.expected_device.type + return batch.unsqueeze(1).repeat(1, 2, 1, 1) + + +class _FakeRoformerModel(torch.nn.Module): + def __init__(self, device: torch.device): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros((), device=device)) + + def forward(self, batch): + assert batch.device == self.anchor.device + return batch.unsqueeze(1).repeat(1, 2, 1, 1) + + +def _mdxc_separator(device: torch.device) -> MDXCSeparator: + separator = object.__new__(MDXCSeparator) + separator.logger = Mock() + separator.torch_device = device + separator.model_run = _FakeMDXCModel(device) + separator.model_data_cfgdict = SimpleNamespace( + training=SimpleNamespace(instruments=["first", "second"], target_instrument=None), + inference=SimpleNamespace(dim_t=5), + audio=SimpleNamespace(hop_length=2), + ) + separator.pitch_shift = 0 + separator.is_roformer = False + separator.segment_size = 5 + separator.overlap = 2 + separator.batch_size = 1 + separator.is_primary_stem_main_target = False + return separator + + +def _roformer_separator(device: torch.device) -> MDXCSeparator: + separator = _mdxc_separator(device) + separator.model_run = _FakeRoformerModel(device) + separator.model_data_cfgdict.model = SimpleNamespace(stft_hop_length=2) + separator.model_data_cfgdict.audio.sample_rate = 1 + separator.is_roformer = True + separator.overlap = 8 + return separator + + +def test_demucs_keeps_full_track_input_on_cpu_for_cpu_inference(): + separator = _demucs_separator(torch.device("cpu")) + + def fake_apply_model(*, mix, **kwargs): + assert mix.device.type == "cpu" + return torch.zeros(1, 4, 2, mix.shape[-1], device=mix.device) + + with patch("audio_separator.separator.architectures.demucs_separator.apply_model", side_effect=fake_apply_model): + result = separator.demix_demucs(_mix()) + + assert result.shape == (4, 2, 128) + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available") +def test_demucs_keeps_full_track_input_on_mps(): + separator = _demucs_separator(torch.device("mps")) + + def fake_apply_model(*, mix, **kwargs): + assert mix.device.type == "mps" + return torch.zeros(1, 4, 2, mix.shape[-1], device=mix.device) + + with patch("audio_separator.separator.architectures.demucs_separator.apply_model", side_effect=fake_apply_model): + result = separator.demix_demucs(_mix()) + + assert result.shape == (4, 2, 128) + + +@pytest.mark.parametrize("device_type", ["cpu", "mps"]) +def test_mdxc_chunk_buffers_share_the_selected_accumulation_device(device_type): + if device_type == "mps" and not torch.backends.mps.is_available(): + pytest.skip("MPS is not available") + + device = torch.device(device_type) + separator = _mdxc_separator(device) + allocated_devices = [] + torch_zeros = torch.zeros + + def tracked_zeros(*args, **kwargs): + tensor = torch_zeros(*args, **kwargs) + allocated_devices.append(tensor.device.type) + return tensor + + with patch("audio_separator.separator.architectures.mdxc_separator.torch.zeros", side_effect=tracked_zeros): + result = separator.demix(_mix(), override_model_segment_size=True) + + assert set(result) == {"first", "second"} + assert all(stem.shape == (2, 128) for stem in result.values()) + assert allocated_devices + assert set(allocated_devices) == {device_type} + + +@pytest.mark.parametrize("device_type", ["cpu", "mps"]) +def test_roformer_overlap_add_buffers_use_the_selected_accumulation_device(device_type): + if device_type == "mps" and not torch.backends.mps.is_available(): + pytest.skip("MPS is not available") + + separator = _roformer_separator(torch.device(device_type)) + window_devices = [] + zero_devices = [] + torch_tensor = torch.tensor + torch_zeros = torch.zeros + + def tracked_tensor(*args, **kwargs): + tensor = torch_tensor(*args, **kwargs) + if kwargs.get("device") is not None: + window_devices.append(tensor.device.type) + return tensor + + def tracked_zeros(*args, **kwargs): + tensor = torch_zeros(*args, **kwargs) + zero_devices.append(tensor.device.type) + return tensor + + with ( + patch("audio_separator.separator.architectures.mdxc_separator.torch.tensor", side_effect=tracked_tensor), + patch("audio_separator.separator.architectures.mdxc_separator.torch.zeros", side_effect=tracked_zeros), + ): + result = separator.demix(_mix(), override_model_segment_size=True) + + assert set(result) == {"first", "second"} + assert window_devices == [device_type] + assert zero_devices == [device_type, device_type] diff --git a/tests/unit/test_mps_native_fp16.py b/tests/unit/test_mps_native_fp16.py new file mode 100644 index 0000000..5d8f919 --- /dev/null +++ b/tests/unit/test_mps_native_fp16.py @@ -0,0 +1,304 @@ +import copy +import platform +import subprocess +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch +from rotary_embedding_torch import RotaryEmbedding + +from audio_separator.separator.architectures.mdxc_separator import MDXCSeparator +from audio_separator.separator.execution_policy import AUTOCAST, FP32, NATIVE_FP16 +from audio_separator.separator.separator import Separator +from audio_separator.separator.uvr_lib_v5.roformer import mel_band_roformer as mel_module + + +def _apple_gpu_is_virtualized() -> bool: + """Detect a paravirtualized Metal device (hosted CI Macs report VirtualMac*).""" + if platform.system() != "Darwin": + return False + try: + result = subprocess.run(["/usr/sbin/sysctl", "-n", "hw.model"], capture_output=True, text=True, timeout=5, check=False) + except (OSError, subprocess.TimeoutExpired): + return False + if result.returncode != 0: + return False + return result.stdout.strip().startswith("VirtualMac") + + +class MelBandRoformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.projection = torch.nn.Linear(2, 2) + self.rotary_embedding = RotaryEmbedding(dim=8) + + +def _separator(device="mps", use_native_fp16=True, model_type="mel_band_roformer"): + separator = object.__new__(MDXCSeparator) + separator.logger = Mock() + separator.model_run = MelBandRoformer() + separator.roformer_model_type = model_type + separator.torch_device = torch.device(device) + separator.requested_torch_device = separator.torch_device + separator.use_autocast = False + separator.use_native_fp16 = use_native_fp16 + separator.use_torch_compile = False + separator.is_native_fp16 = False + return separator + + +def _dispatch_separator(*, effective_precision, device="mps", separator_device=None): + separator = object.__new__(Separator) + separator.chunk_duration = None + separator.logger = Mock() + separator.normalization_threshold = 1.0 + separator.amplification_threshold = 0.0 + separator.use_autocast = effective_precision == AUTOCAST + separator.use_native_fp16 = effective_precision == NATIVE_FP16 + separator.torch_device = torch.device(separator_device or device) + separator.model_instance = Mock(effective_precision=effective_precision, torch_device=torch.device(device)) + separator.model_instance.separate.return_value = ["output.wav"] + separator.print_uvr_vip_message = Mock() + return separator + + +def test_mel_band_roformer_uses_explicit_native_fp16_on_mps(): + separator = _separator() + separator.model_run.rotary_embedding.rotate_queries_or_keys(torch.randn(1, 4, 8)) + expected_frequencies = separator.model_run.rotary_embedding.freqs.detach().clone() + assert separator.model_run.rotary_embedding.cached_freqs is not None + + separator._configure_model_precision() + + assert separator.is_native_fp16 is True + assert separator.model_run.projection.weight.dtype == torch.float16 + assert separator.model_run.rotary_embedding.freqs.dtype == torch.float32 + assert separator.model_run.rotary_embedding.cached_freqs is None + torch.testing.assert_close(separator.model_run.rotary_embedding.freqs, expected_frequencies, rtol=0, atol=0) + + output = separator.model_run.rotary_embedding.rotate_queries_or_keys(torch.randn(1, 1101, 8).half()) + + assert output.dtype == torch.float16 + assert torch.isfinite(output).all() + assert separator.model_run.rotary_embedding.cached_freqs.dtype == torch.float32 + + +@pytest.mark.parametrize( + ("device", "use_native_fp16", "model_type"), + [ + ("cpu", True, "mel_band_roformer"), + ("mps", False, "mel_band_roformer"), + ], +) +def test_native_fp16_requires_verified_device_and_model(device, use_native_fp16, model_type): + separator = _separator(device=device, use_native_fp16=use_native_fp16, model_type=model_type) + + separator._configure_model_precision() + + assert separator.is_native_fp16 is False + assert separator.model_run.projection.weight.dtype == torch.float32 + + +def test_native_fp16_supports_legacy_loader_class_detection(): + separator = _separator(model_type=None) + + separator._configure_model_precision() + + assert separator.is_native_fp16 is True + + +@pytest.mark.parametrize( + ("device", "use_native_fp16", "expected_load_device"), + [ + ("mps", True, "cpu"), + ("mps", False, "mps"), + ("cpu", True, "cpu"), + ("cuda", True, "cuda"), + ("privateuseone", True, "privateuseone"), + ], +) +def test_roformer_loads_on_cpu_only_before_native_mps_conversion(device, use_native_fp16, expected_load_device): + separator = object.__new__(MDXCSeparator) + separator.logger = Mock() + separator.is_roformer = True + separator.model_data = {} + separator.model_path = "/tmp/model.ckpt" + separator.torch_device = torch.device(device) + separator.requested_torch_device = separator.torch_device + separator.use_autocast = False + separator.use_native_fp16 = use_native_fp16 + separator.use_torch_compile = False + loaded_model = Mock() + loaded_model.to.return_value = loaded_model + separator.roformer_loader = Mock( + load_model=Mock( + return_value=SimpleNamespace( + success=True, + model=loaded_model, + model_info={"model_type": "mel_band_roformer"}, + ) + ) + ) + + with ( + patch.object(separator, "_configure_model_precision") as configure_precision, + patch.object(separator, "_configure_model_compilation") as configure_compilation, + ): + separator.load_model() + + separator.roformer_loader.load_model.assert_called_once_with( + model_path="/tmp/model.ckpt", + config={}, + device=expected_load_device, + ) + configure_precision.assert_called_once_with() + configure_compilation.assert_called_once_with() + loaded_model.to.assert_called_once_with(torch.device(device)) + loaded_model.eval.assert_called_once_with() + + +def test_separator_bypasses_autocast_for_native_fp16(): + separator = _dispatch_separator(effective_precision=NATIVE_FP16) + + with ( + patch("audio_separator.separator.separator.autocast_mode.autocast", return_value=nullcontext()) as autocast, + ): + output_files = separator._separate_file("input.wav") + + assert output_files == ["output.wav"] + autocast.assert_not_called() + separator.logger.debug.assert_any_call("Using native float16 inference.") + + +def test_separator_preserves_non_native_autocast(): + separator = _dispatch_separator(effective_precision=AUTOCAST, device="cpu", separator_device="mps") + + with ( + patch("audio_separator.separator.separator.autocast_mode.autocast", return_value=nullcontext()) as autocast, + ): + separator._separate_file("input.wav") + + autocast.assert_called_once_with("cpu") + + +def test_separator_uses_float32_for_fallback_policy(): + separator = _dispatch_separator(effective_precision=FP32, device="cpu") + + with patch("audio_separator.separator.separator.autocast_mode.autocast", return_value=nullcontext()) as autocast: + separator._separate_file("input.wav") + + autocast.assert_not_called() + separator.logger.debug.assert_any_call("Using float32 inference.") + + +def _tiny_mel_band_roformer(): + torch.manual_seed(0) + return mel_module.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() + + +def _half_preserving_rotary_frequencies(model): + rotary_frequencies = [module.freqs.detach().clone() for module in model.modules() if isinstance(module, RotaryEmbedding)] + + model.half() + + rotary_modules = [module for module in model.modules() if isinstance(module, RotaryEmbedding)] + for rotary, frequencies in zip(rotary_modules, rotary_frequencies): + rotary.freqs.data = frequencies.to(rotary.freqs.device) + rotary.cached_freqs = None + + return model + + +def test_mel_band_half_forward_keeps_cpu_silence_finite_and_zero(): + model = _half_preserving_rotary_frequencies(_tiny_mel_band_roformer()) + + with torch.no_grad(): + output = model(torch.zeros(1, 8192)) + + assert output.shape == (1, 1, 8192) + assert output.dtype == torch.float32 + assert torch.isfinite(output).all() + assert torch.count_nonzero(output) == 0 + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available") +@pytest.mark.parametrize("force_cpu_complex", [False, True]) +def test_native_fp16_mps_forward_handles_silence(force_cpu_complex): + separator = object.__new__(MDXCSeparator) + separator.logger = Mock() + separator.model_run = _tiny_mel_band_roformer() + separator.roformer_model_type = "mel_band_roformer" + separator.torch_device = torch.device("mps") + separator.requested_torch_device = separator.torch_device + separator.use_autocast = False + separator.use_native_fp16 = True + separator.use_torch_compile = False + separator.is_native_fp16 = False + separator._configure_model_precision() + separator.model_run.to(separator.torch_device) + audio = torch.zeros(1, 8192, device=separator.torch_device) + + with patch.object(mel_module, "should_fallback_to_cpu_for_complex_ops", return_value=force_cpu_complex), torch.no_grad(): + output = separator.model_run(audio) + + assert output.shape == audio.unsqueeze(1).shape + assert torch.isfinite(output).all() + assert torch.count_nonzero(output) == 0 + assert output.device.type == ("cpu" if force_cpu_complex else "mps") + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available") +@pytest.mark.skipif( + _apple_gpu_is_virtualized(), + reason="fp16 SNR gates need a real Apple GPU: virtualized Metal degrades half-precision accumulation", +) +@pytest.mark.parametrize("force_cpu_complex", [False, True]) +def test_native_fp16_mps_forward_matches_cpu_fp32_for_non_silent_audio(force_cpu_complex): + cpu_model = _tiny_mel_band_roformer() + mps_model = copy.deepcopy(cpu_model) + separator = object.__new__(MDXCSeparator) + separator.logger = Mock() + separator.model_run = mps_model + separator.roformer_model_type = "mel_band_roformer" + separator.torch_device = torch.device("mps") + separator.requested_torch_device = separator.torch_device + separator.use_autocast = False + separator.use_native_fp16 = True + separator.use_torch_compile = False + separator.is_native_fp16 = False + separator._configure_model_precision() + separator.model_run.to(separator.torch_device) + + sample_indices = torch.arange(8192, dtype=torch.float32) + audio = ( + 0.35 * torch.sin(2 * torch.pi * 440 * sample_indices / 44100) + 0.15 * torch.sin(2 * torch.pi * 880 * sample_indices / 44100) + ).unsqueeze(0) + + with torch.no_grad(): + reference = cpu_model(audio) + with patch.object(mel_module, "should_fallback_to_cpu_for_complex_ops", return_value=force_cpu_complex): + output = separator.model_run(audio.to(separator.torch_device)) + + assert output.device.type == ("cpu" if force_cpu_complex else "mps") + output = output.cpu().float() + error = output - reference + snr = 20 * torch.log10(reference.square().mean().sqrt() / error.square().mean().sqrt()) + + assert output.shape == reference.shape + assert torch.isfinite(output).all() + assert reference.square().mean().sqrt().item() > 1e-4 + assert snr.item() > 30 + torch.testing.assert_close(output, reference, rtol=0.1, atol=1e-3) diff --git a/tests/unit/test_mps_stft_helpers.py b/tests/unit/test_mps_stft_helpers.py new file mode 100644 index 0000000..cf5fc48 --- /dev/null +++ b/tests/unit/test_mps_stft_helpers.py @@ -0,0 +1,104 @@ +from unittest.mock import Mock, patch + +import pytest +import torch + +from audio_separator.separator.uvr_lib_v5 import device_utils +from audio_separator.separator.uvr_lib_v5.stft import STFT as CommonSTFT +from audio_separator.separator.uvr_lib_v5.tfc_tdf_v3 import STFT as TFCSTFT + + +def _common_stft(device): + return CommonSTFT(Mock(), n_fft=256, hop_length=64, dim_f=129, device=device) + + +def _tfc_stft(device): + return TFCSTFT(n_fft=256, hop_length=64, dim_f=129, device=device) + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available") +@pytest.mark.filterwarnings("ignore:stft with return_complex=False is deprecated:UserWarning") +@pytest.mark.parametrize( + ("factory", "module_name"), + [ + pytest.param( + _common_stft, + "audio_separator.separator.uvr_lib_v5.stft", + id="common-stft", + ), + pytest.param( + _tfc_stft, + "audio_separator.separator.uvr_lib_v5.tfc_tdf_v3", + id="tfc-tdf-stft", + ), + ], +) +@pytest.mark.parametrize("force_cpu_fallback", [False, True]) +def test_mps_stft_round_trip_uses_selected_compute_device( + factory, + module_name, + force_cpu_fallback, + monkeypatch, +): + monkeypatch.delenv("AUDIO_SEPARATOR_FORCE_CPU_COMPLEX", raising=False) + mps_device = torch.device("mps") + + if not force_cpu_fallback and device_utils.should_fallback_to_cpu_for_complex_ops(mps_device): + pytest.skip("This MPS runtime does not support the native spectral path") + + indices = torch.arange(2048, dtype=torch.float32) + audio = torch.stack( + ( + 0.5 * torch.sin(2 * torch.pi * 440 * indices / 44100), + 0.3 * torch.cos(2 * torch.pi * 880 * indices / 44100), + ) + ).unsqueeze(0) + + cpu_stft = factory(torch.device("cpu")) + reference_spectrum = cpu_stft(audio) + reference_audio = cpu_stft.inverse(reference_spectrum) + + real_stft = torch.stft + real_istft = torch.istft + stft_devices = [] + istft_devices = [] + + def tracked_stft(input_tensor, *args, **kwargs): + stft_devices.append(input_tensor.device.type) + return real_stft(input_tensor, *args, **kwargs) + + def tracked_istft(input_tensor, *args, **kwargs): + istft_devices.append(input_tensor.device.type) + return real_istft(input_tensor, *args, **kwargs) + + with ( + patch( + f"{module_name}.should_fallback_to_cpu_for_complex_ops", + return_value=force_cpu_fallback, + ) as fallback, + patch(f"{module_name}.torch.stft", side_effect=tracked_stft), + patch(f"{module_name}.torch.istft", side_effect=tracked_istft), + ): + mps_stft = factory(mps_device) + spectrum = mps_stft(audio.to(mps_device)) + reconstructed = mps_stft.inverse(spectrum) + + compute_device = "cpu" if force_cpu_fallback else "mps" + assert stft_devices == [compute_device] + assert istft_devices == [compute_device] + assert spectrum.device.type == "mps" + assert reconstructed.device.type == "mps" + assert fallback.call_count == 2 + + torch.testing.assert_close( + spectrum.cpu(), + reference_spectrum, + rtol=2e-4, + atol=2e-5, + ) + torch.testing.assert_close( + reconstructed.cpu(), + reference_audio, + rtol=2e-4, + atol=2e-5, + ) diff --git a/tests/unit/test_mps_torch_compile.py b/tests/unit/test_mps_torch_compile.py new file mode 100644 index 0000000..4ac260c --- /dev/null +++ b/tests/unit/test_mps_torch_compile.py @@ -0,0 +1,206 @@ +from unittest.mock import Mock, patch + +from packaging import version +import pytest +import torch + +from audio_separator.separator.architectures.mdxc_separator import MDXCSeparator +from audio_separator.separator.uvr_lib_v5.roformer.bs_roformer import Transformer as BSTransformer +from audio_separator.separator.uvr_lib_v5.roformer.mel_band_roformer import Transformer as MelBandTransformer +from rotary_embedding_torch import RotaryEmbedding + + +def _separator(*, use_torch_compile=True, policy_allows_compile=True): + separator = object.__new__(MDXCSeparator) + separator.logger = Mock() + separator.use_torch_compile = use_torch_compile + separator._should_torch_compile = use_torch_compile and policy_allows_compile + separator.effective_torch_compile = False + separator.torch_device = torch.device("mps") + separator.model_run = Mock() + separator.model_run.layers = torch.nn.ModuleList([torch.nn.ModuleList([torch.nn.Identity(), torch.nn.Identity()])]) + return separator + + +@pytest.mark.parametrize( + "transformer", + [ + MelBandTransformer(dim=16, depth=1, dim_head=4, heads=2, rotary_embed=RotaryEmbedding(dim=4), flash_attn=True), + BSTransformer(dim=16, depth=1, dim_head=4, heads=2, rotary_embed=RotaryEmbedding(dim=4), flash_attn=True), + BSTransformer(dim=16, depth=1, dim_head=4, heads=2, flash_attn=True, linear_attn=True), + ], + ids=["mel-band", "bs-rotary", "bs-linear"], +) +def test_regional_transformer_is_captured_as_one_dynamo_graph(transformer): + if version.parse(torch.__version__.split("+")[0]) < version.parse("2.6"): + pytest.skip("Regional compilation requires PyTorch 2.6 or newer") + + explanation = torch._dynamo.explain(transformer.eval())(torch.randn(2, 8, 16)) + + assert explanation.graph_count == 1 + assert explanation.graph_break_count == 0 + + +def test_regional_compile_wraps_repeated_transformers(): + separator = _separator() + + with patch.object(torch.nn.Module, "compile", autospec=True) as compile_module: + separator._configure_model_compilation() + + assert separator.is_torch_compiled is True + assert compile_module.call_count == 2 + separator.logger.warning.assert_not_called() + + +def test_regional_compile_skips_when_policy_does_not_enable_it(): + separator = _separator(policy_allows_compile=False) + + with patch.object(torch.nn.Module, "compile", autospec=True) as compile_module: + separator._configure_model_compilation() + + assert separator.is_torch_compiled is False + assert separator.effective_torch_compile is False + compile_module.assert_not_called() + separator.logger.warning.assert_not_called() + + +def test_disabled_regional_compile_is_silent(): + separator = _separator(use_torch_compile=False) + + with patch.object(torch.nn.Module, "compile", autospec=True) as compile_module: + separator._configure_model_compilation() + + assert separator.is_torch_compiled is False + compile_module.assert_not_called() + separator.logger.warning.assert_not_called() + + +def test_regional_compile_requires_restorable_module_calls(): + separator = _separator() + transformer = Mock(spec=["compile"]) + separator.model_run.layers = [[transformer]] + + separator._configure_model_compilation() + + assert separator.is_torch_compiled is False + transformer.compile.assert_not_called() + separator.logger.warning.assert_called_once() + + +def test_regional_compile_falls_back_to_eager_when_compilation_fails(): + separator = _separator() + transformer_blocks = [transformer for layer in separator.model_run.layers for transformer in layer] + existing_call = Mock() + transformer_blocks[0]._compiled_call_impl = existing_call + compile_calls = 0 + + def compile_then_fail(module): + nonlocal compile_calls + compile_calls += 1 + if compile_calls == 1: + module._compiled_call_impl = Mock() + return + raise RuntimeError("unsupported") + + with patch.object(torch.nn.Module, "compile", autospec=True, side_effect=compile_then_fail): + separator._configure_model_compilation() + + assert separator.is_torch_compiled is False + assert transformer_blocks[0]._compiled_call_impl is existing_call + assert transformer_blocks[1]._compiled_call_impl is None + separator.logger.warning.assert_called_once() + + +def test_regional_compile_restores_original_calls_after_lazy_failure(): + separator = _separator() + transformer_blocks = separator._regional_compile_targets() + original_calls = [Mock(), Mock()] + for transformer, original_call in zip(transformer_blocks, original_calls, strict=True): + transformer._compiled_call_impl = original_call + + def install_compiled_call(module): + module._compiled_call_impl = Mock() + + with patch.object(torch.nn.Module, "compile", autospec=True, side_effect=install_compiled_call): + separator._configure_model_compilation() + + assert separator.is_torch_compiled is True + assert all( + transformer._compiled_call_impl is not original_call + for transformer, original_call in zip(transformer_blocks, original_calls, strict=True) + ) + + expected = torch.ones(1, 2, 8) + separator.model_run.side_effect = [RuntimeError("backend compilation failed"), (expected,)] + + assert separator._run_roformer_model(torch.zeros(2, 8)) is expected + assert separator.is_torch_compiled is False + assert all( + transformer._compiled_call_impl is original_call + for transformer, original_call in zip(transformer_blocks, original_calls, strict=True) + ) + separator.logger.warning.assert_called_once() + + +def test_regional_compile_retries_eager_when_lazy_compilation_fails(): + separator = _separator() + transformer_blocks = separator._regional_compile_targets() + for transformer in transformer_blocks: + transformer._compiled_call_impl = Mock() + separator.is_torch_compiled = True + expected = torch.ones(1, 2, 8) + separator.model_run.side_effect = [RuntimeError("backend compilation failed"), (expected,)] + + result = separator._run_roformer_model(torch.zeros(2, 8)) + + assert result is expected + assert separator.model_run.call_count == 2 + assert separator.is_torch_compiled is False + assert all(transformer._compiled_call_impl is None for transformer in transformer_blocks) + separator.logger.warning.assert_called_once() + + +def test_regional_compile_can_fall_back_after_an_earlier_successful_forward(): + separator = _separator() + transformer_blocks = separator._regional_compile_targets() + for transformer in transformer_blocks: + transformer._compiled_call_impl = Mock() + separator.is_torch_compiled = True + first = torch.ones(1, 2, 8) + fallback = torch.full((1, 2, 8), 2.0) + separator.model_run.side_effect = [(first,), RuntimeError("recompile failed"), (fallback,)] + + assert separator._run_roformer_model(torch.zeros(2, 8)) is first + assert separator._run_roformer_model(torch.zeros(2, 8)) is fallback + + assert separator.model_run.call_count == 3 + assert separator.is_torch_compiled is False + assert all(transformer._compiled_call_impl is None for transformer in transformer_blocks) + separator.logger.warning.assert_called_once() + + +def test_regional_compile_preserves_an_eager_retry_error(): + separator = _separator() + for transformer in separator._regional_compile_targets(): + transformer._compiled_call_impl = Mock() + separator.is_torch_compiled = True + separator.model_run.side_effect = [RuntimeError("compiled path failed"), ValueError("model failed")] + + with pytest.raises(ValueError, match="model failed"): + separator._run_roformer_model(torch.zeros(2, 8)) + + assert separator.model_run.call_count == 2 + assert separator.is_torch_compiled is False + separator.logger.warning.assert_not_called() + + +def test_eager_roformer_errors_are_not_retried(): + separator = _separator() + separator.is_torch_compiled = False + separator.model_run.side_effect = RuntimeError("model failed") + + with pytest.raises(RuntimeError, match="model failed"): + separator._run_roformer_model(torch.zeros(2, 8)) + + separator.model_run.assert_called_once() + separator.logger.warning.assert_not_called() diff --git a/tests/unit/test_roformer_dml_forward.py b/tests/unit/test_roformer_dml_forward.py index b33e1d3..b1ea6df 100644 --- a/tests/unit/test_roformer_dml_forward.py +++ b/tests/unit/test_roformer_dml_forward.py @@ -11,9 +11,11 @@ import pytest import torch from unittest.mock import patch +from packaging import version 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 +from audio_separator.separator.uvr_lib_v5.roformer import rotary as rotary_mod def _tiny_bs_roformer(): @@ -140,6 +142,59 @@ def test_sliced_attention_matches_unsliced(self): assert torch.allclose(unsliced, sliced, atol=1e-6), "batch-sliced attention diverges" +class TestAttendCompilation: + @pytest.mark.parametrize( + ("config", "expected"), + [ + ( + (True, True, True), + [ + torch.nn.attention.SDPBackend.FLASH_ATTENTION, + torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, + torch.nn.attention.SDPBackend.MATH, + torch.nn.attention.SDPBackend.CUDNN_ATTENTION, + ], + ), + ( + (True, False, False), + [ + torch.nn.attention.SDPBackend.FLASH_ATTENTION, + torch.nn.attention.SDPBackend.CUDNN_ATTENTION, + ], + ), + ( + (False, True, True), + [ + torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, + torch.nn.attention.SDPBackend.MATH, + torch.nn.attention.SDPBackend.CUDNN_ATTENTION, + ], + ), + ], + ) + def test_sdpa_backend_translation_preserves_legacy_flags(self, config, expected): + from audio_separator.separator.uvr_lib_v5.roformer import attend as attend_mod + + assert attend_mod._sdpa_backends(attend_mod.FlashAttentionConfig(*config)) == expected + + def test_sdpa_path_has_no_dynamo_graph_break(self): + from audio_separator.separator.uvr_lib_v5.roformer import attend as attend_mod + + if version.parse(torch.__version__.split("+")[0]) < version.parse("2.6"): + pytest.skip("Dynamo learned to trace torch.nn.attention.sdpa_kernel in PyTorch 2.6") + + torch.manual_seed(0) + att = attend_mod.Attend(flash=True).eval() + q = torch.randn(2, 4, 16, 32) + k = torch.randn(2, 4, 16, 32) + v = torch.randn(2, 4, 16, 32) + + explanation = torch._dynamo.explain(att)(q, k, v) + + assert explanation.graph_count == 1 + assert explanation.graph_break_count == 0 + + class TestRotaryNoCatEquivalence: """The DML rotary path computes t*cos + rotate_half(t)*sin directly, skipping the (empty) edge concat torch-directml rejects. It must match @@ -153,9 +208,9 @@ def test_manual_rotation_matches_library(self): t = torch.randn(2, 8, 16, 64) library = rotary.rotate_queries_or_keys(t) - with patch.object(bs_mod, "_is_dml_device", return_value=True): + with patch.object(rotary_mod, "_is_dml_device", return_value=True): manual_bs = bs_mod._rotate_queries_or_keys(rotary, t) - with patch.object(mel_mod, "_is_dml_device", return_value=True): + with patch.object(rotary_mod, "_is_dml_device", return_value=True): manual_mel = mel_mod._rotate_queries_or_keys(rotary, t) assert torch.allclose(library, manual_bs, atol=1e-6) diff --git a/tests/unit/test_roformer_rotary.py b/tests/unit/test_roformer_rotary.py new file mode 100644 index 0000000..98dfb34 --- /dev/null +++ b/tests/unit/test_roformer_rotary.py @@ -0,0 +1,94 @@ +from unittest.mock import patch + +import pytest +import torch +from rotary_embedding_torch import RotaryEmbedding +from rotary_embedding_torch.rotary_embedding_torch import rotate_half + +from audio_separator.separator.uvr_lib_v5 import device_utils +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 +from audio_separator.separator.uvr_lib_v5.roformer import rotary as rotary_mod + + +def _float32_reference(rotary: RotaryEmbedding, tensor: torch.Tensor) -> torch.Tensor: + positions = torch.arange(tensor.shape[-2], device=tensor.device, dtype=torch.float32) + angles = torch.einsum("n, f -> n f", positions, rotary.freqs.float()) + angles = torch.repeat_interleave(angles, 2, dim=-1) + rotated = tensor.float() * angles.cos() + rotate_half(tensor.float()) * angles.sin() + return rotated.to(tensor.dtype) + + +def _seed_low_precision_cache(rotary: RotaryEmbedding, *, seq_len: int, device: torch.device, dtype: torch.dtype) -> None: + rotary.cached_freqs = torch.zeros(seq_len, rotary.freqs.numel() * 2, device=device, dtype=dtype) + + +def test_mel_and_bs_use_the_shared_rotary_helper(): + assert mel_mod._rotate_queries_or_keys is rotary_mod.rotate_queries_or_keys + assert bs_mod._rotate_queries_or_keys is rotary_mod.rotate_queries_or_keys + + +def test_cpu_autocast_replaces_low_precision_cache_and_preserves_input_dtype(): + torch.manual_seed(0) + rotary = RotaryEmbedding(dim=64) + tensor = torch.randn(2, 4, 1101, 64, dtype=torch.bfloat16) + _seed_low_precision_cache(rotary, seq_len=1101, device=tensor.device, dtype=torch.bfloat16) + + with torch.autocast(device_type="cpu", dtype=torch.bfloat16): + actual = rotary_mod.rotate_queries_or_keys(rotary, tensor) + + expected = _float32_reference(rotary, tensor) + assert actual.dtype == tensor.dtype + assert rotary.cached_freqs.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + cached_freqs = rotary.cached_freqs + with torch.autocast(device_type="cpu", dtype=torch.bfloat16): + repeated = rotary_mod.rotate_queries_or_keys(rotary, tensor) + + assert rotary.cached_freqs is cached_freqs + torch.testing.assert_close(repeated, expected, rtol=0, atol=0) + + +def test_compiled_rotation_bypasses_mutable_frequency_cache(): + torch.manual_seed(0) + rotary = RotaryEmbedding(dim=8) + tensor = torch.randn(1, 2, 11, 8) + cached_freqs = torch.zeros(11, 8) + rotary.cached_freqs = cached_freqs + + with patch.object(torch.compiler, "is_compiling", return_value=True): + actual = rotary_mod.rotate_queries_or_keys(rotary, tensor) + + assert rotary.cached_freqs is cached_freqs + torch.testing.assert_close(actual, _float32_reference(rotary, tensor)) + + +def test_rotation_does_not_open_autocast_for_an_unsupported_backend(): + torch.manual_seed(0) + rotary = RotaryEmbedding(dim=8) + tensor = torch.randn(1, 4, 8) + + with ( + patch.object(device_utils, "supports_autocast", return_value=False), + patch.object(device_utils.torch, "autocast", side_effect=AssertionError("autocast must not be opened")), + ): + actual = rotary_mod.rotate_queries_or_keys(rotary, tensor) + + assert actual.shape == tensor.shape + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available") +def test_mps_autocast_uses_float32_angles_and_replaces_low_precision_cache(): + torch.manual_seed(0) + rotary = RotaryEmbedding(dim=64).to("mps") + tensor = torch.randn(2, 4, 1101, 64, device="mps", dtype=torch.float16) + _seed_low_precision_cache(rotary, seq_len=1101, device=tensor.device, dtype=torch.float16) + + with torch.autocast(device_type="mps", dtype=torch.float16): + actual = rotary_mod.rotate_queries_or_keys(rotary, tensor) + + expected = _float32_reference(rotary, tensor) + assert actual.dtype == tensor.dtype + assert rotary.cached_freqs.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/unit/test_separator_api_compatibility.py b/tests/unit/test_separator_api_compatibility.py new file mode 100644 index 0000000..2875957 --- /dev/null +++ b/tests/unit/test_separator_api_compatibility.py @@ -0,0 +1,38 @@ +import inspect + +from audio_separator.separator import Separator + + +def test_execution_options_are_appended_to_constructor_signature(): + parameters = list(inspect.signature(Separator.__init__).parameters) + + assert parameters == [ + "self", + "log_level", + "log_formatter", + "model_file_dir", + "output_dir", + "output_format", + "output_bitrate", + "normalization_threshold", + "amplification_threshold", + "output_single_stem", + "invert_using_spec", + "sample_rate", + "use_soundfile", + "use_autocast", + "use_directml", + "chunk_duration", + "mdx_params", + "vr_params", + "demucs_params", + "mdxc_params", + "ensemble_algorithm", + "ensemble_weights", + "ensemble_preset", + "info_only", + "use_torch_compile", + "use_native_fp16", + ] + assert inspect.signature(Separator.__init__).parameters["use_torch_compile"].default is False + assert inspect.signature(Separator.__init__).parameters["use_native_fp16"].default is False