Use GPU processing for UVR5 and input audio loading and resampling where possible to improve inference efficiency and reduce CPU usage

i18n/locale/en_US.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[Index training][Skipped] trained index already exists: %s",
     "当前设备：%s | 推理精度：%s": "Current device: %s | Inference precision: %s",
     "正在预热CUDA Graph": "Warming up CUDA Graph",
-    "CUDA Graph预热完成": "CUDA Graph warm-up complete"
+    "CUDA Graph预热完成": "CUDA Graph warm-up complete",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Data extraction started: start_time=%.6f, requested concurrency=%s, actual concurrency limit=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Data extraction finished: end_time=%.6f, total elapsed=%.3f seconds"
 }

i18n/locale/es_ES.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[Entrenamiento de índice][Omitido] El índice trained ya existe: %s",
     "当前设备：%s | 推理精度：%s": "Dispositivo actual: %s | Precisión de inferencia: %s",
     "正在预热CUDA Graph": "Preparando CUDA Graph",
-    "CUDA Graph预热完成": "Preparación de CUDA Graph completada"
+    "CUDA Graph预热完成": "Preparación de CUDA Graph completada",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Extracción de datos iniciada: start_time=%.6f, concurrencia solicitada=%s, límite real=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extracción de datos finalizada: end_time=%.6f, tiempo total=%.3f segundos"
 }

i18n/locale/fr_FR.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[Entraînement de l’index][Ignoré] L’index trained existe déjà : %s",
     "当前设备：%s | 推理精度：%s": "Périphérique actuel : %s | Précision d’inférence : %s",
     "正在预热CUDA Graph": "Préchauffage de CUDA Graph",
-    "CUDA Graph预热完成": "Préchauffage de CUDA Graph terminé"
+    "CUDA Graph预热完成": "Préchauffage de CUDA Graph terminé",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Extraction des données démarrée : start_time=%.6f, parallélisme demandé=%s, limite réelle=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extraction des données terminée : end_time=%.6f, durée totale=%.3f secondes"
 }

i18n/locale/it_IT.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[Addestramento indice][Saltato] L’indice trained esiste già: %s",
     "当前设备：%s | 推理精度：%s": "Dispositivo corrente: %s | Precisione di inferenza: %s",
     "正在预热CUDA Graph": "Riscaldamento di CUDA Graph",
-    "CUDA Graph预热完成": "Riscaldamento di CUDA Graph completato"
+    "CUDA Graph预热完成": "Riscaldamento di CUDA Graph completato",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Estrazione dati avviata: start_time=%.6f, concorrenza richiesta=%s, limite effettivo=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Estrazione dati completata: end_time=%.6f, tempo totale=%.3f secondi"
 }

i18n/locale/ja_JP.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[インデックス学習][スキップ] trainedインデックスは既に存在します：%s",
     "当前设备：%s | 推理精度：%s": "現在のデバイス：%s | 推論精度：%s",
     "正在预热CUDA Graph": "CUDA Graphをウォームアップしています",
-    "CUDA Graph预热完成": "CUDA Graphのウォームアップが完了しました"
+    "CUDA Graph预热完成": "CUDA Graphのウォームアップが完了しました",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "データ抽出開始：start_time=%.6f、要求並列数=%s、実並列上限=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "データ抽出完了：end_time=%.6f、合計所要時間=%.3f秒"
 }

i18n/locale/ko_KR.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[인덱스 학습][건너뜀] trained 인덱스가 이미 있습니다: %s",
     "当前设备：%s | 推理精度：%s": "현재 장치: %s | 추론 정밀도: %s",
     "正在预热CUDA Graph": "CUDA Graph 워밍업 중",
-    "CUDA Graph预热完成": "CUDA Graph 워밍업 완료"
+    "CUDA Graph预热完成": "CUDA Graph 워밍업 완료",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "데이터 추출 시작: start_time=%.6f, 요청 병렬 수=%s, 실제 병렬 상한=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "데이터 추출 완료: end_time=%.6f, 총 소요 시간=%.3f초"
 }

i18n/locale/pt_BR.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[Treinamento de índice][Ignorado] O índice trained já existe: %s",
     "当前设备：%s | 推理精度：%s": "Dispositivo atual: %s | Precisão de inferência: %s",
     "正在预热CUDA Graph": "Aquecendo o CUDA Graph",
-    "CUDA Graph预热完成": "Aquecimento do CUDA Graph concluído"
+    "CUDA Graph预热完成": "Aquecimento do CUDA Graph concluído",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Extração de dados iniciada: start_time=%.6f, concorrência solicitada=%s, limite real=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extração de dados concluída: end_time=%.6f, tempo total=%.3f segundos"
 }

i18n/locale/ru_RU.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[Обучение индекса][Пропущено] Индекс trained уже существует: %s",
     "当前设备：%s | 推理精度：%s": "Текущее устройство: %s | Точность вывода: %s",
     "正在预热CUDA Graph": "Прогрев CUDA Graph",
-    "CUDA Graph预热完成": "Прогрев CUDA Graph завершён"
+    "CUDA Graph预热完成": "Прогрев CUDA Graph завершён",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Извлечение данных начато: start_time=%.6f, запрошенный параллелизм=%s, фактический предел=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Извлечение данных завершено: end_time=%.6f, общее время=%.3f с"
 }

i18n/locale/tr_TR.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[Dizin eğitimi][Atlandı] trained dizini zaten var: %s",
     "当前设备：%s | 推理精度：%s": "Geçerli cihaz: %s | Çıkarım hassasiyeti: %s",
     "正在预热CUDA Graph": "CUDA Graph ısınıyor",
-    "CUDA Graph预热完成": "CUDA Graph ısınması tamamlandı"
+    "CUDA Graph预热完成": "CUDA Graph ısınması tamamlandı",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Veri çıkarma başladı: start_time=%.6f, istenen eşzamanlılık=%s, gerçek üst sınır=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Veri çıkarma tamamlandı: end_time=%.6f, toplam süre=%.3f saniye"
 }

i18n/locale/zh_CN.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[索引训练][跳过] trained索引已存在：%s",
     "当前设备：%s | 推理精度：%s": "当前设备：%s | 推理精度：%s",
     "正在预热CUDA Graph": "正在预热CUDA Graph",
-    "CUDA Graph预热完成": "CUDA Graph预热完成"
+    "CUDA Graph预热完成": "CUDA Graph预热完成",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "数据提取结束：end_time=%.6f，总耗时=%.3f秒"
 }

i18n/locale/zh_HK.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[索引訓練][跳過] trained索引已存在：%s",
     "当前设备：%s | 推理精度：%s": "目前裝置：%s | 推理精度：%s",
     "正在预热CUDA Graph": "正在預熱 CUDA Graph",
-    "CUDA Graph预热完成": "CUDA Graph 預熱完成"
+    "CUDA Graph预热完成": "CUDA Graph 預熱完成",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "資料擷取開始：start_time=%.6f，請求並行數=%s，實際並行數上限=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "資料擷取結束：end_time=%.6f，總耗時=%.3f秒"
 }

i18n/locale/zh_SG.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[索引训练][跳过] trained索引已存在：%s",
     "当前设备：%s | 推理精度：%s": "当前设备：%s | 推理精度：%s",
     "正在预热CUDA Graph": "正在预热CUDA Graph",
-    "CUDA Graph预热完成": "CUDA Graph预热完成"
+    "CUDA Graph预热完成": "CUDA Graph预热完成",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "数据提取结束：end_time=%.6f，总耗时=%.3f秒"
 }

i18n/locale/zh_TW.json

@@ -259,5 +259,7 @@
     "[索引训练][跳过] trained索引已存在：%s": "[索引訓練][跳過] trained索引已存在：%s",
     "当前设备：%s | 推理精度：%s": "目前裝置：%s | 推論精度：%s",
     "正在预热CUDA Graph": "正在預熱 CUDA Graph",
-    "CUDA Graph预热完成": "CUDA Graph 預熱完成"
+    "CUDA Graph预热完成": "CUDA Graph 預熱完成",
+    "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "資料擷取開始：start_time=%.6f，請求並行數=%s，實際並行數上限=%s",
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "資料擷取結束：end_time=%.6f，總耗時=%.3f秒"
 }

infer/audio.py

@@ -3,6 +3,43 @@ import ffmpeg
 import numpy as np
 import av
 from io import BytesIO
+import threading
+
+
+_USE_TORCHAUDIO_GPU = False
+_AUDIO_DEVICE = None
+_AUDIO_DTYPE = None
+_TORCH = None
+_TORCHAUDIO = None
+_TORCHAUDIO_RESAMPLE = None
+_RESAMPLE_TRANSFORMS = {}
+_RESAMPLE_LOCK = threading.Lock()
+_FORCE_CPU_AUDIO = os.environ.get("RVC_AUDIO_FORCE_CPU", "0") == "1"
+
+# Select the accelerated loader once, when this module is imported.  The CUDA
+# device and dtype come from the project's shared automatic selection rules.
+# CPU and DirectML keep the original FFmpeg path.  Import failures (including
+# missing torchaudio DLLs) also leave FFmpeg selected.
+if not _FORCE_CPU_AUDIO:
+    try:
+        import torch as _TORCH
+        import torchaudio as _TORCHAUDIO
+        from torchaudio.transforms import Resample as _TORCHAUDIO_RESAMPLE
+
+        from configs.config import infer_device as _AUDIO_DEVICE
+        from configs.config import infer_dtype as _AUDIO_DTYPE
+
+        _USE_TORCHAUDIO_GPU = (
+            getattr(_AUDIO_DEVICE, "type", None) == "cuda"
+            and _TORCH.cuda.is_available()
+        )
+    except Exception:
+        _USE_TORCHAUDIO_GPU = False
+
+AUDIO_LOAD_BACKEND = "torchaudio_cuda" if _USE_TORCHAUDIO_GPU else "ffmpeg"
+TORCHAUDIO_GPU_ENABLED = _USE_TORCHAUDIO_GPU
+AUDIO_DEVICE = _AUDIO_DEVICE
+AUDIO_DTYPE = _AUDIO_DTYPE
 
 
 def wav2(i, o, format):
@@ -28,21 +65,243 @@ def wav2(i, o, format):
     inp.close()
 
 
-def load_audio(file, sr):
-    try:
-        # https://github.com/openai/whisper/blob/main/whisper/audio.py#L26
-        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
-        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
-        file = clean_path(file)  # 防止小白拷路径头尾带了空格和"和回车
+def _probe_audio(file):
+    info = ffmpeg.probe(file, cmd="ffprobe")
+    stream = next(
+        item for item in info.get("streams", []) if item.get("codec_type") == "audio"
+    )
+    return int(stream["sample_rate"]), int(stream.get("channels", 1))
+
+
+def _decode_audio_ffmpeg(file):
+    """Decode one audio stream without changing its sample rate."""
+    source_sr, channels = _probe_audio(file)
+    channels = max(1, channels)
     out, _ = (
         ffmpeg.input(file, threads=0)
-            .output("-", format="f32le", acodec="pcm_f32le", ac=1, ar=sr)
+        .output(
+            "-",
+            format="f32le",
+            acodec="pcm_f32le",
+            ac=channels,
+            ar=source_sr,
+        )
         .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
     )
-    except Exception as e:
-        raise RuntimeError(f"Failed to load audio: {e}")
+    samples = np.frombuffer(out, np.float32)
+    usable = samples.size - samples.size % channels
+    samples = samples[:usable]
+    return samples.reshape(-1, channels).T.copy(), source_sr
 
-    return np.frombuffer(out, np.float32).flatten()
+
+def _load_audio_ffmpeg(file, sr, force_mono=True):
+    # https://github.com/openai/whisper/blob/main/whisper/audio.py#L26
+    # Keep the original FFmpeg path for CPU and unsupported decoders.  When
+    # stereo is requested, retain all source channels instead of forcing ac=1.
+    channels = 1
+    if not force_mono:
+        _, channels = _probe_audio(file)
+        channels = max(1, channels)
+    output = {
+        "format": "f32le",
+        "acodec": "pcm_f32le",
+        "ar": sr,
+    }
+    if force_mono:
+        output["ac"] = 1
+    else:
+        output["ac"] = channels
+    out, _ = (
+        ffmpeg.input(file, threads=0)
+        .output("-", **output)
+        .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
+    )
+    samples = np.frombuffer(out, np.float32)
+    if force_mono:
+        return samples.flatten()
+    usable = samples.size - samples.size % channels
+    return samples[:usable].reshape(-1, channels).T.copy()
+
+
+def _channel_first_tensor(audio):
+    if not _TORCH.is_tensor(audio):
+        audio = _TORCH.from_numpy(np.ascontiguousarray(audio, dtype=np.float32))
+    if audio.ndim == 1:
+        audio = audio.unsqueeze(0)
+    if audio.ndim != 2:
+        raise ValueError("Audio data must be one-dimensional or channel-first two-dimensional")
+    return audio
+
+
+def _format_audio_tensor(audio, force_mono=True, keep_on_device=False):
+    audio = _channel_first_tensor(audio)
+    if force_mono:
+        if audio.shape[0] > 1:
+            audio = audio.mean(dim=0, keepdim=True)
+        else:
+            audio = audio[:1]
+    if keep_on_device:
+        audio = audio.to(device=_AUDIO_DEVICE, dtype=_AUDIO_DTYPE)
+        return audio.flatten() if force_mono else audio
+    audio = audio.detach().to(device="cpu", dtype=_TORCH.float32).contiguous().numpy()
+    return audio[0].flatten() if force_mono else audio
+
+
+def _get_gpu_resampler(source_sr, target_sr):
+    key = (source_sr, target_sr, str(_AUDIO_DEVICE), _AUDIO_DTYPE)
+    transform = _RESAMPLE_TRANSFORMS.get(key)
+    if transform is None:
+        with _RESAMPLE_LOCK:
+            transform = _RESAMPLE_TRANSFORMS.get(key)
+            if transform is None:
+                transform = _TORCHAUDIO_RESAMPLE(
+                    source_sr,
+                    target_sr,
+                    dtype=_AUDIO_DTYPE,
+                ).to(device=_AUDIO_DEVICE, dtype=_AUDIO_DTYPE)
+                _RESAMPLE_TRANSFORMS[key] = transform
+    return transform
+
+
+def _resample_tensor_gpu(audio, source_sr, target_sr, force_mono=True, keep_on_device=False):
+    audio = _channel_first_tensor(audio)
+    if source_sr == target_sr:
+        return _format_audio_tensor(audio, force_mono, keep_on_device)
+    with _TORCH.inference_mode():
+        audio = audio.to(device=_AUDIO_DEVICE, dtype=_AUDIO_DTYPE)
+        if force_mono and audio.shape[0] > 1:
+            audio = audio.mean(dim=0, keepdim=True)
+        elif force_mono:
+            audio = audio[:1]
+        audio = _get_gpu_resampler(source_sr, target_sr)(audio)
+        if keep_on_device:
+            return audio.flatten() if force_mono else audio
+        audio = audio.detach().to(device="cpu", dtype=_TORCH.float32).contiguous().numpy()
+        return audio[0].flatten() if force_mono else audio
+
+
+def _load_audio_torchaudio_gpu(file, sr, force_mono=True, keep_on_device=False):
+    audio, source_sr = _TORCHAUDIO.load(file)
+    return _resample_tensor_gpu(audio, source_sr, sr, force_mono, keep_on_device)
+
+
+def resample_audio(audio, source_sr, target_sr, force_mono=False, res_type=None):
+    """Resample channel-first audio, using the same CUDA path as load_audio."""
+    tensor = _channel_first_tensor(audio)
+    if _USE_TORCHAUDIO_GPU:
+        try:
+            return _resample_tensor_gpu(
+                tensor, source_sr, target_sr, force_mono, keep_on_device=False
+            )
+        except Exception:
+            try:
+                _TORCH.cuda.empty_cache()
+            except Exception:
+                pass
+
+    tensor = _channel_first_tensor(tensor)
+    if force_mono:
+        if tensor.shape[0] > 1:
+            tensor = tensor.mean(dim=0, keepdim=True)
+        else:
+            tensor = tensor[:1]
+    array = tensor.detach().to(device="cpu", dtype=_TORCH.float32).numpy()
+    if source_sr != target_sr:
+        import librosa
+
+        kwargs = {}
+        if res_type is not None:
+            kwargs["res_type"] = res_type
+        array = librosa.resample(
+            array,
+            orig_sr=source_sr,
+            target_sr=target_sr,
+            axis=-1,
+            **kwargs,
+        )
+    return array[0].flatten() if force_mono else array
+
+
+def resample_audio_tensor(audio, source_sr, target_sr, force_mono=False):
+    """Resample channel-first audio and keep the result on the selected GPU."""
+    if _USE_TORCHAUDIO_GPU:
+        try:
+            return _resample_tensor_gpu(
+                audio,
+                source_sr,
+                target_sr,
+                force_mono=force_mono,
+                keep_on_device=True,
+            )
+        except Exception:
+            try:
+                _TORCH.cuda.empty_cache()
+            except Exception:
+                pass
+    result = resample_audio(
+        audio,
+        source_sr,
+        target_sr,
+        force_mono=force_mono,
+    )
+    tensor = _TORCH.from_numpy(np.ascontiguousarray(result, dtype=np.float32))
+    if _USE_TORCHAUDIO_GPU:
+        tensor = tensor.to(device=_AUDIO_DEVICE, dtype=_AUDIO_DTYPE)
+    return tensor
+
+
+def load_audio_tensor(file, sr, force_mono=True):
+    """Load audio as a tensor; stereo mode returns channel-first data."""
+    file = clean_path(file)
+    if not _USE_TORCHAUDIO_GPU:
+        return _TORCH.from_numpy(_load_audio_ffmpeg(file, sr, force_mono))
+    try:
+        return _load_audio_torchaudio_gpu(
+            file, sr, force_mono=force_mono, keep_on_device=True
+        )
+    except Exception as torchaudio_error:
+        try:
+            audio, source_sr = _decode_audio_ffmpeg(file)
+            return _resample_tensor_gpu(
+                audio, source_sr, sr, force_mono=force_mono, keep_on_device=True
+            )
+        except Exception as ffmpeg_error:
+            try:
+                audio = _TORCH.from_numpy(_load_audio_ffmpeg(file, sr, force_mono))
+                return audio.to(device=_AUDIO_DEVICE, dtype=_AUDIO_DTYPE)
+            except Exception:
+                raise RuntimeError(
+                    "Failed to load audio with torchaudio (%s) and FFmpeg (%s)"
+                    % (torchaudio_error, ffmpeg_error)
+                ) from ffmpeg_error
+
+
+def load_audio(file, sr, force_mono=True):
+    """Load float32 audio; mono is [T], preserved channels are [C, T]."""
+    file = clean_path(file)  # 防止小白拷路径头尾带了空格和"和回车
+    if _USE_TORCHAUDIO_GPU:
+        try:
+            return _load_audio_torchaudio_gpu(file, sr, force_mono=force_mono)
+        except Exception as torchaudio_error:
+            try:
+                audio, source_sr = _decode_audio_ffmpeg(file)
+                return _resample_tensor_gpu(
+                    audio, source_sr, sr, force_mono=force_mono
+                )
+            except Exception as decode_error:
+                # Preserve the old format coverage if both torchaudio decode
+                # and GPU processing are unavailable for this file.
+                try:
+                    return _load_audio_ffmpeg(file, sr, force_mono)
+                except Exception as ffmpeg_error:
+                    raise RuntimeError(
+                        "Failed to load audio with torchaudio (%s), raw FFmpeg (%s), and FFmpeg (%s)"
+                        % (torchaudio_error, decode_error, ffmpeg_error)
+                    ) from ffmpeg_error
+    try:
+        return _load_audio_ffmpeg(file, sr, force_mono)
+    except Exception as error:
+        raise RuntimeError("Failed to load audio: %s" % error) from error
 
 
 def clean_path(path_str):

tools/uvr5/bs_roformer/bs_roformer.py

@@ -186,6 +186,7 @@ class BSRoformer(Module):
         self.final_norm = RMSNorm(dim)
         self.stft_kwargs = dict(n_fft=stft_n_fft, hop_length=stft_hop_length, win_length=stft_win_length, normalized=stft_normalized)
         self.stft_window_fn = partial(default(stft_window_fn, torch.hann_window), stft_win_length)
+        self._stft_windows = {}
         freqs = torch.stft(torch.randn(1, 4096), **self.stft_kwargs, window=torch.ones(stft_win_length), return_complex=True).shape[1]
         assert len(freqs_per_bands) > 1
         assert sum(freqs_per_bands) == freqs, f'the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}'
@@ -201,6 +202,14 @@ class BSRoformer(Module):
         self.multi_stft_window_fn = multi_stft_window_fn
         self.multi_stft_kwargs = dict(hop_length=multi_stft_hop_size, normalized=multi_stft_normalized)
 
+    def _get_stft_window(self, device):
+        key = str(device)
+        window = self._stft_windows.get(key)
+        if window is None:
+            window = self.stft_window_fn(device=device, dtype=torch.float32)
+            self._stft_windows[key] = window
+        return window
+
     def forward(self, raw_audio, target=None, return_loss_breakdown=False):
         """
         einops
@@ -224,7 +233,7 @@ class BSRoformer(Module):
         if x_is_dml:
             # DirectML has no complex/STFT kernels.  Keep only the spectral
             # boundary on CPU and move its real representation to DirectML.
-            stft_window = self.stft_window_fn(device='cpu')
+            stft_window = self._get_stft_window('cpu')
             stft_complex = torch.stft(
                 raw_audio.cpu(),
                 **self.stft_kwargs,
@@ -240,7 +249,7 @@ class BSRoformer(Module):
             )
             stft_repr = stft_repr_cpu.to(device)
         else:
-            stft_window = self.stft_window_fn(device=device)
+            stft_window = self._get_stft_window(device)
             try:
                 stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
             except:

tools/uvr5/bs_roformer/mel_band_roformer.py

@@ -189,6 +189,7 @@ class MelBandRoformer(Module):
             tran_modules.append(Transformer(depth=freq_transformer_depth, rotary_embed=freq_rotary_embed, **transformer_kwargs))
             self.layers.append(nn.ModuleList(tran_modules))
         self.stft_window_fn = partial(default(stft_window_fn, torch.hann_window), stft_win_length)
+        self._stft_windows = {}
         self.stft_kwargs = dict(n_fft=stft_n_fft, hop_length=stft_hop_length, win_length=stft_win_length, normalized=stft_normalized)
         freqs = torch.stft(torch.randn(1, 4096), **self.stft_kwargs, window=torch.ones(stft_n_fft), return_complex=True).shape[1]
         mel_filter_bank_numpy = filters.mel(sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands)
@@ -222,6 +223,14 @@ class MelBandRoformer(Module):
         self.multi_stft_kwargs = dict(hop_length=multi_stft_hop_size, normalized=multi_stft_normalized)
         self.match_input_audio_length = match_input_audio_length
 
+    def _get_stft_window(self, device):
+        key = str(device)
+        window = self._stft_windows.get(key)
+        if window is None:
+            window = self.stft_window_fn(device=device, dtype=torch.float32)
+            self._stft_windows[key] = window
+        return window
+
     def forward(self, raw_audio, target=None, return_loss_breakdown=False):
         """
         einops
@@ -245,7 +254,7 @@ class MelBandRoformer(Module):
         if x_is_dml:
             # DirectML has no STFT or complex tensor support.  Build the real
             # spectral features on CPU, then run the learned network on DML.
-            stft_window = self.stft_window_fn(device='cpu')
+            stft_window = self._get_stft_window('cpu')
             stft_complex = torch.stft(
                 raw_audio.cpu(),
                 **self.stft_kwargs,
@@ -257,16 +266,14 @@ class MelBandRoformer(Module):
                 stft_repr, batch_audio_channel_packed_shape, '* f t c'
             )
             stft_repr = rearrange(stft_repr, 'b s f t c -> b (f s) t c')
-            batch_arange = torch.arange(batch)[..., None]
-            x = stft_repr[batch_arange, self.freq_indices.cpu()].to(device)
+            x = stft_repr[:, self.freq_indices.cpu()].to(device)
         else:
-            stft_window = self.stft_window_fn(device=device)
+            stft_window = self._get_stft_window(device)
             stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
             stft_repr = torch.view_as_real(stft_repr)
             stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, '* f t c')
             stft_repr = rearrange(stft_repr, 'b s f t c -> b (f s) t c')
-            batch_arange = torch.arange(batch, device=device)[..., None]
-            x = stft_repr[batch_arange, self.freq_indices]
+            x = stft_repr[:, self.freq_indices]
         x = rearrange(x, 'b f t c -> b t (f c)')
         if self.use_torch_checkpoint:
             x = checkpoint(self.band_split, x, use_reentrant=False)
@@ -316,9 +323,9 @@ class MelBandRoformer(Module):
         masks = torch.view_as_complex(masks.contiguous())
         masks = masks.type(stft_repr.dtype)
         freq_indices = self.freq_indices.cpu() if x_is_dml else self.freq_indices
-        scatter_indices = repeat(freq_indices, 'f -> b n f t', b=batch, n=num_stems, t=stft_repr.shape[-1])
         stft_repr_expanded_stems = repeat(stft_repr, 'b 1 ... -> b n ...', n=num_stems)
-        masks_summed = torch.zeros_like(stft_repr_expanded_stems).scatter_add_(2, scatter_indices, masks)
+        masks_summed = torch.zeros_like(stft_repr_expanded_stems)
+        masks_summed.index_add_(2, freq_indices, masks)
         num_bands_per_freq = self.num_bands_per_freq.cpu() if x_is_dml else self.num_bands_per_freq
         denom = repeat(num_bands_per_freq, 'f -> (f r) 1', r=channels)
         masks_averaged = masks_summed / denom.clamp(min=1e-08)

tools/uvr5/bsroformer.py

@@ -3,14 +3,13 @@ import os
 import warnings
 from contextlib import nullcontext
 
-import librosa
 import numpy as np
 import soundfile as sf
 import torch
 import torch.nn as nn
 import yaml
 
-from tools.cuda_graph import run_cuda_graph
+from infer.audio import TORCHAUDIO_GPU_ENABLED, load_audio, load_audio_tensor
 from tqdm import tqdm
 from tools.file_io import read_text
 from i18n.i18n import I18nAuto
@@ -120,27 +119,67 @@ class Roformer_Loader:
         batch_size = self.config["inference"]["batch_size"]
 
         length_init = mix.shape[-1]
-        progress_bar = tqdm(total=length_init // step + 1, desc="Processing", leave=False)
 
         # Do pad from the beginning and end to account floating window results better
         if length_init > 2 * border and (border > 0):
             mix = nn.functional.pad(mix, (border, border), mode="reflect")
+        total_windows = (mix.shape[-1] + step - 1) // step
+        progress_bar = tqdm(total=total_windows, desc="Processing", leave=False)
 
-        # Prepare windows arrays (do 1 time for speed up). This trick repairs click problems on the edges of segment
-        window_size = C
-        fadein = torch.linspace(0, 1, fade_size)
-        fadeout = torch.linspace(1, 0, fade_size)
-        window_start = torch.ones(window_size)
-        window_middle = torch.ones(window_size)
-        window_finish = torch.ones(window_size)
-        window_start[-fade_size:] *= fadeout  # First audio chunk, no fadein
-        window_finish[:fade_size] *= fadein  # Last audio chunk, no fadeout
+        parsed_device = device if isinstance(device, torch.device) else torch.device(device)
+        device_type = parsed_device.type
+        if self.config["training"]["target_instrument"] is None:
+            source_count = len(self.config["training"]["instruments"])
+        else:
+            source_count = 1
+        req_shape = (source_count,) + tuple(mix.shape)
+
+        accumulation_device = torch.device("cpu")
+        if device_type == "cuda":
+            required_bytes = int(np.prod(req_shape)) * 4 + mix.shape[-1] * 4
+            free_bytes, _ = torch.cuda.mem_get_info(parsed_device)
+            limit = min(1024**3, int(free_bytes * 0.22))
+            if required_bytes <= limit:
+                accumulation_device = parsed_device
+
+        try:
+            result = torch.zeros(
+                req_shape,
+                dtype=torch.float32,
+                device=accumulation_device,
+            )
+            counter = torch.zeros(
+                mix.shape[-1],
+                dtype=torch.float32,
+                device=accumulation_device,
+            )
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            accumulation_device = torch.device("cpu")
+            result = torch.zeros(req_shape, dtype=torch.float32)
+            counter = torch.zeros(mix.shape[-1], dtype=torch.float32)
+
+        # The overlap-add window lives beside the accumulator.  A short file
+        # with one window uses the all-ones window to avoid a zero denominator.
+        fadein = torch.linspace(
+            0, 1, fade_size, device=accumulation_device, dtype=torch.float32
+        )
+        fadeout = torch.linspace(
+            1, 0, fade_size, device=accumulation_device, dtype=torch.float32
+        )
+        window_full = torch.ones(C, device=accumulation_device)
+        window_start = window_full.clone()
+        window_middle = window_full.clone()
+        window_finish = window_full.clone()
+        window_start[-fade_size:] *= fadeout
+        window_finish[:fade_size] *= fadein
         window_middle[-fade_size:] *= fadeout
         window_middle[:fade_size] *= fadein
 
-        device_type = device.type if isinstance(device, torch.device) else torch.device(device).type
         amp_context = (
-            torch.amp.autocast("cuda") if device_type == "cuda" else nullcontext()
+            torch.amp.autocast("cuda", enabled=self.is_half)
+            if device_type == "cuda"
+            else nullcontext()
         )
         grad_context = (
             torch.no_grad()
@@ -152,62 +191,62 @@ class Roformer_Loader:
             # therefore needs no_grad rather than inference_mode. CUDA and CPU
             # retain the existing inference-mode path.
             with grad_context:
-                if self.config["training"]["target_instrument"] is None:
-                    req_shape = (len(self.config["training"]["instruments"]),) + tuple(mix.shape)
-                else:
-                    req_shape = (1,) + tuple(mix.shape)
-
-                result = torch.zeros(req_shape, dtype=torch.float32)
-                counter = torch.zeros(req_shape, dtype=torch.float32)
+                model_dtype = next(model.parameters()).dtype
                 i = 0
                 batch_data = []
                 batch_locations = []
                 while i < mix.shape[1]:
-                    part = mix[:, i : i + C].to(device)
+                    part = mix[:, i : i + C]
                     length = part.shape[-1]
                     if length < C:
                         if length > C // 2 + 1:
                             part = nn.functional.pad(input=part, pad=(0, C - length), mode="reflect")
                         else:
                             part = nn.functional.pad(input=part, pad=(0, C - length, 0, 0), mode="constant", value=0)
-                    if self.is_half:
-                        part = part.half()
                     batch_data.append(part)
                     batch_locations.append((i, length))
                     i += step
                     progress_bar.update(1)
 
                     if len(batch_data) >= batch_size or (i >= mix.shape[1]):
-                        arr = torch.stack(batch_data, dim=0)
-                        # print(23333333,arr.dtype)
-                        x = run_cuda_graph(
-                            model,
-                            "uvr-bsroformer",
-                            lambda audio: model(audio),
-                            arr,
+                        arr = torch.stack(batch_data, dim=0).to(
+                            device=parsed_device,
+                            dtype=model_dtype,
+                        )
+                        # Torch STFT/ISTFT cannot be captured reliably by a
+                        # CUDA Graph on the supported runtime, so keep this
+                        # model call eager while all tensors remain on CUDA.
+                        x = model(arr)
+                        x_for_accumulation = (
+                            x.float()
+                            if accumulation_device.type == "cuda"
+                            else x.float().cpu()
                         )
-
-                        window = window_middle
-                        if i - step == 0:  # First audio chunk, no fadein
-                            window = window_start
-                        elif i >= mix.shape[1]:  # Last audio chunk, no fadeout
-                            window = window_finish
-
                         for j in range(len(batch_locations)):
                             start, l = batch_locations[j]
-                            result[..., start : start + l] += x[j][..., :l].cpu() * window[..., :l]
-                            counter[..., start : start + l] += window[..., :l]
+                            is_first = start == 0
+                            is_last = start + l >= mix.shape[1]
+                            if is_first and is_last:
+                                window = window_full
+                            elif is_first:
+                                window = window_start
+                            elif is_last:
+                                window = window_finish
+                            else:
+                                window = window_middle
+                            result[..., start : start + l].add_(
+                                x_for_accumulation[j][..., :l] * window[:l]
+                            )
+                            counter[start : start + l].add_(window[:l])
 
                         batch_data = []
                         batch_locations = []
 
-                estimated_sources = result / counter
-                estimated_sources = estimated_sources.cpu().numpy()
-                np.nan_to_num(estimated_sources, copy=False, nan=0.0)
-
+                result.div_(counter.clamp_min(1e-8))
+                torch.nan_to_num_(result)
                 if length_init > 2 * border and (border > 0):
-                    # Remove pad
-                    estimated_sources = estimated_sources[..., border:-border]
+                    result = result[..., border:-border]
+                estimated_sources = result.cpu().numpy()
 
         progress_bar.close()
 
@@ -227,22 +266,63 @@ class Roformer_Loader:
         if "sample_rate" in self.config["audio"]:
             sample_rate = self.config["audio"]["sample_rate"]
 
+        isstereo = self.config["model"].get("stereo", True)
+        device_type = (
+            self.device.type
+            if isinstance(self.device, torch.device)
+            else torch.device(self.device).type
+        )
         try:
-            mix, sr = librosa.load(path, sr=sample_rate, mono=False)
+            if device_type == "cuda" and TORCHAUDIO_GPU_ENABLED:
+                mix = load_audio_tensor(
+                    path, sample_rate, force_mono=not isstereo
+                )
+            else:
+                mix = load_audio(path, sample_rate, force_mono=not isstereo)
+            sr = sample_rate
         except Exception as e:
             print(i18n("无法读取音频：%s") % path)
             print(i18n("错误信息：%s") % str(e))
             return
 
-        # in case if model only supports mono tracks
-        isstereo = self.config["model"].get("stereo", True)
-        if not isstereo and len(mix.shape) != 1:
-            mix = np.mean(mix, axis=0)  # if more than 2 channels, take mean
+        if isstereo:
+            if mix.ndim == 1:
+                mix = mix.unsqueeze(0) if torch.is_tensor(mix) else mix[np.newaxis, :]
+            if mix.shape[0] == 1:
+                mix = mix.repeat(2, 1) if torch.is_tensor(mix) else np.repeat(mix, 2, axis=0)
+            elif mix.shape[0] > 2:
+                mix = mix[:2].contiguous() if torch.is_tensor(mix) else np.ascontiguousarray(mix[:2])
+        else:
+            if mix.ndim == 1:
+                mix = mix.unsqueeze(0) if torch.is_tensor(mix) else mix[np.newaxis, :]
+            elif mix.shape[0] > 1:
+                mix = (
+                    mix.mean(dim=0, keepdim=True)
+                    if torch.is_tensor(mix)
+                    else np.mean(mix, axis=0, keepdims=True)
+                )
                 print(i18n("音频包含多个声道，但模型仅支持单声道，将对所有声道取平均值"))
 
-        mix_orig = mix.copy()
-
-        mixture = torch.tensor(mix, dtype=torch.float32)
+        if torch.is_tensor(mix):
+            keep_on_gpu = mix.device.type == "cuda"
+            if keep_on_gpu:
+                free_bytes, _ = torch.cuda.mem_get_info(mix.device)
+                input_bytes = mix.numel() * mix.element_size()
+                keep_on_gpu = input_bytes <= min(
+                    512 * 1024 * 1024,
+                    int(free_bytes * 0.10),
+                )
+            if keep_on_gpu:
+                mixture = mix
+                mix_orig = mix.detach().float().cpu().numpy()
+            else:
+                mixture = mix.detach().float().cpu()
+                mix_orig = mixture.numpy()
+                del mix
+        else:
+            mix = np.ascontiguousarray(mix, dtype=np.float32)
+            mix_orig = mix
+            mixture = torch.from_numpy(mix)
         res = self.demix_track(self.model, mixture, self.device)
 
         if self.config["training"]["target_instrument"] is not None:
@@ -250,7 +330,8 @@ class Roformer_Loader:
             # other instruments are caculated by subtracting target instrument from mixture
             target_instrument = self.config["training"]["target_instrument"]
             other_instruments = [i for i in self.config["training"]["instruments"] if i != target_instrument]
-            other = mix_orig - res[target_instrument]  # caculate other instruments
+            np.subtract(mix_orig, res[target_instrument], out=mix_orig)
+            other = mix_orig
 
             path_vocal = "{}/{}_{}.wav".format(vocal_root, file_base_name, target_instrument)
             path_other = "{}/{}_{}.wav".format(others_root, file_base_name, other_instruments[0])

tools/uvr5/lib/lib_v5/spec_utils.py

@@ -1,12 +1,55 @@
-import hashlib
-import json
 import math
-import os
 
 import librosa
 import numpy as np
-import soundfile as sf
-from tqdm import tqdm
+import torch
+from infer.audio import resample_audio, resample_audio_tensor
+
+
+_STFT_WINDOWS = {}
+
+
+def _stft_window(n_fft, device):
+    key = (n_fft, str(device))
+    window = _STFT_WINDOWS.get(key)
+    if window is None:
+        window = torch.hann_window(
+            n_fft,
+            periodic=True,
+            device=device,
+            dtype=torch.float32,
+        )
+        _STFT_WINDOWS[key] = window
+    return window
+
+
+def _wave_to_spectrogram_torch(
+    wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False
+):
+    wave = wave.to(dtype=torch.float32)
+    if reverse:
+        transformed = torch.flip(wave[:2], dims=(-1,))
+    elif mid_side:
+        transformed = torch.stack(
+            ((wave[0] + wave[1]) / 2, wave[0] - wave[1])
+        )
+    elif mid_side_b2:
+        transformed = torch.stack(
+            (wave[1] + wave[0] * 0.5, wave[0] - wave[1] * 0.5)
+        )
+    else:
+        transformed = wave[:2]
+    return torch.stft(
+        transformed,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        window=_stft_window(n_fft, transformed.device),
+        center=True,
+        pad_mode="constant",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
 
 
 def crop_center(h1, h2):
@@ -27,29 +70,11 @@ def crop_center(h1, h2):
     return h1
 
 
-def wave_to_spectrogram(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False):
-    if reverse:
-        wave_left = np.flip(np.asfortranarray(wave[0]))
-        wave_right = np.flip(np.asfortranarray(wave[1]))
-    elif mid_side:
-        wave_left = np.asfortranarray(np.add(wave[0], wave[1]) / 2)
-        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1]))
-    elif mid_side_b2:
-        wave_left = np.asfortranarray(np.add(wave[1], wave[0] * 0.5))
-        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1] * 0.5))
-    else:
-        wave_left = np.asfortranarray(wave[0])
-        wave_right = np.asfortranarray(wave[1])
-
-    spec_left = librosa.stft(wave_left, n_fft=n_fft, hop_length=hop_length)
-    spec_right = librosa.stft(wave_right, n_fft=n_fft, hop_length=hop_length)
-
-    spec = np.asfortranarray([spec_left, spec_right])
-
-    return spec
-
-
 def wave_to_spectrogram_mt(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False):
+    if torch.is_tensor(wave):
+        return _wave_to_spectrogram_torch(
+            wave, hop_length, n_fft, mid_side, mid_side_b2, reverse
+        )
     import threading
 
     if reverse:
@@ -84,6 +109,14 @@ def wave_to_spectrogram_mt(wave, hop_length, n_fft, mid_side=False, mid_side_b2=
 
 def combine_spectrograms(specs, mp):
     l = min([specs[i].shape[2] for i in specs])
+    first = specs[next(iter(specs))]
+    if torch.is_tensor(first):
+        spec_c = torch.zeros(
+            (2, mp.param["bins"] + 1, l),
+            dtype=torch.complex64,
+            device=first.device,
+        )
+    else:
         spec_c = np.zeros(shape=(2, mp.param["bins"] + 1, l), dtype=np.complex64)
     offset = 0
     bands_n = len(mp.param["band"])
@@ -109,48 +142,62 @@ def combine_spectrograms(specs, mp):
                 gp = g
                 spec_c[:, b, :] *= g
 
+    if torch.is_tensor(spec_c):
+        return spec_c.contiguous()
     return np.asfortranarray(spec_c)
 
 
-def spectrogram_to_image(spec, mode="magnitude"):
-    if mode == "magnitude":
-        if np.iscomplexobj(spec):
-            y = np.abs(spec)
-        else:
-            y = spec
-        y = np.log10(y**2 + 1e-8)
-    elif mode == "phase":
-        if np.iscomplexobj(spec):
-            y = np.angle(spec)
-        else:
-            y = spec
-
-    y -= y.min()
-    y *= 255 / y.max()
-    img = np.uint8(y)
-
-    if y.ndim == 3:
-        img = img.transpose(1, 2, 0)
-        img = np.concatenate([np.max(img, axis=2, keepdims=True), img], axis=2)
-
-    return img
-
-
-def reduce_vocal_aggressively(X, y, softmask):
-    v = X - y
-    y_mag_tmp = np.abs(y)
-    v_mag_tmp = np.abs(v)
-
-    v_mask = v_mag_tmp > y_mag_tmp
-    y_mag = np.clip(y_mag_tmp - v_mag_tmp * v_mask * softmask, 0, np.inf)
-
-    return y_mag * np.exp(1.0j * np.angle(y))
-
-
 def mask_silence(mag, ref, thres=0.2, min_range=64, fade_size=32):
     if min_range < fade_size * 2:
         raise ValueError("min_range must be >= fade_area * 2")
 
+    if torch.is_tensor(mag):
+        mag = mag.clone()
+        idx = torch.where(ref.mean(dim=(0, 1)) < thres)[0]
+        if idx.numel() == 0:
+            return mag
+        breaks = torch.where(torch.diff(idx) != 1)[0]
+        starts = torch.cat((idx[:1], idx[breaks + 1]))
+        ends = torch.cat((idx[breaks], idx[-1:]))
+        informative = torch.where(ends - starts > min_range)[0]
+        old_e = None
+        for position in informative.tolist():
+            s = int(starts[position].item())
+            e = int(ends[position].item())
+            if old_e is not None and s - old_e < fade_size:
+                s = old_e - fade_size * 2
+            if s != 0:
+                weight = torch.linspace(
+                    0,
+                    1,
+                    fade_size,
+                    device=mag.device,
+                    dtype=mag.dtype,
+                )
+                mag[:, :, s : s + fade_size] += (
+                    weight * ref[:, :, s : s + fade_size]
+                )
+            else:
+                s -= fade_size
+            if e != mag.shape[2]:
+                weight = torch.linspace(
+                    1,
+                    0,
+                    fade_size,
+                    device=mag.device,
+                    dtype=mag.dtype,
+                )
+                mag[:, :, e - fade_size : e] += (
+                    weight * ref[:, :, e - fade_size : e]
+                )
+            else:
+                e += fade_size
+            mag[:, :, s + fade_size : e - fade_size] += ref[
+                :, :, s + fade_size : e - fade_size
+            ]
+            old_e = e
+        return mag
+
     mag = mag.copy()
 
     idx = np.where(ref.mean(axis=(0, 1)) < thres)[0]
@@ -183,96 +230,34 @@ def mask_silence(mag, ref, thres=0.2, min_range=64, fade_size=32):
     return mag
 
 
-def align_wave_head_and_tail(a, b):
-    l = min([a[0].size, b[0].size])
-
-    return a[:l, :l], b[:l, :l]
-
-
-def cache_or_load(mix_path, inst_path, mp):
-    mix_basename = os.path.splitext(os.path.basename(mix_path))[0]
-    inst_basename = os.path.splitext(os.path.basename(inst_path))[0]
-
-    cache_dir = "mph{}".format(hashlib.sha1(json.dumps(mp.param, sort_keys=True).encode("utf-8")).hexdigest())
-    mix_cache_dir = os.path.join("cache", cache_dir)
-    inst_cache_dir = os.path.join("cache", cache_dir)
-
-    os.makedirs(mix_cache_dir, exist_ok=True)
-    os.makedirs(inst_cache_dir, exist_ok=True)
-
-    mix_cache_path = os.path.join(mix_cache_dir, mix_basename + ".npy")
-    inst_cache_path = os.path.join(inst_cache_dir, inst_basename + ".npy")
-
-    if os.path.exists(mix_cache_path) and os.path.exists(inst_cache_path):
-        X_spec_m = np.load(mix_cache_path)
-        y_spec_m = np.load(inst_cache_path)
-    else:
-        X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}
-
-        for d in range(len(mp.param["band"]), 0, -1):
-            bp = mp.param["band"][d]
-
-            if d == len(mp.param["band"]):  # high-end band
-                X_wave[d], _ = librosa.load(
-                    mix_path, sr=bp["sr"], mono=False, dtype=np.float32, res_type=bp["res_type"]
-                )
-                y_wave[d], _ = librosa.load(
-                    inst_path,
-                    sr=bp["sr"],
-                    mono=False,
-                    dtype=np.float32,
-                    res_type=bp["res_type"],
-                )
-            else:  # lower bands
-                X_wave[d] = librosa.resample(
-                    X_wave[d + 1],
-                    orig_sr=mp.param["band"][d + 1]["sr"],
-                    target_sr=bp["sr"],
-                    res_type=bp["res_type"],
-                )
-                y_wave[d] = librosa.resample(
-                    y_wave[d + 1],
-                    orig_sr=mp.param["band"][d + 1]["sr"],
-                    target_sr=bp["sr"],
-                    res_type=bp["res_type"],
-                )
-
-            X_wave[d], y_wave[d] = align_wave_head_and_tail(X_wave[d], y_wave[d])
-
-            X_spec_s[d] = wave_to_spectrogram(
-                X_wave[d],
-                bp["hl"],
-                bp["n_fft"],
-                mp.param["mid_side"],
-                mp.param["mid_side_b2"],
-                mp.param["reverse"],
-            )
-            y_spec_s[d] = wave_to_spectrogram(
-                y_wave[d],
-                bp["hl"],
-                bp["n_fft"],
-                mp.param["mid_side"],
-                mp.param["mid_side_b2"],
-                mp.param["reverse"],
-            )
-
-        del X_wave, y_wave
-
-        X_spec_m = combine_spectrograms(X_spec_s, mp)
-        y_spec_m = combine_spectrograms(y_spec_s, mp)
-
-        if X_spec_m.shape != y_spec_m.shape:
-            raise ValueError("The combined spectrograms are different: " + mix_path)
-
-        _, ext = os.path.splitext(mix_path)
-
-        np.save(mix_cache_path, X_spec_m)
-        np.save(inst_cache_path, y_spec_m)
-
-    return X_spec_m, y_spec_m
-
-
 def spectrogram_to_wave(spec, hop_length, mid_side, mid_side_b2, reverse):
+    if torch.is_tensor(spec):
+        n_fft = (spec.shape[1] - 1) * 2
+        wave = torch.istft(
+            spec.to(dtype=torch.complex64),
+            n_fft=n_fft,
+            hop_length=hop_length,
+            window=_stft_window(n_fft, spec.device),
+            center=True,
+            normalized=False,
+            onesided=True,
+            return_complex=False,
+        )
+        wave_left, wave_right = wave[0], wave[1]
+        if reverse:
+            return torch.stack(
+                (torch.flip(wave_left, dims=(-1,)), torch.flip(wave_right, dims=(-1,)))
+            )
+        if mid_side:
+            return torch.stack(
+                (wave_left + wave_right / 2, wave_left - wave_right / 2)
+            )
+        if mid_side_b2:
+            return torch.stack(
+                (wave_right / 1.25 + 0.4 * wave_left, wave_left / 1.25 - 0.4 * wave_right)
+            )
+        return wave
+
     spec_left = np.asfortranarray(spec[0])
     spec_right = np.asfortranarray(spec[1])
 
@@ -294,36 +279,6 @@ def spectrogram_to_wave(spec, hop_length, mid_side, mid_side_b2, reverse):
         return np.asfortranarray([wave_left, wave_right])
 
 
-def spectrogram_to_wave_mt(spec, hop_length, mid_side, reverse, mid_side_b2):
-    import threading
-
-    spec_left = np.asfortranarray(spec[0])
-    spec_right = np.asfortranarray(spec[1])
-
-    def run_thread(**kwargs):
-        global wave_left
-        wave_left = librosa.istft(**kwargs)
-
-    thread = threading.Thread(target=run_thread, kwargs={"stft_matrix": spec_left, "hop_length": hop_length})
-    thread.start()
-    wave_right = librosa.istft(spec_right, hop_length=hop_length)
-    thread.join()
-
-    if reverse:
-        return np.asfortranarray([np.flip(wave_left), np.flip(wave_right)])
-    elif mid_side:
-        return np.asfortranarray([np.add(wave_left, wave_right / 2), np.subtract(wave_left, wave_right / 2)])
-    elif mid_side_b2:
-        return np.asfortranarray(
-            [
-                np.add(wave_right / 1.25, 0.4 * wave_left),
-                np.subtract(wave_left / 1.25, 0.4 * wave_right),
-            ]
-        )
-    else:
-        return np.asfortranarray([wave_left, wave_right])
-
-
 def cmb_spectrogram_to_wave(spec_m, mp, extra_bins_h=None, extra_bins=None):
     wave_band = {}
     bands_n = len(mp.param["band"])
@@ -331,7 +286,11 @@ def cmb_spectrogram_to_wave(spec_m, mp, extra_bins_h=None, extra_bins=None):
 
     for d in range(1, bands_n + 1):
         bp = mp.param["band"][d]
-        spec_s = np.ndarray(shape=(2, bp["n_fft"] // 2 + 1, spec_m.shape[2]), dtype=complex)
+        shape = (2, bp["n_fft"] // 2 + 1, spec_m.shape[2])
+        if torch.is_tensor(spec_m):
+            spec_s = torch.zeros(shape, dtype=spec_m.dtype, device=spec_m.device)
+        else:
+            spec_s = np.ndarray(shape=shape, dtype=complex)
         h = bp["crop_stop"] - bp["crop_start"]
         spec_s[:, bp["crop_start"] : bp["crop_stop"], :] = spec_m[:, offset : offset + h, :]
 
@@ -351,49 +310,60 @@ def cmb_spectrogram_to_wave(spec_m, mp, extra_bins_h=None, extra_bins=None):
                     mp.param["reverse"],
                 )
             else:
-                wave = np.add(
-                    wave,
-                    spectrogram_to_wave(
+                wave = wave + spectrogram_to_wave(
                         spec_s,
                         bp["hl"],
                         mp.param["mid_side"],
                         mp.param["mid_side_b2"],
                         mp.param["reverse"],
-                    ),
                     )
         else:
             sr = mp.param["band"][d + 1]["sr"]
             if d == 1:  # lower
                 spec_s = fft_lp_filter(spec_s, bp["lpf_start"], bp["lpf_stop"])
-                wave = librosa.resample(
-                    spectrogram_to_wave(
+                band_wave = spectrogram_to_wave(
                         spec_s,
                         bp["hl"],
                         mp.param["mid_side"],
                         mp.param["mid_side_b2"],
                         mp.param["reverse"],
-                    ),
-                    orig_sr=bp["sr"],
-                    target_sr=sr,
+                    )
+                if torch.is_tensor(band_wave):
+                    wave = resample_audio_tensor(
+                        band_wave, bp["sr"], sr, force_mono=False
+                    )
+                else:
+                    wave = resample_audio(
+                        band_wave,
+                        bp["sr"],
+                        sr,
+                        force_mono=False,
                         res_type="sinc_fastest",
                     )
             else:  # mid
                 spec_s = fft_hp_filter(spec_s, bp["hpf_start"], bp["hpf_stop"] - 1)
                 spec_s = fft_lp_filter(spec_s, bp["lpf_start"], bp["lpf_stop"])
-                wave2 = np.add(
-                    wave,
-                    spectrogram_to_wave(
+                wave2 = wave + spectrogram_to_wave(
                         spec_s,
                         bp["hl"],
                         mp.param["mid_side"],
                         mp.param["mid_side_b2"],
                         mp.param["reverse"],
-                    ),
                     )
-                # wave = librosa.core.resample(wave2, orig_sr=bp['sr'], target_sr=sr, res_type="sinc_fastest")
-                wave = librosa.core.resample(wave2, orig_sr=bp["sr"], target_sr=sr, res_type="scipy")
+                if torch.is_tensor(wave2):
+                    wave = resample_audio_tensor(
+                        wave2, bp["sr"], sr, force_mono=False
+                    )
+                else:
+                    wave = resample_audio(
+                        wave2,
+                        bp["sr"],
+                        sr,
+                        force_mono=False,
+                        res_type="scipy",
+                    )
 
-    return wave.T
+    return wave.transpose(0, 1) if torch.is_tensor(wave) else wave.T
 
 
 def fft_lp_filter(spec, bin_start, bin_stop):
@@ -419,6 +389,31 @@ def fft_hp_filter(spec, bin_start, bin_stop):
 
 
 def mirroring(a, spec_m, input_high_end, mp):
+    if torch.is_tensor(spec_m):
+        source = spec_m[
+            :,
+            mp.param["pre_filter_start"]
+            - 10
+            - input_high_end.shape[1] : mp.param["pre_filter_start"]
+            - 10,
+            :,
+        ]
+        mirror = torch.flip(torch.abs(source), dims=(1,))
+        if "mirroring" == a:
+            mirror = torch.polar(mirror, torch.angle(input_high_end))
+            return torch.where(
+                torch.abs(input_high_end) <= torch.abs(mirror),
+                input_high_end,
+                mirror,
+            )
+        if "mirroring2" == a:
+            mirror = mirror * input_high_end * 1.7
+            return torch.where(
+                torch.abs(input_high_end) <= torch.abs(mirror),
+                input_high_end,
+                mirror,
+            )
+
     if "mirroring" == a:
         mirror = np.flip(
             np.abs(
@@ -448,190 +443,3 @@ def mirroring(a, spec_m, input_high_end, mp):
         mi = np.multiply(mirror, input_high_end * 1.7)
 
         return np.where(np.abs(input_high_end) <= np.abs(mi), input_high_end, mi)
-
-
-def ensembling(a, specs):
-    for i in range(1, len(specs)):
-        if i == 1:
-            spec = specs[0]
-
-        ln = min([spec.shape[2], specs[i].shape[2]])
-        spec = spec[:, :, :ln]
-        specs[i] = specs[i][:, :, :ln]
-
-        if "min_mag" == a:
-            spec = np.where(np.abs(specs[i]) <= np.abs(spec), specs[i], spec)
-        if "max_mag" == a:
-            spec = np.where(np.abs(specs[i]) >= np.abs(spec), specs[i], spec)
-
-    return spec
-
-
-def stft(wave, nfft, hl):
-    wave_left = np.asfortranarray(wave[0])
-    wave_right = np.asfortranarray(wave[1])
-    spec_left = librosa.stft(wave_left, n_fft=nfft, hop_length=hl)
-    spec_right = librosa.stft(wave_right, n_fft=nfft, hop_length=hl)
-    spec = np.asfortranarray([spec_left, spec_right])
-
-    return spec
-
-
-def istft(spec, hl):
-    spec_left = np.asfortranarray(spec[0])
-    spec_right = np.asfortranarray(spec[1])
-
-    wave_left = librosa.istft(spec_left, hop_length=hl)
-    wave_right = librosa.istft(spec_right, hop_length=hl)
-    wave = np.asfortranarray([wave_left, wave_right])
-
-    return wave
-
-
-if __name__ == "__main__":
-    import argparse
-    import time
-
-    import cv2
-    from model_param_init import ModelParameters
-
-    p = argparse.ArgumentParser()
-    p.add_argument(
-        "--algorithm",
-        "-a",
-        type=str,
-        choices=["invert", "invert_p", "min_mag", "max_mag", "deep", "align"],
-        default="min_mag",
-    )
-    p.add_argument(
-        "--model_params",
-        "-m",
-        type=str,
-        default=os.path.join("modelparams", "1band_sr44100_hl512.json"),
-    )
-    p.add_argument("--output_name", "-o", type=str, default="output")
-    p.add_argument("--vocals_only", "-v", action="store_true")
-    p.add_argument("input", nargs="+")
-    args = p.parse_args()
-
-    start_time = time.time()
-
-    if args.algorithm.startswith("invert") and len(args.input) != 2:
-        raise ValueError("There should be two input files.")
-
-    if not args.algorithm.startswith("invert") and len(args.input) < 2:
-        raise ValueError("There must be at least two input files.")
-
-    wave, specs = {}, {}
-    mp = ModelParameters(args.model_params)
-
-    for i in range(len(args.input)):
-        spec = {}
-
-        for d in range(len(mp.param["band"]), 0, -1):
-            bp = mp.param["band"][d]
-
-            if d == len(mp.param["band"]):  # high-end band
-                wave[d], _ = librosa.load(
-                    args.input[i],
-                    sr=bp["sr"],
-                    mono=False,
-                    dtype=np.float32,
-                    res_type=bp["res_type"],
-                )
-
-                if len(wave[d].shape) == 1:  # mono to stereo
-                    wave[d] = np.array([wave[d], wave[d]])
-            else:  # lower bands
-                wave[d] = librosa.resample(
-                    wave[d + 1],
-                    orig_sr=mp.param["band"][d + 1]["sr"],
-                    target_sr=bp["sr"],
-                    res_type=bp["res_type"],
-                )
-
-            spec[d] = wave_to_spectrogram(
-                wave[d],
-                bp["hl"],
-                bp["n_fft"],
-                mp.param["mid_side"],
-                mp.param["mid_side_b2"],
-                mp.param["reverse"],
-            )
-
-        specs[i] = combine_spectrograms(spec, mp)
-
-    del wave
-
-    if args.algorithm == "deep":
-        d_spec = np.where(np.abs(specs[0]) <= np.abs(spec[1]), specs[0], spec[1])
-        v_spec = d_spec - specs[1]
-        sf.write(
-            os.path.join("{}.wav".format(args.output_name)),
-            cmb_spectrogram_to_wave(v_spec, mp),
-            mp.param["sr"],
-        )
-
-    if args.algorithm.startswith("invert"):
-        ln = min([specs[0].shape[2], specs[1].shape[2]])
-        specs[0] = specs[0][:, :, :ln]
-        specs[1] = specs[1][:, :, :ln]
-
-        if "invert_p" == args.algorithm:
-            X_mag = np.abs(specs[0])
-            y_mag = np.abs(specs[1])
-            max_mag = np.where(X_mag >= y_mag, X_mag, y_mag)
-            v_spec = specs[1] - max_mag * np.exp(1.0j * np.angle(specs[0]))
-        else:
-            specs[1] = reduce_vocal_aggressively(specs[0], specs[1], 0.2)
-            v_spec = specs[0] - specs[1]
-
-            if not args.vocals_only:
-                X_mag = np.abs(specs[0])
-                y_mag = np.abs(specs[1])
-                v_mag = np.abs(v_spec)
-
-                X_image = spectrogram_to_image(X_mag)
-                y_image = spectrogram_to_image(y_mag)
-                v_image = spectrogram_to_image(v_mag)
-
-                cv2.imwrite("{}_X.png".format(args.output_name), X_image)
-                cv2.imwrite("{}_y.png".format(args.output_name), y_image)
-                cv2.imwrite("{}_v.png".format(args.output_name), v_image)
-
-                sf.write(
-                    "{}_X.wav".format(args.output_name),
-                    cmb_spectrogram_to_wave(specs[0], mp),
-                    mp.param["sr"],
-                )
-                sf.write(
-                    "{}_y.wav".format(args.output_name),
-                    cmb_spectrogram_to_wave(specs[1], mp),
-                    mp.param["sr"],
-                )
-
-        sf.write(
-            "{}_v.wav".format(args.output_name),
-            cmb_spectrogram_to_wave(v_spec, mp),
-            mp.param["sr"],
-        )
-    else:
-        if not args.algorithm == "deep":
-            sf.write(
-                os.path.join("ensembled", "{}.wav".format(args.output_name)),
-                cmb_spectrogram_to_wave(ensembling(args.algorithm, specs), mp),
-                mp.param["sr"],
-            )
-
-    if args.algorithm == "align":
-        trackalignment = [
-            {
-                "file1": '"{}"'.format(args.input[0]),
-                "file2": '"{}"'.format(args.input[1]),
-            }
-        ]
-
-        for i, e in tqdm(enumerate(trackalignment), desc="Performing Alignment..."):
-            os.system(f"python lib/align_tracks.py {e['file1']} {e['file2']}")
-
-    # print('Total time: {0:.{1}f}s'.format(time.time() - start_time, 1))

tools/uvr5/lib/utils.py

@@ -1,6 +1,7 @@
 import numpy as np
 import torch
-from tools.cuda_graph import run_cuda_graph
+import torch.nn.functional as F
+from tools.cuda_graph import clear_cuda_graph_cache, run_cuda_graph
 from tqdm import tqdm
 
 
@@ -14,11 +15,120 @@ def make_padding(width, cropsize, offset):
     return left, right, roi_size
 
 
+def _execute_torch_windows(
+    X_mag_pad,
+    roi_size,
+    n_window,
+    device,
+    model,
+    aggressiveness,
+    data,
+    batch_size,
+):
+    windows = X_mag_pad.unfold(
+        2,
+        data["window_size"],
+        roi_size,
+    )[:, :, :n_window, :]
+    model_dtype = next(model.parameters()).dtype
+    predictions = None
+    write_offset = 0
+    with torch.inference_mode():
+        for start in tqdm(range(0, n_window, batch_size)):
+            end = min(start + batch_size, n_window)
+            batch = (
+                windows[:, :, start:end, :]
+                .permute(2, 0, 1, 3)
+                .contiguous()
+                .to(device=device, dtype=model_dtype)
+            )
+            prediction = run_cuda_graph(
+                model,
+                "uvr-vr-%s" % repr(aggressiveness),
+                lambda window: model.predict(window, aggressiveness),
+                batch,
+            )
+            prediction = prediction.float().permute(1, 2, 0, 3).reshape(
+                prediction.shape[1], prediction.shape[2], -1
+            )
+            if predictions is None:
+                predictions = torch.empty(
+                    prediction.shape[0],
+                    prediction.shape[1],
+                    n_window * roi_size,
+                    device=prediction.device,
+                    dtype=torch.float32,
+                )
+            end_offset = write_offset + prediction.shape[2]
+            predictions[:, :, write_offset:end_offset].copy_(prediction)
+            write_offset = end_offset
+    return predictions[:, :, :write_offset]
+
+
+def _torch_batch_size(device):
+    free_bytes, _ = torch.cuda.mem_get_info(device)
+    free_gb = free_bytes / (1024**3)
+    if free_gb > 20:
+        return 8
+    if free_gb > 12:
+        return 4
+    if free_gb > 8:
+        return 2
+    return 1
+
+
+def _inference_torch(X_spec, device, model, aggressiveness, data):
+    X_spec = X_spec.to(device)
+    X_mag = torch.abs(X_spec)
+    coef = X_mag.max().clamp_min(1e-8)
+    X_mag_pre = X_mag / coef
+    n_frame = X_mag_pre.shape[2]
+    pad_l, pad_r, roi_size = make_padding(
+        n_frame, data["window_size"], model.offset
+    )
+    n_window = int(np.ceil(n_frame / roi_size))
+
+    def execute(pad_left, pad_right, windows_count):
+        padded = F.pad(X_mag_pre, (pad_left, pad_right))
+        batch_size = _torch_batch_size(device)
+        while True:
+            try:
+                return _execute_torch_windows(
+                    padded,
+                    roi_size,
+                    windows_count,
+                    device,
+                    model,
+                    aggressiveness,
+                    data,
+                    batch_size,
+                )
+            except torch.cuda.OutOfMemoryError:
+                clear_cuda_graph_cache(model)
+                torch.cuda.empty_cache()
+                if batch_size == 1:
+                    raise
+                batch_size = max(1, batch_size // 2)
+
+    pred = execute(pad_l, pad_r, n_window)[:, :, :n_frame]
+    if data["tta"]:
+        pad_l += roi_size // 2
+        pad_r += roi_size // 2
+        n_window += 1
+        pred_tta = execute(pad_l, pad_r, n_window)
+        pred_tta = pred_tta[:, :, roi_size // 2 :][:, :, :n_frame]
+        pred = (pred + pred_tta) * 0.5
+    return pred * coef, X_mag, None
+
+
 def inference(X_spec, device, model, aggressiveness, data):
     """
     data : dic configs
     """
 
+    if torch.is_tensor(X_spec) and X_spec.device.type == "cuda":
+        return _inference_torch(X_spec, device, model, aggressiveness, data)
+
     def _execute(X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half=True):
         model.eval()
         with torch.no_grad():

tools/uvr5/mdxnet.py

@@ -4,11 +4,11 @@ import sysconfig
 
 logger = logging.getLogger(__name__)
 
-import librosa
 import numpy as np
 import soundfile as sf
 import torch
 from tqdm import tqdm
+from infer.audio import load_audio, load_audio_tensor
 
 
 _ORT_CUDA_DLL_HANDLES = []
@@ -61,9 +61,7 @@ cpu = torch.device("cpu")
 
 
 class ConvTDFNetTrim:
-    def __init__(self, device, model_name, target_name, L, dim_f, dim_t, n_fft, hop=1024):
-        super(ConvTDFNetTrim, self).__init__()
-
+    def __init__(self, device, dim_f, dim_t, n_fft, hop=1024):
         self.dim_f = dim_f
         self.dim_t = 2**dim_t
         self.n_fft = n_fft
@@ -71,14 +69,11 @@ class ConvTDFNetTrim:
         self.n_bins = self.n_fft // 2 + 1
         self.chunk_size = hop * (self.dim_t - 1)
         self.window = torch.hann_window(window_length=self.n_fft, periodic=True).to(device)
-        self.target_name = target_name
-        self.blender = "blender" in model_name
-
         self.dim_c = 4
-        out_c = self.dim_c * 4 if target_name == "*" else self.dim_c
-        self.freq_pad = torch.zeros([1, out_c, self.n_bins - self.dim_f, self.dim_t]).to(device)
-
-        self.n = L // 2
+        self.freq_pad = torch.zeros(
+            [1, self.dim_c, self.n_bins - self.dim_f, self.dim_t],
+            device=device,
+        )
 
     def stft(self, x):
         x = x.reshape([-1, self.chunk_size])
@@ -95,10 +90,10 @@ class ConvTDFNetTrim:
         x = x.reshape([-1, 2, 2, self.n_bins, self.dim_t]).reshape([-1, self.dim_c, self.n_bins, self.dim_t])
         return x[:, :, : self.dim_f]
 
-    def istft(self, x, freq_pad=None):
-        freq_pad = self.freq_pad.repeat([x.shape[0], 1, 1, 1]) if freq_pad is None else freq_pad
+    def istft(self, x):
+        freq_pad = self.freq_pad.expand(x.shape[0], -1, -1, -1)
         x = torch.cat([x, freq_pad], -2)
-        c = 4 * 2 if self.target_name == "*" else 2
+        c = 2
         x = x.reshape([-1, c, 2, self.n_bins, self.dim_t]).reshape([-1, 2, self.n_bins, self.dim_t])
         x = x.permute([0, 2, 3, 1])
         x = x.contiguous()
@@ -110,9 +105,6 @@ class ConvTDFNetTrim:
 def get_models(device, dim_f, dim_t, n_fft):
     return ConvTDFNetTrim(
         device=device,
-        model_name="Conv-TDF",
-        target_name="vocals",
-        L=11,
         dim_f=dim_f,
         dim_t=dim_t,
         n_fft=n_fft,
@@ -124,9 +116,13 @@ class Predictor:
         import onnxruntime as ort
 
         available_providers = ort.get_available_providers()
+        requested_providers = [
+            provider[0] if isinstance(provider, (tuple, list)) else provider
+            for provider in args.providers
+        ]
         logger.info("ONNX Runtime available providers: %s", available_providers)
         if (
-            "CUDAExecutionProvider" in args.providers
+            "CUDAExecutionProvider" in requested_providers
             and "CUDAExecutionProvider" not in available_providers
         ):
             raise RuntimeError(
@@ -136,7 +132,7 @@ class Predictor:
                 "project's runtime Python."
             )
         if (
-            "DmlExecutionProvider" in args.providers
+            "DmlExecutionProvider" in requested_providers
             and "DmlExecutionProvider" not in available_providers
         ):
             raise RuntimeError(
@@ -145,15 +141,30 @@ class Predictor:
                 "requirments_cpu_py312.txt with this project's runtime Python."
             )
         self.args = args
-        self.model_ = get_models(device=cpu, dim_f=args.dim_f, dim_t=args.dim_t, n_fft=args.n_fft)
+        try:
+            requested_torch_device = torch.device(args.device)
+        except Exception:
+            requested_torch_device = cpu
+        if requested_torch_device.type == "cuda" and requested_torch_device.index is None:
+            requested_torch_device = torch.device("cuda:0")
+        # DirectML and CPU keep the established NumPy/CPU STFT path.  The
+        # Torch CUDA path is enabled only after the ORT session confirms that
+        # its CUDA provider really became the primary provider.
+        model_device = requested_torch_device if requested_torch_device.type == "cuda" else cpu
+        self.model_ = get_models(
+            device=model_device,
+            dim_f=args.dim_f,
+            dim_t=args.dim_t,
+            n_fft=args.n_fft,
+        )
         self.model = ort.InferenceSession(
-            os.path.join(args.onnx, self.model_.target_name + ".onnx"),
+            os.path.join(args.onnx, "vocals.onnx"),
             providers=args.providers,
         )
         active_providers = self.model.get_providers()
         logger.info("ONNX Runtime active providers: %s", active_providers)
         if (
-            "CUDAExecutionProvider" in args.providers
+            "CUDAExecutionProvider" in requested_providers
             and (
                 not active_providers
                 or active_providers[0] != "CUDAExecutionProvider"
@@ -164,7 +175,7 @@ class Predictor:
                 "check the CUDA 11/cuDNN 8 DLL installation."
             )
         if (
-            "DmlExecutionProvider" in args.providers
+            "DmlExecutionProvider" in requested_providers
             and (
                 not active_providers
                 or active_providers[0] != "DmlExecutionProvider"
@@ -174,7 +185,67 @@ class Predictor:
                 "The FoxJoy ONNX model did not activate DmlExecutionProvider; "
                 "check the ONNX Runtime DirectML installation."
             )
-        logger.info("ONNX load done")
+        self.cuda_pipeline = bool(
+            requested_torch_device.type == "cuda"
+            and active_providers
+            and active_providers[0] == "CUDAExecutionProvider"
+        )
+        self.torch_device = requested_torch_device if self.cuda_pipeline else cpu
+        logger.info(
+            "ONNX load done; FoxJoy tensor pipeline=%s, torch device=%s",
+            "cuda" if self.cuda_pipeline else "cpu-compatible",
+            self.torch_device,
+        )
+
+    def _run_ort_cuda(self, input_tensor, output_tensor):
+        input_tensor = input_tensor.contiguous()
+        if input_tensor.dtype != torch.float32:
+            input_tensor = input_tensor.float()
+        if not output_tensor.is_contiguous() or output_tensor.dtype != torch.float32:
+            raise RuntimeError("FoxJoy CUDA output buffer must be contiguous float32")
+
+        device_id = self.torch_device.index
+        io_binding = self.model.io_binding()
+        io_binding.bind_input(
+            name=self.model.get_inputs()[0].name,
+            device_type="cuda",
+            device_id=device_id,
+            element_type=np.float32,
+            shape=tuple(input_tensor.shape),
+            buffer_ptr=input_tensor.data_ptr(),
+        )
+        io_binding.bind_output(
+            name=self.model.get_outputs()[0].name,
+            device_type="cuda",
+            device_id=device_id,
+            element_type=np.float32,
+            shape=tuple(output_tensor.shape),
+            buffer_ptr=output_tensor.data_ptr(),
+        )
+        # ORT owns a separate CUDA stream by default.  Explicit boundaries
+        # guarantee that it sees the completed Torch STFT and that Torch sees
+        # the completed output without staging either tensor through NumPy.
+        torch.cuda.synchronize(self.torch_device)
+        self.model.run_with_iobinding(io_binding)
+        torch.cuda.synchronize(self.torch_device)
+        return input_tensor
+
+    def _infer_cuda(self, spek):
+        spek = spek.contiguous().float()
+        output = torch.empty_like(spek)
+        if self.args.denoise:
+            # Reuse both the ORT output allocation and the input allocation
+            # for the negative/positive passes.  Only the accumulator is
+            # separate because the second ORT run overwrites its output.
+            spek.neg_()
+            spek = self._run_ort_cuda(spek, output)
+            prediction = output * -0.5
+            spek.neg_()
+            spek = self._run_ort_cuda(spek, output)
+            prediction.add_(output, alpha=0.5)
+            return prediction
+        self._run_ort_cuda(spek, output)
+        return output
 
     def demix(self, mix):
         samples = mix.shape[-1]
@@ -198,7 +269,10 @@ class Predictor:
 
             start = skip - s_margin
 
-            segmented_mix[skip] = mix[:, start:end].copy()
+            segment = mix[:, start:end]
+            # CUDA segments are views of the already resident decoded audio;
+            # copying every segment would almost double long-file VRAM use.
+            segmented_mix[skip] = segment if torch.is_tensor(segment) else segment.copy()
             if end == samples:
                 break
 
@@ -222,31 +296,78 @@ class Predictor:
             trim = model.n_fft // 2
             gen_size = model.chunk_size - 2 * trim
             pad = gen_size - n_sample % gen_size
-            mix_p = np.concatenate((np.zeros((2, trim)), cmix, np.zeros((2, pad)), np.zeros((2, trim))), 1)
+            if self.cuda_pipeline and torch.is_tensor(cmix):
+                cmix = cmix.to(self.torch_device, dtype=torch.float32)
+                mix_p = torch.cat(
+                    (
+                        cmix.new_zeros((2, trim)),
+                        cmix,
+                        cmix.new_zeros((2, pad)),
+                        cmix.new_zeros((2, trim)),
+                    ),
+                    1,
+                )
+            else:
+                mix_p = np.concatenate(
+                    (
+                        np.zeros((2, trim)),
+                        cmix,
+                        np.zeros((2, pad)),
+                        np.zeros((2, trim)),
+                    ),
+                    1,
+                )
             mix_waves = []
             i = 0
             while i < n_sample + pad:
-                waves = np.array(mix_p[:, i : i + model.chunk_size])
+                waves = mix_p[:, i : i + model.chunk_size]
+                if not torch.is_tensor(waves):
+                    waves = np.array(waves)
                 mix_waves.append(waves)
                 i += gen_size
-            mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(cpu)
+            if torch.is_tensor(mix_waves[0]):
+                mix_waves = torch.stack(mix_waves).float()
+            else:
+                mix_waves = torch.from_numpy(np.asarray(mix_waves, dtype=np.float32))
             with torch.no_grad():
                 _ort = self.model
+                if self.cuda_pipeline:
+                    # One H2D for all windows in this outer segment. STFT,
+                    # both denoise passes and ISTFT remain on the selected
+                    # CUDA device; only the finished waveform returns to CPU.
+                    if mix_waves.device != self.torch_device:
+                        mix_waves = mix_waves.to(self.torch_device, non_blocking=True)
+                    spek = model.stft(mix_waves)
+                    spec_pred = self._infer_cuda(spek)
+                    tar_waves = model.istft(spec_pred)
+                    tar_signal = (
+                        tar_waves[:, :, trim:-trim]
+                        .transpose(0, 1)
+                        .reshape(2, -1)[:, :-pad]
+                        .cpu()
+                        .numpy()
+                    )
+                else:
                     spek = model.stft(mix_waves)
                     if self.args.denoise:
+                        spek_numpy = spek.numpy()
                         spec_pred = (
-                        -_ort.run(None, {"input": -spek.cpu().numpy()})[0] * 0.5
-                        + _ort.run(None, {"input": spek.cpu().numpy()})[0] * 0.5
+                            -_ort.run(None, {"input": -spek_numpy})[0] * 0.5
+                            + _ort.run(None, {"input": spek_numpy})[0] * 0.5
                         )
-                    tar_waves = model.istft(torch.tensor(spec_pred))
+                        tar_waves = model.istft(torch.from_numpy(spec_pred))
                     else:
-                    tar_waves = model.istft(torch.tensor(_ort.run(None, {"input": spek.cpu().numpy()})[0]))
-                tar_signal = tar_waves[:, :, trim:-trim].transpose(0, 1).reshape(2, -1).numpy()[:, :-pad]
+                        spec_pred = _ort.run(None, {"input": spek.numpy()})[0]
+                        tar_waves = model.istft(torch.from_numpy(spec_pred))
+                    tar_signal = (
+                        tar_waves[:, :, trim:-trim]
+                        .transpose(0, 1)
+                        .reshape(2, -1)
+                        .numpy()[:, :-pad]
+                    )
 
                 start = 0 if mix == 0 else margin_size
                 end = None if mix == list(mixes.keys())[::-1][0] else -margin_size
-                if margin_size == 0:
-                    end = None
                 sources.append(tar_signal[:, start:end])
 
                 progress_bar.update(1)
@@ -261,12 +382,24 @@ class Predictor:
         os.makedirs(vocal_root, exist_ok=True)
         os.makedirs(others_root, exist_ok=True)
         basename = os.path.basename(m)
-        mix, rate = librosa.load(m, mono=False, sr=44100)
+        if self.cuda_pipeline:
+            mix = load_audio_tensor(m, 44100, force_mono=False)
+            mix = mix.to(self.torch_device)
+        else:
+            mix = load_audio(m, 44100, force_mono=False)
+        rate = 44100
         if mix.ndim == 1:
-            mix = np.asfortranarray([mix, mix])
-        mix = mix.T
-        sources = self.demix(mix.T)
+            mix = mix.unsqueeze(0) if torch.is_tensor(mix) else mix[np.newaxis, :]
+        if mix.shape[0] == 1:
+            mix = mix.repeat(2, 1) if torch.is_tensor(mix) else np.repeat(mix, 2, axis=0)
+        elif mix.shape[0] > 2:
+            mix = mix[:2].contiguous() if torch.is_tensor(mix) else np.ascontiguousarray(mix[:2])
+        sources = self.demix(mix)
         opt = sources[0].T
+        if torch.is_tensor(mix):
+            mix = mix.transpose(0, 1).float().cpu().numpy()
+        else:
+            mix = mix.T
         if format in ["wav", "flac"]:
             sf.write("%s/%s_main_vocal.%s" % (vocal_root, basename, format), mix - opt, rate)
             sf.write("%s/%s_others.%s" % (others_root, basename, format), opt, rate)
@@ -294,22 +427,20 @@ class Predictor:
 
 
 class MDXNetDereverb:
-    def __init__(self, chunks, providers):
+    def __init__(self, chunks, providers, device="cpu"):
         self.onnx = os.path.join(
             os.getenv("weight_uvr5_root", "assets/uvr5_weights"),
             "onnx_dereverb_By_FoxJoy",
         )
-        self.shifts = 10  # 'Predict with randomised equivariant stabilisation'
-        self.mixing = "min_mag"  # ['default','min_mag','max_mag']
         self.chunks = chunks
         self.providers = providers
+        self.device = device
         self.margin = 44100
         self.dim_t = 9
         self.dim_f = 3072
         self.n_fft = 6144
         self.denoise = True
         self.pred = Predictor(self)
-        self.device = cpu
 
     def _path_audio_(self, input, others_root, vocal_root, format, is_hp3=False):
         self.pred.prediction(input, vocal_root, others_root, format)

tools/uvr5/vr.py

@@ -5,10 +5,16 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-import librosa
 import numpy as np
 import soundfile as sf
 import torch
+from infer.audio import (
+    TORCHAUDIO_GPU_ENABLED,
+    load_audio,
+    load_audio_tensor,
+    resample_audio,
+    resample_audio_tensor,
+)
 from tools.uvr5.lib.lib_v5 import nets_61968KB as Nets
 from tools.uvr5.lib.lib_v5 import spec_utils
 from tools.uvr5.lib.lib_v5.model_param_init import ModelParameters
@@ -16,6 +22,152 @@ from tools.uvr5.lib.lib_v5.nets_new import CascadedNet
 from tools.uvr5.lib.utils import inference
 
 
+def _ensure_stereo(audio):
+    audio = np.asarray(audio, dtype=np.float32)
+    if audio.ndim == 1:
+        audio = audio[np.newaxis, :]
+    if audio.shape[0] == 1:
+        return np.repeat(audio, 2, axis=0)
+    if audio.shape[0] > 2:
+        return np.ascontiguousarray(audio[:2])
+    return audio
+
+
+def _ensure_stereo_tensor(audio, device):
+    if audio.ndim == 1:
+        audio = audio.unsqueeze(0)
+    if audio.shape[0] == 1:
+        audio = audio.repeat(2, 1)
+    elif audio.shape[0] > 2:
+        audio = audio[:2]
+    return audio.to(device=device)
+
+
+def _cuda_device(device):
+    parsed = device if isinstance(device, torch.device) else torch.device(device)
+    return parsed if parsed.type == "cuda" else None
+
+
+def _vr_gpu_memory_fits(audio, mp, device):
+    highest_band = len(mp.param["band"])
+    frames = max(
+        1,
+        int(audio.shape[-1] // mp.param["band"][highest_band]["hl"] + 1),
+    )
+    band_bins = sum(
+        mp.param["band"][band]["n_fft"] // 2 + 1
+        for band in mp.param["band"]
+    )
+    combined_bins = mp.param["bins"] + 1
+    # Complex band spectra + combined/target spectra + magnitude/prediction.
+    estimated = frames * 2 * (
+        band_bins * 8 + combined_bins * (8 * 3 + 4 * 3)
+    )
+    free_bytes, _ = torch.cuda.mem_get_info(device)
+    return estimated <= int(free_bytes * 0.42)
+
+
+def _prepare_spectrogram(music_file, mp, data, device, allow_gpu=True):
+    cuda_device = _cuda_device(device)
+    use_gpu = bool(
+        allow_gpu and cuda_device is not None and TORCHAUDIO_GPU_ENABLED
+    )
+    if use_gpu:
+        try:
+            high_sr = mp.param["band"][len(mp.param["band"])]["sr"]
+            high_wave = _ensure_stereo_tensor(
+                load_audio_tensor(music_file, high_sr, force_mono=False),
+                cuda_device,
+            )
+            if not _vr_gpu_memory_fits(high_wave, mp, cuda_device):
+                use_gpu = False
+                high_wave = high_wave.float().cpu().numpy()
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            use_gpu = False
+            high_wave = None
+    else:
+        high_wave = None
+
+    input_high_end_h = None
+    input_high_end = None
+    X_spec_s = {}
+    bands_n = len(mp.param["band"])
+    previous_wave = None
+    for d in range(bands_n, 0, -1):
+        bp = mp.param["band"][d]
+        if d == bands_n:
+            if high_wave is None:
+                current_wave = _ensure_stereo(
+                    load_audio(music_file, bp["sr"], force_mono=False)
+                )
+            else:
+                current_wave = high_wave
+        elif use_gpu:
+            current_wave = resample_audio_tensor(
+                previous_wave,
+                mp.param["band"][d + 1]["sr"],
+                bp["sr"],
+                force_mono=False,
+            )
+        else:
+            current_wave = resample_audio(
+                previous_wave,
+                mp.param["band"][d + 1]["sr"],
+                bp["sr"],
+                force_mono=False,
+                res_type=bp["res_type"],
+            )
+        X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(
+            current_wave,
+            bp["hl"],
+            bp["n_fft"],
+            mp.param["mid_side"],
+            mp.param["mid_side_b2"],
+            mp.param["reverse"],
+        )
+        if d == bands_n and data["high_end_process"] != "none":
+            input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (
+                mp.param["pre_filter_stop"] - mp.param["pre_filter_start"]
+            )
+            input_high_end = X_spec_s[d][
+                :, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, :
+            ]
+            if torch.is_tensor(input_high_end):
+                input_high_end = input_high_end.clone()
+        previous_wave = current_wave
+
+    X_spec_m = spec_utils.combine_spectrograms(X_spec_s, mp)
+    del previous_wave, X_spec_s
+    return X_spec_m, input_high_end_h, input_high_end
+
+
+def _wave_for_write(wave):
+    if torch.is_tensor(wave):
+        return wave.detach().to(device="cpu", dtype=torch.float32).numpy()
+    return np.asarray(wave)
+
+
+def _separate_spectrogram(X_spec_m, device, model, aggressiveness, data):
+    with torch.no_grad():
+        pred, X_mag, X_phase = inference(
+            X_spec_m, device, model, aggressiveness, data
+        )
+    if data["postprocess"]:
+        if torch.is_tensor(pred):
+            pred_inv = torch.clamp(X_mag - pred, min=0)
+        else:
+            pred_inv = np.clip(X_mag - pred, 0, np.inf)
+        pred = spec_utils.mask_silence(pred, pred_inv)
+    if torch.is_tensor(X_spec_m):
+        ratio = pred.float() / X_mag.clamp_min(1e-8)
+        ratio = torch.nan_to_num(ratio)
+        y_spec_m = X_spec_m * ratio
+    else:
+        y_spec_m = pred * X_phase
+    return y_spec_m
+
+
 class AudioPre:
     def __init__(self, agg, model_path, device, is_half, tta=False):
         self.model_path = model_path
@@ -50,61 +202,32 @@ class AudioPre:
             os.makedirs(ins_root, exist_ok=True)
         if vocal_root is not None:
             os.makedirs(vocal_root, exist_ok=True)
-        X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}
-        bands_n = len(self.mp.param["band"])
-        # print(bands_n)
-        for d in range(bands_n, 0, -1):
-            bp = self.mp.param["band"][d]
-            if d == bands_n:  # high-end band
-                (
-                    X_wave[d],
-                    _,
-                ) = librosa.core.load(  # 理论上librosa读取可能对某些音频有bug，应该上ffmpeg读取，但是太麻烦了弃坑
-                    music_file,
-                    sr=bp["sr"],
-                    mono=False,
-                    dtype=np.float32,
-                    res_type=bp["res_type"],
-                )
-                if X_wave[d].ndim == 1:
-                    X_wave[d] = np.asfortranarray([X_wave[d], X_wave[d]])
-            else:  # lower bands
-                X_wave[d] = librosa.core.resample(
-                    X_wave[d + 1],
-                    orig_sr=self.mp.param["band"][d + 1]["sr"],
-                    target_sr=bp["sr"],
-                    res_type=bp["res_type"],
-                )
-            # Stft of wave source
-            X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(
-                X_wave[d],
-                bp["hl"],
-                bp["n_fft"],
-                self.mp.param["mid_side"],
-                self.mp.param["mid_side_b2"],
-                self.mp.param["reverse"],
-            )
-            # pdb.set_trace()
-            if d == bands_n and self.data["high_end_process"] != "none":
-                input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (
-                    self.mp.param["pre_filter_stop"] - self.mp.param["pre_filter_start"]
-                )
-                input_high_end = X_spec_s[d][:, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, :]
-
-        X_spec_m = spec_utils.combine_spectrograms(X_spec_s, self.mp)
         aggresive_set = float(self.data["agg"] / 100)
         aggressiveness = {
             "value": aggresive_set,
             "split_bin": self.mp.param["band"][1]["crop_stop"],
         }
-        with torch.no_grad():
-            pred, X_mag, X_phase = inference(X_spec_m, self.device, self.model, aggressiveness, self.data)
-        # Postprocess
-        if self.data["postprocess"]:
-            pred_inv = np.clip(X_mag - pred, 0, np.inf)
-            pred = spec_utils.mask_silence(pred, pred_inv)
-        y_spec_m = pred * X_phase
-        v_spec_m = X_spec_m - y_spec_m
+        gpu_oom = False
+        try:
+            X_spec_m, input_high_end_h, input_high_end = _prepare_spectrogram(
+                music_file, self.mp, self.data, self.device
+            )
+            y_spec_m = _separate_spectrogram(
+                X_spec_m, self.device, self.model, aggressiveness, self.data
+            )
+        except torch.cuda.OutOfMemoryError:
+            X_spec_m = None
+            input_high_end = None
+            y_spec_m = None
+            gpu_oom = True
+        if gpu_oom:
+            torch.cuda.empty_cache()
+            X_spec_m, input_high_end_h, input_high_end = _prepare_spectrogram(
+                music_file, self.mp, self.data, self.device, allow_gpu=False
+            )
+            y_spec_m = _separate_spectrogram(
+                X_spec_m, self.device, self.model, aggressiveness, self.data
+            )
 
         if is_hp3 == True:
             ins_root, vocal_root = vocal_root, ins_root
@@ -128,14 +251,14 @@ class AudioPre:
                         ins_root,
                         head + "{}_{}.{}".format(name, self.data["agg"], format),
                     ),
-                    (np.array(wav_instrument) * 32768).astype("int16"),
+                    (_wave_for_write(wav_instrument) * 32768).astype("int16"),
                     self.mp.param["sr"],
                 )  #
             else:
                 path = os.path.join(ins_root, head + "{}_{}.wav".format(name, self.data["agg"]))
                 sf.write(
                     path,
-                    (np.array(wav_instrument) * 32768).astype("int16"),
+                    (_wave_for_write(wav_instrument) * 32768).astype("int16"),
                     self.mp.param["sr"],
                 )
                 if os.path.exists(path):
@@ -149,6 +272,12 @@ class AudioPre:
                         except:
                             pass
         if vocal_root is not None:
+            if torch.is_tensor(y_spec_m):
+                y_spec_m.neg_().add_(X_spec_m)
+                v_spec_m = y_spec_m
+            else:
+                np.subtract(X_spec_m, y_spec_m, out=y_spec_m)
+                v_spec_m = y_spec_m
             if is_hp3 == True:
                 head = "instrument_"
             else:
@@ -165,14 +294,14 @@ class AudioPre:
                         vocal_root,
                         head + "{}_{}.{}".format(name, self.data["agg"], format),
                     ),
-                    (np.array(wav_vocals) * 32768).astype("int16"),
+                    (_wave_for_write(wav_vocals) * 32768).astype("int16"),
                     self.mp.param["sr"],
                 )
             else:
                 path = os.path.join(vocal_root, head + "{}_{}.wav".format(name, self.data["agg"]))
                 sf.write(
                     path,
-                    (np.array(wav_vocals) * 32768).astype("int16"),
+                    (_wave_for_write(wav_vocals) * 32768).astype("int16"),
                     self.mp.param["sr"],
                 )
                 if os.path.exists(path):
@@ -224,61 +353,32 @@ class AudioPreDeEcho:
             os.makedirs(ins_root, exist_ok=True)
         if vocal_root is not None:
             os.makedirs(vocal_root, exist_ok=True)
-        X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}
-        bands_n = len(self.mp.param["band"])
-        # print(bands_n)
-        for d in range(bands_n, 0, -1):
-            bp = self.mp.param["band"][d]
-            if d == bands_n:  # high-end band
-                (
-                    X_wave[d],
-                    _,
-                ) = librosa.core.load(  # 理论上librosa读取可能对某些音频有bug，应该上ffmpeg读取，但是太麻烦了弃坑
-                    music_file,
-                    sr=bp["sr"],
-                    mono=False,
-                    dtype=np.float32,
-                    res_type=bp["res_type"],
-                )
-                if X_wave[d].ndim == 1:
-                    X_wave[d] = np.asfortranarray([X_wave[d], X_wave[d]])
-            else:  # lower bands
-                X_wave[d] = librosa.core.resample(
-                    X_wave[d + 1],
-                    orig_sr=self.mp.param["band"][d + 1]["sr"],
-                    target_sr=bp["sr"],
-                    res_type=bp["res_type"],
-                )
-            # Stft of wave source
-            X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(
-                X_wave[d],
-                bp["hl"],
-                bp["n_fft"],
-                self.mp.param["mid_side"],
-                self.mp.param["mid_side_b2"],
-                self.mp.param["reverse"],
-            )
-            # pdb.set_trace()
-            if d == bands_n and self.data["high_end_process"] != "none":
-                input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (
-                    self.mp.param["pre_filter_stop"] - self.mp.param["pre_filter_start"]
-                )
-                input_high_end = X_spec_s[d][:, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, :]
-
-        X_spec_m = spec_utils.combine_spectrograms(X_spec_s, self.mp)
         aggresive_set = float(self.data["agg"] / 100)
         aggressiveness = {
             "value": aggresive_set,
             "split_bin": self.mp.param["band"][1]["crop_stop"],
         }
-        with torch.no_grad():
-            pred, X_mag, X_phase = inference(X_spec_m, self.device, self.model, aggressiveness, self.data)
-        # Postprocess
-        if self.data["postprocess"]:
-            pred_inv = np.clip(X_mag - pred, 0, np.inf)
-            pred = spec_utils.mask_silence(pred, pred_inv)
-        y_spec_m = pred * X_phase
-        v_spec_m = X_spec_m - y_spec_m
+        gpu_oom = False
+        try:
+            X_spec_m, input_high_end_h, input_high_end = _prepare_spectrogram(
+                music_file, self.mp, self.data, self.device
+            )
+            y_spec_m = _separate_spectrogram(
+                X_spec_m, self.device, self.model, aggressiveness, self.data
+            )
+        except torch.cuda.OutOfMemoryError:
+            X_spec_m = None
+            input_high_end = None
+            y_spec_m = None
+            gpu_oom = True
+        if gpu_oom:
+            torch.cuda.empty_cache()
+            X_spec_m, input_high_end_h, input_high_end = _prepare_spectrogram(
+                music_file, self.mp, self.data, self.device, allow_gpu=False
+            )
+            y_spec_m = _separate_spectrogram(
+                X_spec_m, self.device, self.model, aggressiveness, self.data
+            )
 
         if ins_root is not None:
             if self.data["high_end_process"].startswith("mirroring"):
@@ -295,14 +395,14 @@ class AudioPreDeEcho:
                         ins_root,
                         "vocal_{}_{}.{}".format(name, self.data["agg"], format),
                     ),
-                    (np.array(wav_instrument) * 32768).astype("int16"),
+                    (_wave_for_write(wav_instrument) * 32768).astype("int16"),
                     self.mp.param["sr"],
                 )  #
             else:
                 path = os.path.join(ins_root, "vocal_{}_{}.wav".format(name, self.data["agg"]))
                 sf.write(
                     path,
-                    (np.array(wav_instrument) * 32768).astype("int16"),
+                    (_wave_for_write(wav_instrument) * 32768).astype("int16"),
                     self.mp.param["sr"],
                 )
                 if os.path.exists(path):
@@ -316,6 +416,12 @@ class AudioPreDeEcho:
                         except:
                             pass
         if vocal_root is not None:
+            if torch.is_tensor(y_spec_m):
+                y_spec_m.neg_().add_(X_spec_m)
+                v_spec_m = y_spec_m
+            else:
+                np.subtract(X_spec_m, y_spec_m, out=y_spec_m)
+                v_spec_m = y_spec_m
             if self.data["high_end_process"].startswith("mirroring"):
                 input_high_end_ = spec_utils.mirroring(self.data["high_end_process"], v_spec_m, input_high_end, self.mp)
                 wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp, input_high_end_h, input_high_end_)
@@ -328,14 +434,14 @@ class AudioPreDeEcho:
                         vocal_root,
                         "instrument_{}_{}.{}".format(name, self.data["agg"], format),
                     ),
-                    (np.array(wav_vocals) * 32768).astype("int16"),
+                    (_wave_for_write(wav_vocals) * 32768).astype("int16"),
                     self.mp.param["sr"],
                 )
             else:
                 path = os.path.join(vocal_root, "instrument_{}_{}.wav".format(name, self.data["agg"]))
                 sf.write(
                     path,
-                    (np.array(wav_vocals) * 32768).astype("int16"),
+                    (_wave_for_write(wav_vocals) * 32768).astype("int16"),
                     self.mp.param["sr"],
                 )
                 if os.path.exists(path):

tools/uvr5/webui.py

@@ -2,10 +2,9 @@ import logging
 import os
 import traceback
 
-import ffmpeg
 import torch
 
-from configs.config import Config, IS_GPU
+from configs.config import Config
 from tools.uvr5.bsroformer import Roformer_Loader
 from tools.uvr5.mdxnet import MDXNetDereverb
 from tools.uvr5.vr import AudioPre, AudioPreDeEcho
@@ -35,11 +34,16 @@ def uvr(model_name, inp_root, save_root_vocal, paths, save_root_ins, agg, format
         if model_name == "onnx_dereverb_By_FoxJoy":
             if config.dml:
                 providers = ["DmlExecutionProvider", "CPUExecutionProvider"]
-            elif IS_GPU:
-                providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+            elif torch.device(config.device).type == "cuda":
+                cuda_device = torch.device(config.device)
+                device_id = cuda_device.index if cuda_device.index is not None else 0
+                providers = [
+                    ("CUDAExecutionProvider", {"device_id": str(device_id)}),
+                    "CPUExecutionProvider",
+                ]
             else:
                 providers = ["CPUExecutionProvider"]
-            pre_fun = MDXNetDereverb(15, providers)
+            pre_fun = MDXNetDereverb(15, providers, config.device)
         elif "roformer" in model_name.lower():
             pre_fun = Roformer_Loader(
                 model_path=os.path.join(weight_uvr5_root, model_name + ".ckpt"),
@@ -68,37 +72,11 @@ def uvr(model_name, inp_root, save_root_vocal, paths, save_root_ins, agg, format
             inp_path = os.path.join(inp_root, path)
             if not os.path.isfile(inp_path):
                 continue
-            need_reformat = True
-            done = False
             try:
-                info = ffmpeg.probe(inp_path, cmd="ffprobe")
-                if (
-                    info["streams"][0]["channels"] == 2
-                    and info["streams"][0]["sample_rate"] == "44100"
-                ):
-                    need_reformat = False
-                    pre_fun._path_audio_(
-                        inp_path,
-                        save_root_ins,
-                        save_root_vocal,
-                        format0,
-                        is_hp3,
-                    )
-                    done = True
-            except:
-                traceback.print_exc()
-            if need_reformat:
-                tmp_path = "%s/%s.reformatted.wav" % (
-                    os.environ["TEMP"],
-                    os.path.basename(inp_path),
-                )
-                os.system(
-                    'ffmpeg -i "%s" -vn -acodec pcm_s16le -ac 2 -ar 44100 "%s" -y'
-                    % (inp_path, tmp_path)
-                )
-                inp_path = tmp_path
-            try:
-                if not done:
+                # Let each model loader decode the original file.  Its
+                # torchaudio path can then perform any required 44.1 kHz
+                # conversion on the selected CUDA device instead of hiding it
+                # behind a CPU FFmpeg pre-conversion.
                 pre_fun._path_audio_(
                     inp_path,
                     save_root_ins,

train/dataset/extract_f0.py

@@ -8,7 +8,6 @@ import logging
 
 import numpy as np
 
-from infer.audio import load_audio
 from i18n.i18n import I18nAuto
 from tools.progress import should_report
 
@@ -44,6 +43,9 @@ elif mode in ("dml", "directml"):
 else:
     raise ValueError("Unsupported F0 extraction mode: %s" % mode)
 
+# CUDA_VISIBLE_DEVICES must be set before infer.audio imports torch/configs.
+from infer.audio import load_audio
+
 f = open("%s/extract_f0_feature.log" % exp_dir, "a", encoding="utf8")
 
 

train/preprocess.py

@@ -10,13 +10,13 @@ n_p = int(sys.argv[3])
 exp_dir = sys.argv[4]
 noparallel = sys.argv[5] == "True"
 per = float(sys.argv[6])
-import os
 import traceback
 
 import librosa
 import numpy as np
 from scipy.io import wavfile
 
+os.environ["RVC_AUDIO_FORCE_CPU"] = "1"
 from infer.audio import load_audio
 from train.dataset.slicer2 import Slicer
 from i18n.i18n import I18nAuto
@@ -73,13 +73,13 @@ class PreProcess:
             self.sr,
             tmp_audio.astype(np.float32),
         )
-        tmp_audio = librosa.resample(
+        audio_16k = librosa.resample(
             tmp_audio, orig_sr=self.sr, target_sr=16000
-        )  # , res_type="soxr_vhq"
+        ).astype(np.float32)
         wavfile.write(
             "%s/%s_%s.wav" % (self.wavs16k_dir, idx0, idx1),
             16000,
-            tmp_audio.astype(np.float32),
+            audio_16k,
         )
         return True
 
@@ -140,19 +140,24 @@ class PreProcess:
                 ("%s/%s" % (inp_root, name), idx, total)
                 for idx, name in enumerate(names)
             ]
-            println(i18n("[数据切分] 待处理：%s | 进程数：%s") % (total, n_p))
+            worker_count = max(n_p, 1)
+            worker_count = min(worker_count, max(total, 1))
+            println(
+                i18n("[数据切分] 待处理：%s | 进程数：%s")
+                % (total, worker_count)
+            )
             if noparallel:
-                for i in range(n_p):
-                    self.pipeline_mp(infos[i::n_p])
+                for i in range(worker_count):
+                    self.pipeline_mp(infos[i::worker_count])
             else:
                 ps = []
-                for i in range(n_p):
+                for i in range(worker_count):
                     p = multiprocessing.Process(
-                        target=self.pipeline_mp, args=(infos[i::n_p],)
+                        target=self.pipeline_mp, args=(infos[i::worker_count],)
                     )
                     ps.append(p)
                     p.start()
-                for i in range(n_p):
+                for i in range(worker_count):
                     ps[i].join()
         except Exception:
             println(i18n("[数据切分][失败] %s") % traceback.format_exc())

webui.py

@@ -487,10 +487,28 @@ def run_preprocess_dataset(trainset_dir, exp_dir, sr, n_p, state, format_output=
         config.noparallel,
         config.preprocess_per,
     )
+    extract_start_time = time.time()
+    requested_workers = max(int(n_p), 1)
+    actual_workers = 1 if config.noparallel else requested_workers
+    print(
+        i18n(
+            "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s"
+        )
+        % (extract_start_time, requested_workers, actual_workers),
+        flush=True,
+    )
+    try:
         process = start_train_process(state, cmd)
         yield from wait_train_processes(
             state, [process], log_path, "数据切分", format_output
         )
+    finally:
+        extract_end_time = time.time()
+        print(
+            i18n("数据提取结束：end_time=%.6f，总耗时=%.3f秒")
+            % (extract_end_time, extract_end_time - extract_start_time),
+            flush=True,
+        )
     if not train_task_stopped(state):
         validate_preprocess_outputs(exp_dir)