diff --git a/README.md b/README.md index 8c80b91..c6456c3 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ + 使用少量数据进行训练也能得到较好结果(推荐至少收集10分钟低底噪语音数据) + 可以通过模型融合来改变音色(借助ckpt处理选项卡中的ckpt-merge) + 简单易用的网页界面 -+ 可调用UVR5模型来快速分离人声和伴奏 ++ 可调用pymss/MSST模型来快速分离人声和伴奏 + 使用最先进的[人声音高提取算法InterSpeech2023-RMVPE](#参考项目)根绝哑音问题，速度快、资源占用小 + A卡/I卡使用 CPU 依赖方案；Windows 可使用 DirectML，Linux 使用 CPU @@ -142,7 +142,7 @@ assets/ ├── rmvpe/rmvpe.pt ├── pretrained/ ├── pretrained_v2/ -├── uvr5_weights/ +├── pymss_weights/ ├── weights/ # user RVC .pth models └── indices/ # user .index files logs/ @@ -155,7 +155,7 @@ assets/hubert_base/pytorch_model.bin assets/rmvpe/rmvpe.pt assets/pretrained/*.pth assets/pretrained_v2/*.pth -assets/uvr5_weights/* +assets/pymss_weights/* assets/weights/*.pth assets/indices/*.index logs/mute/* @@ -179,9 +179,9 @@ hf download lj1995/VoiceConversionWebUI mute.zip --revision main \ --local-dir .model-downloads python -m zipfile -e .model-downloads/mute.zip logs -# Required only for UVR5 vocal separation +# Required only for pymss/MSST vocal separation hf download lj1995/VoiceConversionWebUI --revision main \ - --include "uvr5_weights/*" --local-dir assets + --include "pymss_weights/*" --local-dir assets ``` 仅 Windows AMD/Intel DirectML 环境还需要： @@ -222,6 +222,7 @@ python webui.py --noautoopen + [Gradio](https://github.com/gradio-app/gradio) + [FFmpeg](https://github.com/FFmpeg/FFmpeg) + [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui) ++ [pymss-project/pymss](https://github.com/pymss-project/pymss) + [audio-slicer](https://github.com/openvpi/audio-slicer) + [Vocal pitch extraction:RMVPE](https://github.com/Dream-High/RMVPE) + The pretrained model is trained and tested by [yxlllc](https://github.com/yxlllc/RMVPE) and [RVC-Boss](https://github.com/RVC-Boss). diff --git a/assets/pymss_weights/config_mel_band_roformer_karaoke.yaml b/assets/pymss_weights/config_mel_band_roformer_karaoke.yaml new file mode 100644 index 0000000..9052dc0 --- /dev/null +++ b/assets/pymss_weights/config_mel_band_roformer_karaoke.yaml @@ -0,0 +1,71 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 000 + +model: + dim: 384 + depth: 6 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0 + ff_dropout: 0 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + +training: + batch_size: 4 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - karaoke + - other + lr: 1.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: karaoke + num_epochs: 1000 + num_steps: 2000 + augmentation: false # enable augmentations by audiomentations and pedalboard + augmentation_type: null + use_mp3_compress: false # Deprecated + augmentation_mix: false # Mix several stems of the same type with some probability + augmentation_loudness: false # randomly change loudness of each stem + augmentation_loudness_type: 1 # Type 1 or 2 + augmentation_loudness_min: 0 + augmentation_loudness_max: 0 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/assets/pymss_weights/dereverb_mel_band_roformer_anvuew.yaml b/assets/pymss_weights/dereverb_mel_band_roformer_anvuew.yaml new file mode 100644 index 0000000..5a96f43 --- /dev/null +++ b/assets/pymss_weights/dereverb_mel_band_roformer_anvuew.yaml @@ -0,0 +1,76 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 6 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0 + ff_dropout: 0 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + +training: + batch_size: 3 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - noreverb + - reverb + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: noreverb + num_epochs: 1000 + num_steps: 4000 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adamw + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.1 + loudness_max: 1.0 + mixup: false # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 801 + num_overlap: 2 \ No newline at end of file diff --git a/assets/pymss_weights/model_bs_roformer_ep_317_sdr_12.9755.yaml b/assets/pymss_weights/model_bs_roformer_ep_317_sdr_12.9755.yaml new file mode 100644 index 0000000..7f4ce5b --- /dev/null +++ b/assets/pymss_weights/model_bs_roformer_ep_317_sdr_12.9755.yaml @@ -0,0 +1,123 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 801 + hop_length: 441 + min_mean_abs: 0.0 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 +inference: + batch_size: 4 + dim_t: 801 + num_overlap: 2 +model: + attn_dropout: 0.1 + depth: 12 + dim: 512 + dim_freqs_in: 1025 + dim_head: 64 + ff_dropout: 0.1 + flash_attn: true + freq_transformer_depth: 1 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + heads: 8 + linear_transformer_depth: 0 + mask_estimator_depth: 2 + multi_stft_hop_size: 147 + multi_stft_normalized: false + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + num_stems: 1 + stereo: true + stft_hop_length: 441 + stft_n_fft: 2048 + stft_normalized: false + stft_win_length: 2048 + time_transformer_depth: 1 +training: + batch_size: 2 + coarse_loss_clip: true + ema_momentum: 0.999 + grad_clip: 0 + gradient_accumulation_steps: 1 + instruments: + - vocals + - other + lr: 1.0e-05 + num_epochs: 1000 + num_steps: 1000 + optimizer: adam + other_fix: true + patience: 2 + q: 0.95 + reduce_factor: 0.95 + target_instrument: vocals + use_amp: true diff --git a/assets/pymss_weights/model_bs_roformer_ep_368_sdr_12.9628.yaml b/assets/pymss_weights/model_bs_roformer_ep_368_sdr_12.9628.yaml new file mode 100644 index 0000000..d1a439a --- /dev/null +++ b/assets/pymss_weights/model_bs_roformer_ep_368_sdr_12.9628.yaml @@ -0,0 +1,133 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 801 # don't work (use in model) + hop_length: 441 # don't work (use in model) + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + dim: 512 + depth: 12 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + +training: + batch_size: 16 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - instrumental + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + augmentation: false # enable augmentations by audiomentations and pedalboard + augmentation_type: simple1 + use_mp3_compress: false # Deprecated + augmentation_mix: true # Mix several stems of the same type with some probability + augmentation_loudness: true # randomly change loudness of each stem + augmentation_loudness_type: 1 # Type 1 or 2 + augmentation_loudness_min: 0.5 + augmentation_loudness_max: 1.5 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +inference: + batch_size: 1 + dim_t: 901 + num_overlap: 4 \ No newline at end of file diff --git a/docs/en/README.en.md b/docs/en/README.en.md index 5b53c8c..c57bb56 100644 --- a/docs/en/README.en.md +++ b/docs/en/README.en.md @@ -51,7 +51,7 @@ A simple, easy-to-use voice timbre conversion / voice changer framework.

+ Training with a small amounts of data (>=10min low noise speech recommended); + Model fusion to change timbres (using ckpt processing tab->ckpt merge); + Easy-to-use WebUI; -+ UVR5 model to quickly separate vocals and instruments; ++ pymss/MSST model to quickly separate vocals and instruments; + High-pitch Voice Extraction Algorithm [InterSpeech2023-RMVPE](#Credits) to prevent a muted sound problem. Provides the best results (significantly) and is faster with lower resource consumption than Crepe_full; + AMD/Intel systems use the CPU dependency set; Windows may use DirectML and Linux uses CPU; @@ -143,7 +143,7 @@ assets/ ├── rmvpe/rmvpe.pt ├── pretrained/ ├── pretrained_v2/ -├── uvr5_weights/ +├── pymss_weights/ ├── weights/ # user RVC .pth models └── indices/ # user .index files logs/ @@ -156,7 +156,7 @@ assets/hubert_base/pytorch_model.bin assets/rmvpe/rmvpe.pt assets/pretrained/*.pth assets/pretrained_v2/*.pth -assets/uvr5_weights/* +assets/pymss_weights/* assets/weights/*.pth assets/indices/*.index logs/mute/* @@ -180,9 +180,9 @@ hf download lj1995/VoiceConversionWebUI mute.zip --revision main \ --local-dir .model-downloads python -m zipfile -e .model-downloads/mute.zip logs -# Required only for UVR5 vocal separation +# Required only for pymss/MSST vocal separation hf download lj1995/VoiceConversionWebUI --revision main \ - --include "uvr5_weights/*" --local-dir assets + --include "pymss_weights/*" --local-dir assets ``` Windows AMD/Intel DirectML environments additionally need: @@ -221,6 +221,7 @@ The default port is `7865`. Put personal `.pth` models in `assets/weights/` and + [Gradio](https://github.com/gradio-app/gradio) + [FFmpeg](https://github.com/FFmpeg/FFmpeg) + [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui) ++ [pymss-project/pymss](https://github.com/pymss-project/pymss) + [audio-slicer](https://github.com/openvpi/audio-slicer) + [Vocal pitch extraction:RMVPE](https://github.com/Dream-High/RMVPE) + The pretrained model is trained and tested by [yxlllc](https://github.com/yxlllc/RMVPE) and [RVC-Boss](https://github.com/RVC-Boss). diff --git a/docs/fr/README.fr.md b/docs/fr/README.fr.md index 590a249..d87ad04 100644 --- a/docs/fr/README.fr.md +++ b/docs/fr/README.fr.md @@ -35,7 +35,7 @@ Ce dépôt a les caractéristiques suivantes : + Obtient de bons résultats même avec peu de données pour la formation (il est recommandé de collecter au moins 10 minutes de données vocales avec un faible bruit de fond). + Peut changer le timbre vocal en fusionnant des modèles (avec l'aide de l'onglet ckpt-merge). + Interface web simple et facile à utiliser. -+ Peut appeler le modèle UVR5 pour séparer rapidement la voix et l'accompagnement. ++ Peut appeler le modèle pymss/MSST pour séparer rapidement la voix et l'accompagnement. + Utilise l'algorithme de pitch vocal le plus avancé [InterSpeech2023-RMVPE](#projets-référencés) pour éliminer les problèmes de voix muette. Meilleurs résultats, plus rapide que crepe_full, et moins gourmand en ressources. + Les systèmes AMD/Intel utilisent les dépendances CPU ; Windows peut utiliser DirectML et Linux utilise le CPU. @@ -127,7 +127,7 @@ assets/ ├── rmvpe/rmvpe.pt ├── pretrained/ ├── pretrained_v2/ -├── uvr5_weights/ +├── pymss_weights/ ├── weights/ # user RVC .pth models └── indices/ # user .index files logs/ @@ -140,7 +140,7 @@ assets/hubert_base/pytorch_model.bin assets/rmvpe/rmvpe.pt assets/pretrained/*.pth assets/pretrained_v2/*.pth -assets/uvr5_weights/* +assets/pymss_weights/* assets/weights/*.pth assets/indices/*.index logs/mute/* @@ -164,9 +164,9 @@ hf download lj1995/VoiceConversionWebUI mute.zip --revision main \ --local-dir .model-downloads python -m zipfile -e .model-downloads/mute.zip logs -# Required only for UVR5 vocal separation +# Required only for pymss/MSST vocal separation hf download lj1995/VoiceConversionWebUI --revision main \ - --include "uvr5_weights/*" --local-dir assets + --include "pymss_weights/*" --local-dir assets ``` Les environnements Windows AMD/Intel DirectML nécessitent aussi : @@ -205,6 +205,7 @@ Le port par défaut est `7865`. Placez les modèles `.pth` dans `assets/weights/ + [Gradio](https://github.com/gradio-app/gradio) + [FFmpeg](https://github.com/FFmpeg/FFmpeg) + [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui) ++ [pymss-project/pymss](https://github.com/pymss-project/pymss) + [audio-slicer](https://github.com/openvpi/audio-slicer) + [Extraction de la hauteur vocale : RMVPE](https://github.com/Dream-High/RMVPE) + Le modèle pré-entraîné a été formé et testé par [yxlllc](https://github.com/yxlllc/RMVPE) et [RVC-Boss](https://github.com/RVC-Boss). diff --git a/docs/jp/README.ja.md b/docs/jp/README.ja.md index 0512a2e..5c74019 100644 --- a/docs/jp/README.ja.md +++ b/docs/jp/README.ja.md @@ -53,7 +53,7 @@ - 少量のデータセットからでも、比較的良い結果を得ることができます。（10 分以上のノイズの少ない音声を推奨します。） - モデルを融合することで、音声を混ぜることができます。（ckpt processing タブの、ckpt merge を使用します。） - 使いやすい WebUI。 -- UVR5 Model も含んでいるため、人の声と BGM を素早く分離できます。 +- pymss/MSST Model も含んでいるため、人の声と BGM を素早く分離できます。 - 最先端の[人間の声のピッチ抽出アルゴリズム InterSpeech2023-RMVPE](#参照プロジェクト)を使用して無声音問題を解決します。効果は最高（著しく）で、crepe_full よりも速く、リソース使用が少ないです。 - A カードと I カードの加速サポート @@ -147,7 +147,7 @@ assets/ ├── rmvpe/rmvpe.pt ├── pretrained/ ├── pretrained_v2/ -├── uvr5_weights/ +├── pymss_weights/ ├── weights/ # user RVC .pth models └── indices/ # user .index files logs/ @@ -160,7 +160,7 @@ assets/hubert_base/pytorch_model.bin assets/rmvpe/rmvpe.pt assets/pretrained/*.pth assets/pretrained_v2/*.pth -assets/uvr5_weights/* +assets/pymss_weights/* assets/weights/*.pth assets/indices/*.index logs/mute/* @@ -184,9 +184,9 @@ hf download lj1995/VoiceConversionWebUI mute.zip --revision main \ --local-dir .model-downloads python -m zipfile -e .model-downloads/mute.zip logs -# Required only for UVR5 vocal separation +# Required only for pymss/MSST vocal separation hf download lj1995/VoiceConversionWebUI --revision main \ - --include "uvr5_weights/*" --local-dir assets + --include "pymss_weights/*" --local-dir assets ``` Windows の AMD/Intel DirectML 環境では、さらに次のファイルが必要です。 @@ -226,6 +226,7 @@ python webui.py --noautoopen - [Gradio](https://github.com/gradio-app/gradio) - [FFmpeg](https://github.com/FFmpeg/FFmpeg) - [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui) +- [pymss-project/pymss](https://github.com/pymss-project/pymss) - [audio-slicer](https://github.com/openvpi/audio-slicer) - [Vocal pitch extraction:RMVPE](https://github.com/Dream-High/RMVPE) - 事前訓練されたモデルは[yxlllc](https://github.com/yxlllc/RMVPE)と[RVC-Boss](https://github.com/RVC-Boss)によって訓練され、テストされました。 diff --git a/docs/kr/README.ko.han.md b/docs/kr/README.ko.han.md index 6bd9092..615ab3b 100644 --- a/docs/kr/README.ko.han.md +++ b/docs/kr/README.ko.han.md @@ -34,7 +34,7 @@ + 적은量의 데이터로 訓練해도 좋은 結果를 얻을 수 있음 (最小10分以上의 低雜음音聲데이터를 使用하는 것을 勸獎); + 모델融合을通한 音色의 變調可能 (ckpt處理탭->ckpt混合選擇); + 使用하기 쉬운 WebUI (웹 使用者인터페이스); -+ UVR5 모델을 利用하여 목소리와 背景音樂의 빠른 分離; ++ pymss/MSST 모델을 利用하여 목소리와 背景音樂의 빠른 分離; ## 環境의 準備 @@ -117,7 +117,7 @@ assets/ ├── rmvpe/rmvpe.pt ├── pretrained/ ├── pretrained_v2/ -├── uvr5_weights/ +├── pymss_weights/ ├── weights/ # user RVC .pth models └── indices/ # user .index files logs/ @@ -130,7 +130,7 @@ assets/hubert_base/pytorch_model.bin assets/rmvpe/rmvpe.pt assets/pretrained/*.pth assets/pretrained_v2/*.pth -assets/uvr5_weights/* +assets/pymss_weights/* assets/weights/*.pth assets/indices/*.index logs/mute/* @@ -154,9 +154,9 @@ hf download lj1995/VoiceConversionWebUI mute.zip --revision main \ --local-dir .model-downloads python -m zipfile -e .model-downloads/mute.zip logs -# Required only for UVR5 vocal separation +# Required only for pymss/MSST vocal separation hf download lj1995/VoiceConversionWebUI --revision main \ - --include "uvr5_weights/*" --local-dir assets + --include "pymss_weights/*" --local-dir assets ``` Windows AMD/Intel DirectML 환경에는 다음 파일도 필요합니다. @@ -195,6 +195,7 @@ python webui.py --noautoopen + [Gradio](https://github.com/gradio-app/gradio) + [FFmpeg](https://github.com/FFmpeg/FFmpeg) + [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui) ++ [pymss-project/pymss](https://github.com/pymss-project/pymss) + [audio-slicer](https://github.com/openvpi/audio-slicer) ## 모든寄與者분들의勞力에感謝드립니다 diff --git a/docs/kr/README.ko.md b/docs/kr/README.ko.md index b52de8b..21c25ad 100644 --- a/docs/kr/README.ko.md +++ b/docs/kr/README.ko.md @@ -53,7 +53,7 @@ - 적은 양의 데이터로 훈련해도 좋은 결과를 얻을 수 있음 (최소 10분 이상의 저잡음 음성 데이터를 사용하는 것을 권장) - 모델 융합을 통한 음색의 변조 가능 (ckpt 처리 탭->ckpt 병합 선택) - 사용하기 쉬운 WebUI (웹 인터페이스) -- UVR5 모델을 이용하여 목소리와 배경음악의 빠른 분리; +- pymss/MSST 모델을 이용하여 목소리와 배경음악의 빠른 분리; - 최첨단 [음성 피치 추출 알고리즘 InterSpeech2023-RMVPE](#参考项目)을 사용하여 무성음 문제를 해결합니다. 효과는 최고(압도적)이며 crepe_full보다 더 빠르고 리소스 사용이 적음 - A카드와 I카드 가속을 지원 @@ -147,7 +147,7 @@ assets/ ├── rmvpe/rmvpe.pt ├── pretrained/ ├── pretrained_v2/ -├── uvr5_weights/ +├── pymss_weights/ ├── weights/ # user RVC .pth models └── indices/ # user .index files logs/ @@ -160,7 +160,7 @@ assets/hubert_base/pytorch_model.bin assets/rmvpe/rmvpe.pt assets/pretrained/*.pth assets/pretrained_v2/*.pth -assets/uvr5_weights/* +assets/pymss_weights/* assets/weights/*.pth assets/indices/*.index logs/mute/* @@ -184,9 +184,9 @@ hf download lj1995/VoiceConversionWebUI mute.zip --revision main \ --local-dir .model-downloads python -m zipfile -e .model-downloads/mute.zip logs -# Required only for UVR5 vocal separation +# Required only for pymss/MSST vocal separation hf download lj1995/VoiceConversionWebUI --revision main \ - --include "uvr5_weights/*" --local-dir assets + --include "pymss_weights/*" --local-dir assets ``` Windows AMD/Intel DirectML 환경에는 다음 파일도 필요합니다. @@ -226,6 +226,7 @@ python webui.py --noautoopen - [Gradio](https://github.com/gradio-app/gradio) - [FFmpeg](https://github.com/FFmpeg/FFmpeg) - [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui) +- [pymss-project/pymss](https://github.com/pymss-project/pymss) - [audio-slicer](https://github.com/openvpi/audio-slicer) - [Vocal pitch extraction:RMVPE](https://github.com/Dream-High/RMVPE) - 사전 훈련된 모델은 [yxlllc](https://github.com/yxlllc/RMVPE)와 [RVC-Boss](https://github.com/RVC-Boss)에 의해 훈련되고 테스트되었습니다. diff --git a/docs/pt/README.pt.md b/docs/pt/README.pt.md index edb6c61..c73ea6b 100644 --- a/docs/pt/README.pt.md +++ b/docs/pt/README.pt.md @@ -42,7 +42,7 @@ Este repositório possui os seguintes recursos: + Treinar com uma pequena quantidade de dados também obtém resultados relativamente bons (>=10min de áudio com baixo ruído recomendado); + Suporta fusão de modelos para alterar timbres (usando guia de processamento ckpt-> mesclagem ckpt); + Interface Webui fácil de usar; -+ Use o modelo UVR5 para separar rapidamente vocais e instrumentos. ++ Use o modelo pymss/MSST para separar rapidamente vocais e instrumentos. + Use o mais poderoso algoritmo de extração de voz de alta frequência [InterSpeech2023-RMVPE](#Credits) para evitar o problema de som mudo. Fornece os melhores resultados (significativamente) e é mais rápido, com consumo de recursos ainda menor que o Crepe_full. + Sistemas AMD/Intel usam as dependências de CPU; Windows pode usar DirectML e Linux usa CPU. @@ -134,7 +134,7 @@ assets/ ├── rmvpe/rmvpe.pt ├── pretrained/ ├── pretrained_v2/ -├── uvr5_weights/ +├── pymss_weights/ ├── weights/ # user RVC .pth models └── indices/ # user .index files logs/ @@ -147,7 +147,7 @@ assets/hubert_base/pytorch_model.bin assets/rmvpe/rmvpe.pt assets/pretrained/*.pth assets/pretrained_v2/*.pth -assets/uvr5_weights/* +assets/pymss_weights/* assets/weights/*.pth assets/indices/*.index logs/mute/* @@ -171,9 +171,9 @@ hf download lj1995/VoiceConversionWebUI mute.zip --revision main \ --local-dir .model-downloads python -m zipfile -e .model-downloads/mute.zip logs -# Required only for UVR5 vocal separation +# Required only for pymss/MSST vocal separation hf download lj1995/VoiceConversionWebUI --revision main \ - --include "uvr5_weights/*" --local-dir assets + --include "pymss_weights/*" --local-dir assets ``` Ambientes Windows AMD/Intel DirectML também precisam de: @@ -212,6 +212,7 @@ A porta padrão é `7865`. Coloque modelos `.pth` em `assets/weights/` e arquivo + [Gradio](https://github.com/gradio-app/gradio) + [FFmpeg](https://github.com/FFmpeg/FFmpeg) + [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui) ++ [pymss-project/pymss](https://github.com/pymss-project/pymss) + [audio-slicer](https://github.com/openvpi/audio-slicer) + [Vocal pitch extraction:RMVPE](https://github.com/Dream-High/RMVPE) + The pretrained model is trained and tested by [yxlllc](https://github.com/yxlllc/RMVPE) and [RVC-Boss](https://github.com/RVC-Boss). diff --git a/docs/tr/README.tr.md b/docs/tr/README.tr.md index 7f35070..91b2fe8 100644 --- a/docs/tr/README.tr.md +++ b/docs/tr/README.tr.md @@ -37,7 +37,7 @@ Bu depo aşağıdaki özelliklere sahiptir: + Az miktarda veriyle bile nispeten iyi sonuçlar alın (>=10 dakika düşük gürültülü konuşma önerilir); + Timbraları değiştirmek için model birleştirmeyi destekleme (ckpt işleme sekmesi-> ckpt birleştir); + Kullanımı kolay Web arayüzü; -+ UVR5 modelini kullanarak hızla vokalleri ve enstrümanları ayırma. ++ pymss/MSST modelini kullanarak hızla vokalleri ve enstrümanları ayırma. + En güçlü Yüksek tiz Ses Çıkarma Algoritması [InterSpeech2023-RMVPE](#Krediler) sessiz ses sorununu önlemek için kullanılır. En iyi sonuçları (önemli ölçüde) sağlar ve Crepe_full'den daha hızlı çalışır, hatta daha düşük kaynak tüketimi sağlar. + AMD/Intel sistemleri CPU bağımlılıklarını kullanır; Windows DirectML, Linux CPU kullanabilir. @@ -129,7 +129,7 @@ assets/ ├── rmvpe/rmvpe.pt ├── pretrained/ ├── pretrained_v2/ -├── uvr5_weights/ +├── pymss_weights/ ├── weights/ # user RVC .pth models └── indices/ # user .index files logs/ @@ -142,7 +142,7 @@ assets/hubert_base/pytorch_model.bin assets/rmvpe/rmvpe.pt assets/pretrained/*.pth assets/pretrained_v2/*.pth -assets/uvr5_weights/* +assets/pymss_weights/* assets/weights/*.pth assets/indices/*.index logs/mute/* @@ -166,9 +166,9 @@ hf download lj1995/VoiceConversionWebUI mute.zip --revision main \ --local-dir .model-downloads python -m zipfile -e .model-downloads/mute.zip logs -# Required only for UVR5 vocal separation +# Required only for pymss/MSST vocal separation hf download lj1995/VoiceConversionWebUI --revision main \ - --include "uvr5_weights/*" --local-dir assets + --include "pymss_weights/*" --local-dir assets ``` Windows AMD/Intel DirectML ortamlarında ayrıca şu dosya gerekir: @@ -207,6 +207,7 @@ Varsayılan bağlantı noktası `7865`'tir. `.pth` modellerini `assets/weights/` + [Gradio](https://github.com/gradio-app/gradio) + [FFmpeg](https://github.com/FFmpeg/FFmpeg) + [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui) ++ [pymss-project/pymss](https://github.com/pymss-project/pymss) + [audio-slicer](https://github.com/openvpi/audio-slicer) + [Vokal ton çıkarma:RMVPE](https://github.com/Dream-High/RMVPE) + Ön eğitimli model [yxlllc](https://github.com/yxlllc/RMVPE) ve [RVC-Boss](https://github.com/RVC-Boss) tarafından eğitilip test edilmiştir. diff --git a/infer/audio.py b/infer/audio.py index 0879d64..552df36 100644 --- a/infer/audio.py +++ b/infer/audio.py @@ -44,25 +44,56 @@ AUDIO_DTYPE = _AUDIO_DTYPE def wav2(i, o, format): inp = av.open(i, "r") - if format == "m4a": - format = "mp4" - out = av.open(o, "w", format=format) - if format == "ogg": - format = "libvorbis" - if format == "mp4": - format = "aac" - - ostream = out.add_stream(format) - - for frame in inp.decode(audio=0): - for p in ostream.encode(frame): - out.mux(p) - - for p in ostream.encode(None): - out.mux(p) - - out.close() - inp.close() + try: + if format == "m4a": + format = "mp4" + out = av.open(o, "w", format=format) + try: + if format == "ogg": + format = "libvorbis" + if format == "mp4": + format = "aac" + + if not inp.streams.audio: + raise ValueError("Input contains no audio stream") + input_stream = inp.streams.audio[0] + source_rate = input_stream.codec_context.sample_rate + ostream = ( + out.add_stream(format, rate=source_rate) + if source_rate + else out.add_stream(format) + ) + source_channels = input_stream.codec_context.channels + if source_channels == 1: + ostream.layout = "mono" + elif source_channels == 2: + ostream.layout = "stereo" + + for frame in inp.decode(input_stream): + for p in ostream.encode(frame): + out.mux(p) + + for p in ostream.encode(None): + out.mux(p) + finally: + out.close() + finally: + inp.close() + + +def transcode_audio_file(input_path, output_path, format): + """Transcode a WAV path and remove partial compressed output on failure.""" + output_path = os.fspath(output_path) + if os.path.exists(output_path): + os.remove(output_path) + try: + wav2(input_path, output_path, format) + if not os.path.isfile(output_path) or os.path.getsize(output_path) == 0: + raise RuntimeError("Audio transcoding produced no output: %s" % output_path) + except Exception: + if os.path.exists(output_path): + os.remove(output_path) + raise def _probe_audio(file): diff --git a/requirments_cu118_py312.txt b/requirments_cu118_py312.txt index 785aa5c..47162bc 100644 --- a/requirments_cu118_py312.txt +++ b/requirments_cu118_py312.txt @@ -31,6 +31,11 @@ matplotlib>=3.8.2,<4 networkx>=3.2.0,<4 numpy>=1.26.4,<2 +# Five-model MSST inference backend. Install both packages from the index; +# local wheel paths are intentionally not used by this requirements file. +pymss==2.0.14 +pymss-core==0.1.4 + # ONNX Runtime 1.18.x is the CUDA 11 / cuDNN 8 generation with Python 3.12 # Windows wheels. The CUDA DLL packages are pinned because later cuDNN/CUDA # majors are ABI-incompatible with this provider build. @@ -43,7 +48,7 @@ nvidia-cufft-cu11==10.9.0.58 opencv-python-headless>=4.10.0,<5 praat-parselmouth>=0.4.5,<1 -PyYAML>=6.0 +PyYAML>=6.0.1 scikit-learn>=1.6.0,<2 scipy>=1.13.1,<2 sounddevice>=0.5.0,<1 diff --git a/requirments_cu128_py312.txt b/requirments_cu128_py312.txt index 2aa873f..fd619f3 100644 --- a/requirments_cu128_py312.txt +++ b/requirments_cu128_py312.txt @@ -31,16 +31,19 @@ matplotlib>=3.8.2,<4 networkx>=3.2.0,<4 numpy>=1.26.4,<2 +# Five-model MSST inference backend. Install both packages from the index; +# local wheel paths are intentionally not used by this requirements file. +pymss==2.0.14 +pymss-core==0.1.4 + # ONNX Runtime 1.19.x uses the CUDA 12 / cuDNN 9 provider ABI. Later current # releases have moved to CUDA 13, so this compatibility window is intentional. -# The Torch cu128 library directory is added to the Windows DLL search path by -# tools/uvr5/mdxnet.py before the FoxJoy ONNX model is loaded. onnxruntime-gpu>=1.19.2,<1.20 coloredlogs>=15.0,<16 opencv-python-headless>=4.10.0,<5 praat-parselmouth>=0.4.5,<1 -PyYAML>=6.0 +PyYAML>=6.0.1 scikit-learn>=1.6.0,<2 scipy>=1.13.1,<2 sounddevice>=0.5.0,<1 diff --git a/tools/pymss_webui.py b/tools/pymss_webui.py new file mode 100644 index 0000000..42ef0c2 --- /dev/null +++ b/tools/pymss_webui.py @@ -0,0 +1,406 @@ +import gc +import logging +import os +import subprocess +import threading +import time +import traceback +import uuid +from concurrent.futures import ThreadPoolExecutor, wait +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch + +from configs.config import Config + + +logger = logging.getLogger(__name__) +config = Config() +weight_pymss_root = Path(os.getenv("weight_pymss_root", "assets/pymss_weights")) + +MODEL_SAMPLE_RATE = 44100 +FFMPEG_PATH = Path(__file__).resolve().parents[2] / "ffmpeg.exe" +AUDIO_PARAMS = { + "wav_bit_depth": "FLOAT", + "flac_bit_depth": "PCM_24", + "mp3_bit_rate": "320k", + "m4a_bit_rate": "320k", + "m4a_codec": "aac", + "m4a_aac_at_quality": 2, +} + + +@dataclass(frozen=True) +class ModelSpec: + label: str + model_id: str + model_type: str + model_file: str + config_file: str + desired_stem: str + secondary_stem: str + desired_suffix: str + secondary_suffix: str + batch_size: int + overlap_size: int + + +MODEL_SPECS = ( + ModelSpec( + label="去混响", + model_id="dereverb-less-aggressive-18.8050", + model_type="mel_band_roformer", + model_file="dereverb_mel_band_roformer_less_aggressive_anvuew_sdr_18.8050.ckpt", + config_file="dereverb_mel_band_roformer_anvuew.yaml", + desired_stem="noreverb", + secondary_stem="reverb", + desired_suffix="noreverb", + secondary_suffix="reverb", + batch_size=1, + overlap_size=176400, + ), + ModelSpec( + label="去混响（激进）", + model_id="dereverb-anvuew-19.1729", + model_type="mel_band_roformer", + model_file="dereverb_mel_band_roformer_anvuew_sdr_19.1729.ckpt", + config_file="dereverb_mel_band_roformer_anvuew.yaml", + desired_stem="noreverb", + secondary_stem="reverb", + desired_suffix="noreverb", + secondary_suffix="reverb", + batch_size=1, + overlap_size=176400, + ), + ModelSpec( + label="去伴奏", + model_id="vocals-bs-roformer-368", + model_type="bs_roformer", + model_file="model_bs_roformer_ep_368_sdr_12.9628.ckpt", + config_file="model_bs_roformer_ep_368_sdr_12.9628.yaml", + desired_stem="vocals", + secondary_stem="instrumental", + desired_suffix="vocals", + secondary_suffix="instrumental", + batch_size=1, + overlap_size=264600, + ), + ModelSpec( + label="去伴奏（激进）", + model_id="vocals-bs-roformer-317", + model_type="bs_roformer", + model_file="model_bs_roformer_ep_317_sdr_12.9755.ckpt", + config_file="model_bs_roformer_ep_317_sdr_12.9755.yaml", + desired_stem="vocals", + secondary_stem="other", + desired_suffix="vocals", + secondary_suffix="instrumental", + batch_size=4, + overlap_size=176400, + ), + ModelSpec( + label="提主旋律", + model_id="karaoke-mel-roformer-10.1956", + model_type="mel_band_roformer", + model_file="model_mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt", + config_file="config_mel_band_roformer_karaoke.yaml", + desired_stem="karaoke", + secondary_stem="other", + desired_suffix="main_vocal", + secondary_suffix="off_vocal", + batch_size=1, + overlap_size=264600, + ), +) + +MODEL_BY_LABEL = {spec.label: spec for spec in MODEL_SPECS} +MODEL_BY_ID = {spec.model_id: spec for spec in MODEL_SPECS} +PYMSS_MODEL_CHOICES = [spec.label for spec in MODEL_SPECS] +UVR_INFERENCE_LOCK = threading.Lock() + + +def resolve_model(model_name): + if not model_name: + return MODEL_SPECS[0] + spec = MODEL_BY_LABEL.get(model_name) or MODEL_BY_ID.get(model_name) + if spec is None: + raise ValueError("Unknown separation model: %s" % model_name) + return spec + + +def get_model_info(model_name): + spec = resolve_model(model_name) + return "%s | %s" % (spec.model_type, spec.model_id) + + +def clean_path(path): + path = path or "" + if path.endswith(("\\", "/")): + path = path[:-1] + return path.replace("/", os.sep).replace("\\", os.sep).strip(" '\n\"\u202a") + + +def _uploaded_path(item): + if isinstance(item, (str, os.PathLike)): + return os.fspath(item) + if isinstance(item, dict): + return item.get("name") or item.get("path") + return getattr(item, "name", None) + + +def collect_input_paths(inp_root, paths): + inp_root = clean_path(inp_root) + if inp_root: + if os.path.isfile(inp_root): + candidates = [inp_root] + elif os.path.isdir(inp_root): + candidates = [os.path.join(inp_root, name) for name in sorted(os.listdir(inp_root))] + else: + raise FileNotFoundError(inp_root) + else: + candidates = [_uploaded_path(item) for item in (paths or [])] + return [os.path.abspath(path) for path in candidates if path and os.path.isfile(path)] + + +def _write_audio(path, audio, sample_rate, output_format): + audio = np.ascontiguousarray(audio, dtype=np.float32) + if audio.ndim == 1: + channels = 1 + elif audio.ndim == 2 and audio.shape[1] in (1, 2): + channels = audio.shape[1] + else: + raise ValueError("Unsupported audio shape: %s" % (audio.shape,)) + + if output_format == "wav": + sf.write(path, audio, sample_rate, format="WAV", subtype="FLOAT") + return + if output_format == "flac": + sf.write(path, audio, sample_rate, format="FLAC", subtype="PCM_24") + return + + ffmpeg = str(FFMPEG_PATH) if FFMPEG_PATH.is_file() else "ffmpeg" + command = [ + ffmpeg, + "-hide_banner", + "-loglevel", + "error", + "-y", + "-f", + "f32le", + "-ar", + str(sample_rate), + "-ac", + str(channels), + "-i", + "pipe:0", + "-vn", + ] + if output_format == "mp3": + command.extend(("-c:a", "libmp3lame", "-b:a", "320k")) + elif output_format == "m4a": + command.extend(("-c:a", "aac", "-aac_coder", "fast", "-b:a", "320k")) + else: + raise ValueError("Unsupported output format: %s" % output_format) + command.append(path) + completed = subprocess.run( + command, + input=audio.tobytes(), + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0, + ) + if completed.returncode != 0: + detail = completed.stderr.decode("utf-8", errors="replace").strip() + raise RuntimeError("FFmpeg audio encoding failed: %s" % detail) + + +class MSSTBatchSeparator: + def __init__(self, spec, output_format, desired_root, secondary_root): + try: + from pymss import MSSeparator, load_audio + except ImportError as error: + raise RuntimeError( + "缺少 pymss 运行库，请安装对应的 CUDA 版 Python 3.12 requirements" + ) from error + + self.spec = spec + self.output_format = output_format.lower() + if self.output_format not in {"wav", "flac", "mp3", "m4a"}: + raise ValueError("Unsupported output format: %s" % output_format) + desired_root = clean_path(desired_root) + secondary_root = clean_path(secondary_root) + if not desired_root or not secondary_root: + raise ValueError("输出文件夹不能为空") + self.desired_root = os.path.abspath(desired_root) + self.secondary_root = os.path.abspath(secondary_root) + os.makedirs(self.desired_root, exist_ok=True) + os.makedirs(self.secondary_root, exist_ok=True) + + model_path = weight_pymss_root / spec.model_file + config_path = weight_pymss_root / spec.config_file + if not model_path.is_file(): + raise FileNotFoundError(model_path) + if not config_path.is_file(): + raise FileNotFoundError(config_path) + + parsed_device = torch.device(config.device) + use_cuda = parsed_device.type == "cuda" + device_id = parsed_device.index if use_cuda and parsed_device.index is not None else 0 + self._load_audio = load_audio + self.model_load_count = 0 + self.separator = MSSeparator( + model_type=spec.model_type, + model_path=str(model_path), + config_path=str(config_path), + device="cuda" if use_cuda else "cpu", + device_ids=[device_id], + output_format=self.output_format, + use_tta=False, + store_dirs={}, + audio_params=AUDIO_PARAMS, + debug=False, + inference_params={ + "batch_size": spec.batch_size, + "chunk_size": 352800, + "overlap_size": spec.overlap_size, + "standardize": False, + "normalize": False, + "use_amp": bool(config.is_half and use_cuda), + "cuda_attention_backend": "default", + }, + ) + self.separator.config.training.use_amp = bool(config.is_half and use_cuda) + self._save_pool = ThreadPoolExecutor(max_workers=2, thread_name_prefix="rvc-msst-save") + self.model_load_count = 1 + logger.info( + "Loaded MSST model once for batch: %s, device=%s, half=%s", + spec.model_id, + self.separator.device, + bool(config.is_half and use_cuda), + ) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback_value): + self.close() + + def _save_output(self, audio, sample_rate, output_root, file_stem, suffix): + output_path = os.path.join( + output_root, + "%s_%s.%s" % (file_stem, suffix, self.output_format), + ) + temp_path = os.path.join( + output_root, + ".%s_%s.%s.tmp.%s" + % (file_stem, suffix, uuid.uuid4().hex, self.output_format), + ) + started = time.perf_counter() + try: + _write_audio(temp_path, audio, sample_rate, self.output_format) + if not os.path.isfile(temp_path) or os.path.getsize(temp_path) == 0: + raise RuntimeError("音频编码没有生成有效文件: %s" % output_path) + os.replace(temp_path, output_path) + except Exception: + if os.path.exists(temp_path): + os.remove(temp_path) + raise + return output_path, time.perf_counter() - started + + def separate_file(self, input_path): + mix, sample_rate = self._load_audio(input_path, sr=MODEL_SAMPLE_RATE, mono=False) + inference_started = time.perf_counter() + results = self.separator.separate(mix, pbar=True) + inference_seconds = time.perf_counter() - inference_started + missing = { + self.spec.desired_stem, + self.spec.secondary_stem, + }.difference(results) + if missing: + raise RuntimeError("模型缺少输出 stem: %s" % ", ".join(sorted(missing))) + + file_stem = Path(input_path).stem + encode_started = time.perf_counter() + futures = ( + self._save_pool.submit( + self._save_output, + results[self.spec.desired_stem], + sample_rate, + self.desired_root, + file_stem, + self.spec.desired_suffix, + ), + self._save_pool.submit( + self._save_output, + results[self.spec.secondary_stem], + sample_rate, + self.secondary_root, + file_stem, + self.spec.secondary_suffix, + ), + ) + wait(futures) + outputs = [future.result() for future in futures] + encode_seconds = time.perf_counter() - encode_started + del results, mix + return { + "outputs": [path for path, _ in outputs], + "inference_seconds": inference_seconds, + "encode_seconds": encode_seconds, + } + + def close(self): + save_pool = getattr(self, "_save_pool", None) + if save_pool is not None: + save_pool.shutdown(wait=True) + self._save_pool = None + separator = getattr(self, "separator", None) + try: + if separator is not None: + separator.close() + self.separator = None + finally: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def pymss_separate(model_name, inp_root, save_root_vocal, paths, save_root_ins, format0): + infos = [] + spec = resolve_model(model_name) + try: + input_paths = collect_input_paths(inp_root, paths) + if not input_paths: + raise ValueError("没有找到可处理的音频文件") + infos.append("%s | %s | 正在加载模型" % (spec.label, spec.model_id)) + yield "\n".join(infos) + with UVR_INFERENCE_LOCK: + with MSSTBatchSeparator( + spec, + format0, + save_root_vocal, + save_root_ins, + ) as batch: + for input_path in input_paths: + try: + result = batch.separate_file(input_path) + infos.append( + "%s -> 成功 | 推理 %.2fs | 编码 %.2fs" + % ( + os.path.basename(input_path), + result["inference_seconds"], + result["encode_seconds"], + ) + ) + except Exception: + infos.append( + "%s -> 失败\n%s" + % (os.path.basename(input_path), traceback.format_exc()) + ) + yield "\n".join(infos) + except Exception: + infos.append("失败\n%s" % traceback.format_exc()) + yield "\n".join(infos) diff --git a/tools/uvr5/bs_roformer/__init__.py b/tools/uvr5/bs_roformer/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tools/uvr5/bs_roformer/attend.py b/tools/uvr5/bs_roformer/attend.py deleted file mode 100644 index 1782bbd..0000000 --- a/tools/uvr5/bs_roformer/attend.py +++ /dev/null @@ -1,70 +0,0 @@ -from packaging import version -import torch -from torch import nn, einsum -import torch.nn.functional as F - - -def exists(val): - return val is not None - - -def default(v, d): - return v if exists(v) else d - - -class Attend(nn.Module): - def __init__(self, dropout=0.0, flash=False, scale=None): - super().__init__() - self.scale = scale - self.dropout = dropout - self.attn_dropout = nn.Dropout(dropout) - - self.flash = flash - assert not (flash and version.parse(torch.__version__) < version.parse("2.0.0")), ( - "in order to use flash attention, you must be using pytorch 2.0 or above" - ) - - def flash_attn(self, q, k, v): - # _, heads, q_len, _, k_len, is_cuda, device = *q.shape, k.shape[-2], q.is_cuda, q.device - - if exists(self.scale): - default_scale = q.shape[-1] ** -0.5 - q = q * (self.scale / default_scale) - - # pytorch 2.0 flash attn: q, k, v, mask, dropout, softmax_scale - # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True): - return F.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout if self.training else 0.0) - - def forward(self, q, k, v): - """ - einstein notation - b - batch - h - heads - n, i, j - sequence length (base sequence length, source, target) - d - feature dimension - """ - - # q_len, k_len, device = q.shape[-2], k.shape[-2], q.device - - scale = default(self.scale, q.shape[-1] ** -0.5) - - # DirectML does not expose PyTorch's SDPA kernels. Keep the existing - # SDPA path for CUDA/CPU/MPS and use the mathematically equivalent - # einsum implementation below for PrivateUse1 tensors. - if self.flash and q.device.type != "privateuseone": - return self.flash_attn(q, k, v) - - # similarity - - sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale - - # attention - - attn = sim.softmax(dim=-1) - attn = self.attn_dropout(attn) - - # aggregate values - - out = einsum("b h i j, b h j d -> b h i d", attn, v) - - return out diff --git a/tools/uvr5/bs_roformer/bs_roformer.py b/tools/uvr5/bs_roformer/bs_roformer.py deleted file mode 100644 index d2ef735..0000000 --- a/tools/uvr5/bs_roformer/bs_roformer.py +++ /dev/null @@ -1,356 +0,0 @@ -from functools import partial -import torch -from torch import nn -from torch.nn import Module, ModuleList -import torch.nn.functional as F -from tools.uvr5.bs_roformer.attend import Attend -from torch.utils.checkpoint import checkpoint -from typing import Tuple, Optional, Callable -from tools.uvr5.rotary_embedding_torch import RotaryEmbedding -from einops import rearrange, pack, unpack -from einops.layers.torch import Rearrange - -def exists(val): - return val is not None - -def default(v, d): - return v if exists(v) else d - -def pack_one(t, pattern): - return pack([t], pattern) - -def unpack_one(t, ps, pattern): - return unpack(t, ps, pattern)[0] - -def l2norm(t): - return F.normalize(t, dim=-1, p=2) - -class RMSNorm(Module): - - def __init__(self, dim): - super().__init__() - self.scale = dim ** 0.5 - self.gamma = nn.Parameter(torch.ones(dim)) - - def forward(self, x): - return F.normalize(x, dim=-1) * self.scale * self.gamma - -class FeedForward(Module): - - def __init__(self, dim, mult=4, dropout=0.0): - super().__init__() - dim_inner = int(dim * mult) - self.net = nn.Sequential(RMSNorm(dim), nn.Linear(dim, dim_inner), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim_inner, dim), nn.Dropout(dropout)) - - def forward(self, x): - return self.net(x) - -class Attention(Module): - - def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, rotary_embed=None, flash=True): - super().__init__() - self.heads = heads - self.scale = dim_head ** (-0.5) - dim_inner = heads * dim_head - self.rotary_embed = rotary_embed - self.attend = Attend(flash=flash, dropout=dropout) - self.norm = RMSNorm(dim) - self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) - self.to_gates = nn.Linear(dim, heads) - self.to_out = nn.Sequential(nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout)) - - def forward(self, x): - x = self.norm(x) - (q, k, v) = rearrange(self.to_qkv(x), 'b n (qkv h d) -> qkv b h n d', qkv=3, h=self.heads) - if exists(self.rotary_embed): - q = self.rotary_embed.rotate_queries_or_keys(q) - k = self.rotary_embed.rotate_queries_or_keys(k) - out = self.attend(q, k, v) - gates = self.to_gates(x) - out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid() - out = rearrange(out, 'b h n d -> b n (h d)') - return self.to_out(out) - -class LinearAttention(Module): - """ - this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al. - """ - - def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0): - super().__init__() - dim_inner = dim_head * heads - self.norm = RMSNorm(dim) - self.to_qkv = nn.Sequential(nn.Linear(dim, dim_inner * 3, bias=False), Rearrange('b n (qkv h d) -> qkv b h d n', qkv=3, h=heads)) - self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) - self.attend = Attend(scale=scale, dropout=dropout, flash=flash) - self.to_out = nn.Sequential(Rearrange('b h d n -> b n (h d)'), nn.Linear(dim_inner, dim, bias=False)) - - def forward(self, x): - x = self.norm(x) - (q, k, v) = self.to_qkv(x) - (q, k) = map(l2norm, (q, k)) - q = q * self.temperature.exp() - out = self.attend(q, k, v) - return self.to_out(out) - -class Transformer(Module): - - def __init__(self, *, dim, depth, dim_head=64, heads=8, attn_dropout=0.0, ff_dropout=0.0, ff_mult=4, norm_output=True, rotary_embed=None, flash_attn=True, linear_attn=False): - super().__init__() - self.layers = ModuleList([]) - for _ in range(depth): - if linear_attn: - attn = LinearAttention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, flash=flash_attn) - else: - attn = Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, rotary_embed=rotary_embed, flash=flash_attn) - self.layers.append(ModuleList([attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)])) - self.norm = RMSNorm(dim) if norm_output else nn.Identity() - - def forward(self, x): - for (attn, ff) in self.layers: - x = attn(x) + x - x = ff(x) + x - return self.norm(x) - -class BandSplit(Module): - - def __init__(self, dim, dim_inputs): - super().__init__() - self.dim_inputs = dim_inputs - self.to_features = ModuleList([]) - for dim_in in dim_inputs: - net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) - self.to_features.append(net) - - def forward(self, x): - x = x.split(self.dim_inputs, dim=-1) - outs = [] - for (split_input, to_feature) in zip(x, self.to_features): - split_output = to_feature(split_input) - outs.append(split_output) - return torch.stack(outs, dim=-2) - -def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): - dim_hidden = default(dim_hidden, dim_in) - net = [] - dims = (dim_in, *(dim_hidden,) * (depth - 1), dim_out) - for (ind, (layer_dim_in, layer_dim_out)) in enumerate(zip(dims[:-1], dims[1:])): - is_last = ind == len(dims) - 2 - net.append(nn.Linear(layer_dim_in, layer_dim_out)) - if is_last: - continue - net.append(activation()) - return nn.Sequential(*net) - -class MaskEstimator(Module): - - def __init__(self, dim, dim_inputs, depth, mlp_expansion_factor=4): - super().__init__() - self.dim_inputs = dim_inputs - self.to_freqs = ModuleList([]) - dim_hidden = dim * mlp_expansion_factor - for dim_in in dim_inputs: - net = [] - mlp = nn.Sequential(MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1)) - self.to_freqs.append(mlp) - - def forward(self, x): - x = x.unbind(dim=-2) - outs = [] - for (band_features, mlp) in zip(x, self.to_freqs): - freq_out = mlp(band_features) - outs.append(freq_out) - return torch.cat(outs, dim=-1) -DEFAULT_FREQS_PER_BANDS = (2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 12, 12, 12, 12, 12, 12, 12, 12, 24, 24, 24, 24, 24, 24, 24, 24, 48, 48, 48, 48, 48, 48, 48, 48, 128, 129) - -class BSRoformer(Module): - - def __init__(self, dim, *, depth, stereo=False, num_stems=1, time_transformer_depth=2, freq_transformer_depth=2, linear_transformer_depth=0, freqs_per_bands=DEFAULT_FREQS_PER_BANDS, dim_head=64, heads=8, attn_dropout=0.0, ff_dropout=0.0, flash_attn=True, dim_freqs_in=1025, stft_n_fft=2048, stft_hop_length=512, stft_win_length=2048, stft_normalized=False, stft_window_fn=None, mask_estimator_depth=2, multi_stft_resolution_loss_weight=1.0, multi_stft_resolutions_window_sizes=(4096, 2048, 1024, 512, 256), multi_stft_hop_size=147, multi_stft_normalized=False, multi_stft_window_fn=torch.hann_window, mlp_expansion_factor=4, use_torch_checkpoint=False, skip_connection=False): - super().__init__() - self.stereo = stereo - self.audio_channels = 2 if stereo else 1 - self.num_stems = num_stems - self.use_torch_checkpoint = use_torch_checkpoint - self.skip_connection = skip_connection - self.layers = ModuleList([]) - transformer_kwargs = dict(dim=dim, heads=heads, dim_head=dim_head, attn_dropout=attn_dropout, ff_dropout=ff_dropout, flash_attn=flash_attn, norm_output=False) - time_rotary_embed = RotaryEmbedding(dim=dim_head) - freq_rotary_embed = RotaryEmbedding(dim=dim_head) - for _ in range(depth): - tran_modules = [] - if linear_transformer_depth > 0: - tran_modules.append(Transformer(depth=linear_transformer_depth, linear_attn=True, **transformer_kwargs)) - tran_modules.append(Transformer(depth=time_transformer_depth, rotary_embed=time_rotary_embed, **transformer_kwargs)) - tran_modules.append(Transformer(depth=freq_transformer_depth, rotary_embed=freq_rotary_embed, **transformer_kwargs)) - self.layers.append(nn.ModuleList(tran_modules)) - 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)}' - freqs_per_bands_with_complex = tuple((2 * f * self.audio_channels for f in freqs_per_bands)) - self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) - self.mask_estimators = nn.ModuleList([]) - for _ in range(num_stems): - mask_estimator = MaskEstimator(dim=dim, dim_inputs=freqs_per_bands_with_complex, depth=mask_estimator_depth, mlp_expansion_factor=mlp_expansion_factor) - self.mask_estimators.append(mask_estimator) - self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight - self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes - self.multi_stft_n_fft = stft_n_fft - 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 - - b - batch - f - freq - t - time - s - audio channel (1 for mono, 2 for stereo) - n - number of 'stems' - c - complex (2) - d - feature dimension - """ - device = raw_audio.device - x_is_dml = device.type == 'privateuseone' - x_is_mps = True if device.type == 'mps' else False - if raw_audio.ndim == 2: - raw_audio = rearrange(raw_audio, 'b t -> b 1 t') - channels = raw_audio.shape[1] - assert not self.stereo and channels == 1 or (self.stereo and channels == 2), 'stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)' - (raw_audio, batch_audio_channel_packed_shape) = pack_one(raw_audio, '* t') - 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._get_stft_window('cpu') - stft_complex = torch.stft( - raw_audio.cpu(), - **self.stft_kwargs, - window=stft_window, - return_complex=True, - ) - stft_repr_cpu = torch.view_as_real(stft_complex) - stft_repr_cpu = unpack_one( - stft_repr_cpu, batch_audio_channel_packed_shape, '* f t c' - ) - stft_repr_cpu = rearrange( - stft_repr_cpu, 'b s f t c -> b (f s) t c' - ) - stft_repr = stft_repr_cpu.to(device) - else: - stft_window = self._get_stft_window(device) - try: - stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True) - except: - stft_repr = torch.stft(raw_audio.cpu() if x_is_mps else raw_audio, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=True).to(device) - 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') - x = rearrange(stft_repr, 'b f t c -> b t (f c)') - if self.use_torch_checkpoint: - x = checkpoint(self.band_split, x, use_reentrant=False) - else: - x = self.band_split(x) - store = [None] * len(self.layers) - for (i, transformer_block) in enumerate(self.layers): - if len(transformer_block) == 3: - (linear_transformer, time_transformer, freq_transformer) = transformer_block - (x, ft_ps) = pack([x], 'b * d') - if self.use_torch_checkpoint: - x = checkpoint(linear_transformer, x, use_reentrant=False) - else: - x = linear_transformer(x) - (x,) = unpack(x, ft_ps, 'b * d') - else: - (time_transformer, freq_transformer) = transformer_block - if self.skip_connection: - for j in range(i): - x = x + store[j] - x = rearrange(x, 'b t f d -> b f t d') - (x, ps) = pack([x], '* t d') - if self.use_torch_checkpoint: - x = checkpoint(time_transformer, x, use_reentrant=False) - else: - x = time_transformer(x) - (x,) = unpack(x, ps, '* t d') - x = rearrange(x, 'b f t d -> b t f d') - (x, ps) = pack([x], '* f d') - if self.use_torch_checkpoint: - x = checkpoint(freq_transformer, x, use_reentrant=False) - else: - x = freq_transformer(x) - (x,) = unpack(x, ps, '* f d') - if self.skip_connection: - store[i] = x - x = self.final_norm(x) - num_stems = len(self.mask_estimators) - if self.use_torch_checkpoint: - mask = torch.stack([checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators], dim=1) - else: - mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1) - mask = rearrange(mask, 'b n t (f c) -> b n f t c', c=2) - if x_is_dml: - # Complex masking and ISTFT stay on CPU; all learned real-valued - # layers above remain on DirectML. - stft_repr = rearrange(stft_repr_cpu, 'b f t c -> b 1 f t c') - stft_repr = torch.view_as_complex(stft_repr.contiguous()) - mask = torch.view_as_complex(mask.float().cpu().contiguous()) - stft_repr = stft_repr * mask - stft_repr = rearrange( - stft_repr, - 'b n (f s) t -> (b n s) f t', - s=self.audio_channels, - ) - recon_audio = torch.istft( - stft_repr, - **self.stft_kwargs, - window=stft_window, - return_complex=False, - length=raw_audio.shape[-1], - ).to(device) - else: - stft_repr = rearrange(stft_repr, 'b f t c -> b 1 f t c') - stft_repr = torch.view_as_complex(stft_repr) - mask = torch.view_as_complex(mask) - stft_repr = stft_repr * mask - stft_repr = rearrange(stft_repr, 'b n (f s) t -> (b n s) f t', s=self.audio_channels) - try: - recon_audio = torch.istft(stft_repr, **self.stft_kwargs, window=stft_window, return_complex=False, length=raw_audio.shape[-1]) - except: - recon_audio = torch.istft(stft_repr.cpu() if x_is_mps else stft_repr, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=False, length=raw_audio.shape[-1]).to(device) - recon_audio = rearrange(recon_audio, '(b n s) t -> b n s t', s=self.audio_channels, n=num_stems) - if num_stems == 1: - recon_audio = rearrange(recon_audio, 'b 1 s t -> b s t') - if not exists(target): - return recon_audio - if self.num_stems > 1: - assert target.ndim == 4 and target.shape[1] == self.num_stems - if target.ndim == 2: - target = rearrange(target, '... t -> ... 1 t') - target = target[..., :recon_audio.shape[-1]] - loss_audio = recon_audio.cpu() if x_is_dml else recon_audio - loss_target = target.cpu() if x_is_dml else target - loss = F.l1_loss(loss_audio, loss_target) - multi_stft_resolution_loss = 0.0 - for window_size in self.multi_stft_resolutions_window_sizes: - spectral_device = 'cpu' if x_is_dml else device - res_stft_kwargs = dict(n_fft=max(window_size, self.multi_stft_n_fft), win_length=window_size, return_complex=True, window=self.multi_stft_window_fn(window_size, device=spectral_device), **self.multi_stft_kwargs) - recon_Y = torch.stft(rearrange(loss_audio, '... s t -> (... s) t'), **res_stft_kwargs) - target_Y = torch.stft(rearrange(loss_target, '... s t -> (... s) t'), **res_stft_kwargs) - multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(recon_Y, target_Y) - weighted_multi_resolution_loss = multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight - total_loss = loss + weighted_multi_resolution_loss - if not return_loss_breakdown: - return total_loss - return (total_loss, (loss, multi_stft_resolution_loss)) diff --git a/tools/uvr5/bs_roformer/mel_band_roformer.py b/tools/uvr5/bs_roformer/mel_band_roformer.py deleted file mode 100644 index 4e62d2e..0000000 --- a/tools/uvr5/bs_roformer/mel_band_roformer.py +++ /dev/null @@ -1,361 +0,0 @@ -from functools import partial -import torch -from torch import nn -from torch.nn import Module, ModuleList -import torch.nn.functional as F -from tools.uvr5.bs_roformer.attend import Attend -from torch.utils.checkpoint import checkpoint -from typing import Tuple, Optional, Callable -from tools.uvr5.rotary_embedding_torch import RotaryEmbedding -from einops import rearrange, pack, unpack, reduce, repeat -from einops.layers.torch import Rearrange -from librosa import filters - -def exists(val): - return val is not None - -def default(v, d): - return v if exists(v) else d - -def pack_one(t, pattern): - return pack([t], pattern) - -def unpack_one(t, ps, pattern): - return unpack(t, ps, pattern)[0] - -def pad_at_dim(t, pad, dim=-1, value=0.0): - dims_from_right = -dim - 1 if dim < 0 else t.ndim - dim - 1 - zeros = (0, 0) * dims_from_right - return F.pad(t, (*zeros, *pad), value=value) - -def l2norm(t): - return F.normalize(t, dim=-1, p=2) - -class RMSNorm(Module): - - def __init__(self, dim): - super().__init__() - self.scale = dim ** 0.5 - self.gamma = nn.Parameter(torch.ones(dim)) - - def forward(self, x): - return F.normalize(x, dim=-1) * self.scale * self.gamma - -class FeedForward(Module): - - def __init__(self, dim, mult=4, dropout=0.0): - super().__init__() - dim_inner = int(dim * mult) - self.net = nn.Sequential(RMSNorm(dim), nn.Linear(dim, dim_inner), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim_inner, dim), nn.Dropout(dropout)) - - def forward(self, x): - return self.net(x) - -class Attention(Module): - - def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, rotary_embed=None, flash=True): - super().__init__() - self.heads = heads - self.scale = dim_head ** (-0.5) - dim_inner = heads * dim_head - self.rotary_embed = rotary_embed - self.attend = Attend(flash=flash, dropout=dropout) - self.norm = RMSNorm(dim) - self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) - self.to_gates = nn.Linear(dim, heads) - self.to_out = nn.Sequential(nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout)) - - def forward(self, x): - x = self.norm(x) - (q, k, v) = rearrange(self.to_qkv(x), 'b n (qkv h d) -> qkv b h n d', qkv=3, h=self.heads) - if exists(self.rotary_embed): - q = self.rotary_embed.rotate_queries_or_keys(q) - k = self.rotary_embed.rotate_queries_or_keys(k) - out = self.attend(q, k, v) - gates = self.to_gates(x) - out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid() - out = rearrange(out, 'b h n d -> b n (h d)') - return self.to_out(out) - -class LinearAttention(Module): - """ - this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al. - """ - - def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0): - super().__init__() - dim_inner = dim_head * heads - self.norm = RMSNorm(dim) - self.to_qkv = nn.Sequential(nn.Linear(dim, dim_inner * 3, bias=False), Rearrange('b n (qkv h d) -> qkv b h d n', qkv=3, h=heads)) - self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) - self.attend = Attend(scale=scale, dropout=dropout, flash=flash) - self.to_out = nn.Sequential(Rearrange('b h d n -> b n (h d)'), nn.Linear(dim_inner, dim, bias=False)) - - def forward(self, x): - x = self.norm(x) - (q, k, v) = self.to_qkv(x) - (q, k) = map(l2norm, (q, k)) - q = q * self.temperature.exp() - out = self.attend(q, k, v) - return self.to_out(out) - -class Transformer(Module): - - def __init__(self, *, dim, depth, dim_head=64, heads=8, attn_dropout=0.0, ff_dropout=0.0, ff_mult=4, norm_output=True, rotary_embed=None, flash_attn=True, linear_attn=False): - super().__init__() - self.layers = ModuleList([]) - for _ in range(depth): - if linear_attn: - attn = LinearAttention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, flash=flash_attn) - else: - attn = Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, rotary_embed=rotary_embed, flash=flash_attn) - self.layers.append(ModuleList([attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)])) - self.norm = RMSNorm(dim) if norm_output else nn.Identity() - - def forward(self, x): - for (attn, ff) in self.layers: - x = attn(x) + x - x = ff(x) + x - return self.norm(x) - -class BandSplit(Module): - - def __init__(self, dim, dim_inputs): - super().__init__() - self.dim_inputs = dim_inputs - self.to_features = ModuleList([]) - for dim_in in dim_inputs: - net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) - self.to_features.append(net) - - def forward(self, x): - x = x.split(self.dim_inputs, dim=-1) - outs = [] - for (split_input, to_feature) in zip(x, self.to_features): - split_output = to_feature(split_input) - outs.append(split_output) - return torch.stack(outs, dim=-2) - -def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): - dim_hidden = default(dim_hidden, dim_in) - net = [] - dims = (dim_in, *(dim_hidden,) * depth, dim_out) - for (ind, (layer_dim_in, layer_dim_out)) in enumerate(zip(dims[:-1], dims[1:])): - is_last = ind == len(dims) - 2 - net.append(nn.Linear(layer_dim_in, layer_dim_out)) - if is_last: - continue - net.append(activation()) - return nn.Sequential(*net) - -class MaskEstimator(Module): - - def __init__(self, dim, dim_inputs, depth, mlp_expansion_factor=4): - super().__init__() - self.dim_inputs = dim_inputs - self.to_freqs = ModuleList([]) - dim_hidden = dim * mlp_expansion_factor - for dim_in in dim_inputs: - net = [] - mlp = nn.Sequential(MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1)) - self.to_freqs.append(mlp) - - def forward(self, x): - x = x.unbind(dim=-2) - outs = [] - for (band_features, mlp) in zip(x, self.to_freqs): - freq_out = mlp(band_features) - outs.append(freq_out) - return torch.cat(outs, dim=-1) - -class MelBandRoformer(Module): - - def __init__(self, dim, *, depth, stereo=False, num_stems=1, time_transformer_depth=2, freq_transformer_depth=2, linear_transformer_depth=0, num_bands=60, dim_head=64, heads=8, attn_dropout=0.1, ff_dropout=0.1, flash_attn=True, dim_freqs_in=1025, sample_rate=44100, stft_n_fft=2048, stft_hop_length=512, stft_win_length=2048, stft_normalized=False, stft_window_fn=None, mask_estimator_depth=1, multi_stft_resolution_loss_weight=1.0, multi_stft_resolutions_window_sizes=(4096, 2048, 1024, 512, 256), multi_stft_hop_size=147, multi_stft_normalized=False, multi_stft_window_fn=torch.hann_window, match_input_audio_length=False, mlp_expansion_factor=4, use_torch_checkpoint=False, skip_connection=False): - super().__init__() - self.stereo = stereo - self.audio_channels = 2 if stereo else 1 - self.num_stems = num_stems - self.use_torch_checkpoint = use_torch_checkpoint - self.skip_connection = skip_connection - self.layers = ModuleList([]) - transformer_kwargs = dict(dim=dim, heads=heads, dim_head=dim_head, attn_dropout=attn_dropout, ff_dropout=ff_dropout, flash_attn=flash_attn) - time_rotary_embed = RotaryEmbedding(dim=dim_head) - freq_rotary_embed = RotaryEmbedding(dim=dim_head) - for _ in range(depth): - tran_modules = [] - if linear_transformer_depth > 0: - tran_modules.append(Transformer(depth=linear_transformer_depth, linear_attn=True, **transformer_kwargs)) - tran_modules.append(Transformer(depth=time_transformer_depth, rotary_embed=time_rotary_embed, **transformer_kwargs)) - 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) - mel_filter_bank = torch.from_numpy(mel_filter_bank_numpy) - mel_filter_bank[0][0] = 1.0 - mel_filter_bank[-1, -1] = 1.0 - freqs_per_band = mel_filter_bank > 0 - assert freqs_per_band.any(dim=0).all(), 'all frequencies need to be covered by all bands for now' - repeated_freq_indices = repeat(torch.arange(freqs), 'f -> b f', b=num_bands) - freq_indices = repeated_freq_indices[freqs_per_band] - if stereo: - freq_indices = repeat(freq_indices, 'f -> f s', s=2) - freq_indices = freq_indices * 2 + torch.arange(2) - freq_indices = rearrange(freq_indices, 'f s -> (f s)') - self.register_buffer('freq_indices', freq_indices, persistent=False) - self.register_buffer('freqs_per_band', freqs_per_band, persistent=False) - num_freqs_per_band = reduce(freqs_per_band, 'b f -> b', 'sum') - num_bands_per_freq = reduce(freqs_per_band, 'b f -> f', 'sum') - self.register_buffer('num_freqs_per_band', num_freqs_per_band, persistent=False) - self.register_buffer('num_bands_per_freq', num_bands_per_freq, persistent=False) - freqs_per_bands_with_complex = tuple((2 * f * self.audio_channels for f in num_freqs_per_band.tolist())) - self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) - self.mask_estimators = nn.ModuleList([]) - for _ in range(num_stems): - mask_estimator = MaskEstimator(dim=dim, dim_inputs=freqs_per_bands_with_complex, depth=mask_estimator_depth, mlp_expansion_factor=mlp_expansion_factor) - self.mask_estimators.append(mask_estimator) - self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight - self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes - self.multi_stft_n_fft = stft_n_fft - self.multi_stft_window_fn = multi_stft_window_fn - 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 - - b - batch - f - freq - t - time - s - audio channel (1 for mono, 2 for stereo) - n - number of 'stems' - c - complex (2) - d - feature dimension - """ - device = raw_audio.device - x_is_dml = device.type == 'privateuseone' - if raw_audio.ndim == 2: - raw_audio = rearrange(raw_audio, 'b t -> b 1 t') - (batch, channels, raw_audio_length) = raw_audio.shape - istft_length = raw_audio_length if self.match_input_audio_length else None - assert not self.stereo and channels == 1 or (self.stereo and channels == 2), 'stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)' - (raw_audio, batch_audio_channel_packed_shape) = pack_one(raw_audio, '* t') - 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._get_stft_window('cpu') - stft_complex = torch.stft( - raw_audio.cpu(), - **self.stft_kwargs, - window=stft_window, - return_complex=True, - ) - stft_repr = torch.view_as_real(stft_complex) - 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') - x = stft_repr[:, self.freq_indices.cpu()].to(device) - else: - 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') - 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) - else: - x = self.band_split(x) - store = [None] * len(self.layers) - for (i, transformer_block) in enumerate(self.layers): - if len(transformer_block) == 3: - (linear_transformer, time_transformer, freq_transformer) = transformer_block - (x, ft_ps) = pack([x], 'b * d') - if self.use_torch_checkpoint: - x = checkpoint(linear_transformer, x, use_reentrant=False) - else: - x = linear_transformer(x) - (x,) = unpack(x, ft_ps, 'b * d') - else: - (time_transformer, freq_transformer) = transformer_block - if self.skip_connection: - for j in range(i): - x = x + store[j] - x = rearrange(x, 'b t f d -> b f t d') - (x, ps) = pack([x], '* t d') - if self.use_torch_checkpoint: - x = checkpoint(time_transformer, x, use_reentrant=False) - else: - x = time_transformer(x) - (x,) = unpack(x, ps, '* t d') - x = rearrange(x, 'b f t d -> b t f d') - (x, ps) = pack([x], '* f d') - if self.use_torch_checkpoint: - x = checkpoint(freq_transformer, x, use_reentrant=False) - else: - x = freq_transformer(x) - (x,) = unpack(x, ps, '* f d') - if self.skip_connection: - store[i] = x - num_stems = len(self.mask_estimators) - if self.use_torch_checkpoint: - masks = torch.stack([checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators], dim=1) - else: - masks = torch.stack([fn(x) for fn in self.mask_estimators], dim=1) - masks = rearrange(masks, 'b n t (f c) -> b n f t c', c=2) - if x_is_dml: - masks = masks.float().cpu() - stft_repr = rearrange(stft_repr, 'b f t c -> b 1 f t c') - stft_repr = torch.view_as_complex(stft_repr.contiguous()) - 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 - stft_repr_expanded_stems = repeat(stft_repr, 'b 1 ... -> b n ...', n=num_stems) - 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) - stft_repr = stft_repr * masks_averaged - stft_repr = rearrange(stft_repr, 'b n (f s) t -> (b n s) f t', s=self.audio_channels) - recon_audio = torch.istft(stft_repr, **self.stft_kwargs, window=stft_window, return_complex=False, length=istft_length) - if x_is_dml: - recon_audio = recon_audio.to(device) - recon_audio = rearrange(recon_audio, '(b n s) t -> b n s t', b=batch, s=self.audio_channels, n=num_stems) - if num_stems == 1: - recon_audio = rearrange(recon_audio, 'b 1 s t -> b s t') - if not exists(target): - return recon_audio - if self.num_stems > 1: - assert target.ndim == 4 and target.shape[1] == self.num_stems - if target.ndim == 2: - target = rearrange(target, '... t -> ... 1 t') - target = target[..., :recon_audio.shape[-1]] - loss_audio = recon_audio.cpu() if x_is_dml else recon_audio - loss_target = target.cpu() if x_is_dml else target - loss = F.l1_loss(loss_audio, loss_target) - multi_stft_resolution_loss = 0.0 - for window_size in self.multi_stft_resolutions_window_sizes: - spectral_device = 'cpu' if x_is_dml else device - res_stft_kwargs = dict(n_fft=max(window_size, self.multi_stft_n_fft), win_length=window_size, return_complex=True, window=self.multi_stft_window_fn(window_size, device=spectral_device), **self.multi_stft_kwargs) - recon_Y = torch.stft(rearrange(loss_audio, '... s t -> (... s) t'), **res_stft_kwargs) - target_Y = torch.stft(rearrange(loss_target, '... s t -> (... s) t'), **res_stft_kwargs) - multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(recon_Y, target_Y) - weighted_multi_resolution_loss = multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight - total_loss = loss + weighted_multi_resolution_loss - if not return_loss_breakdown: - return total_loss - return (total_loss, (loss, multi_stft_resolution_loss)) diff --git a/tools/uvr5/bsroformer.py b/tools/uvr5/bsroformer.py deleted file mode 100644 index 66e08ac..0000000 --- a/tools/uvr5/bsroformer.py +++ /dev/null @@ -1,405 +0,0 @@ -# This code is modified from https://github.com/ZFTurbo/ -import os -import warnings -from contextlib import nullcontext - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -import yaml - -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 - -warnings.filterwarnings("ignore") -i18n = I18nAuto() - - -class Roformer_Loader: - def get_config(self, config_path): - return yaml.load(read_text(config_path), Loader=yaml.FullLoader) - - def get_default_config(self): - default_config = None - if self.model_type == "bs_roformer": - # Use model_bs_roformer_ep_368_sdr_12.9628.yaml and model_bs_roformer_ep_317_sdr_12.9755.yaml as default configuration files - # Other BS_Roformer models may not be compatible - # fmt: off - default_config = { - "audio": {"chunk_size": 352800, "sample_rate": 44100}, - "model": { - "dim": 512, - "depth": 12, - "stereo": True, - "num_stems": 1, - "time_transformer_depth": 1, - "freq_transformer_depth": 1, - "linear_transformer_depth": 0, - "freqs_per_bands": (2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 12, 12, 12, 12, 12, 12, 12, 12, 24, 24, 24, 24, 24, 24, 24, 24, 48, 48, 48, 48, 48, 48, 48, 48, 128, 129), - "dim_head": 64, - "heads": 8, - "attn_dropout": 0.1, - "ff_dropout": 0.1, - "flash_attn": True, - "dim_freqs_in": 1025, - "stft_n_fft": 2048, - "stft_hop_length": 441, - "stft_win_length": 2048, - "stft_normalized": False, - "mask_estimator_depth": 2, - "multi_stft_resolution_loss_weight": 1.0, - "multi_stft_resolutions_window_sizes": (4096, 2048, 1024, 512, 256), - "multi_stft_hop_size": 147, - "multi_stft_normalized": False, - }, - "training": {"instruments": ["vocals", "other"], "target_instrument": "vocals"}, - "inference": {"batch_size": 2, "num_overlap": 2}, - } - # fmt: on - elif self.model_type == "mel_band_roformer": - # Use model_mel_band_roformer_ep_3005_sdr_11.4360.yaml as default configuration files - # Other Mel_Band_Roformer models may not be compatible - default_config = { - "audio": {"chunk_size": 352800, "sample_rate": 44100}, - "model": { - "dim": 384, - "depth": 12, - "stereo": True, - "num_stems": 1, - "time_transformer_depth": 1, - "freq_transformer_depth": 1, - "linear_transformer_depth": 0, - "num_bands": 60, - "dim_head": 64, - "heads": 8, - "attn_dropout": 0.1, - "ff_dropout": 0.1, - "flash_attn": True, - "dim_freqs_in": 1025, - "sample_rate": 44100, - "stft_n_fft": 2048, - "stft_hop_length": 441, - "stft_win_length": 2048, - "stft_normalized": False, - "mask_estimator_depth": 2, - "multi_stft_resolution_loss_weight": 1.0, - "multi_stft_resolutions_window_sizes": (4096, 2048, 1024, 512, 256), - "multi_stft_hop_size": 147, - "multi_stft_normalized": False, - }, - "training": {"instruments": ["vocals", "other"], "target_instrument": "vocals"}, - "inference": {"batch_size": 2, "num_overlap": 2}, - } - - return default_config - - def get_model_from_config(self): - if self.model_type == "bs_roformer": - from tools.uvr5.bs_roformer.bs_roformer import BSRoformer - - model = BSRoformer(**dict(self.config["model"])) - elif self.model_type == "mel_band_roformer": - from tools.uvr5.bs_roformer.mel_band_roformer import MelBandRoformer - - model = MelBandRoformer(**dict(self.config["model"])) - else: - print(i18n("错误：未知模型：%s") % self.model_type) - model = None - return model - - def demix_track(self, model, mix, device): - C = self.config["audio"]["chunk_size"] # chunk_size - N = self.config["inference"]["num_overlap"] - fade_size = C // 10 - step = int(C // N) - border = C - step - batch_size = self.config["inference"]["batch_size"] - - length_init = mix.shape[-1] - - # 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) - - 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 - - amp_context = ( - torch.amp.autocast("cuda", enabled=self.is_half) - if device_type == "cuda" - else nullcontext() - ) - grad_context = ( - torch.no_grad() - if device_type == "privateuseone" - else torch.inference_mode() - ) - with amp_context: - # DirectML updates version counters in several linear kernels and - # therefore needs no_grad rather than inference_mode. CUDA and CPU - # retain the existing inference-mode path. - with grad_context: - model_dtype = next(model.parameters()).dtype - i = 0 - batch_data = [] - batch_locations = [] - while i < mix.shape[1]: - 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) - 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).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() - ) - for j in range(len(batch_locations)): - start, l = batch_locations[j] - 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 = [] - - result.div_(counter.clamp_min(1e-8)) - torch.nan_to_num_(result) - if length_init > 2 * border and (border > 0): - result = result[..., border:-border] - estimated_sources = result.cpu().numpy() - - progress_bar.close() - - if self.config["training"]["target_instrument"] is None: - return {k: v for k, v in zip(self.config["training"]["instruments"], estimated_sources)} - else: - return {k: v for k, v in zip([self.config["training"]["target_instrument"]], estimated_sources)} - - def run_folder(self, input, vocal_root, others_root, format): - self.model.eval() - path = input - os.makedirs(vocal_root, exist_ok=True) - os.makedirs(others_root, exist_ok=True) - file_base_name = os.path.splitext(os.path.basename(path))[0] - - sample_rate = 44100 - 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: - 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 - - 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("音频包含多个声道，但模型仅支持单声道，将对所有声道取平均值")) - - 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: - # if target instrument is specified, save target instrument as vocal and other instruments as others - # 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] - 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]) - self.save_audio(path_vocal, res[target_instrument].T, sr, format) - self.save_audio(path_other, other.T, sr, format) - else: - # if target instrument is not specified, save the first instrument as vocal and the rest as others - vocal_inst = self.config["training"]["instruments"][0] - path_vocal = "{}/{}_{}.wav".format(vocal_root, file_base_name, vocal_inst) - self.save_audio(path_vocal, res[vocal_inst].T, sr, format) - for other in self.config["training"]["instruments"][1:]: # save other instruments - path_other = "{}/{}_{}.wav".format(others_root, file_base_name, other) - self.save_audio(path_other, res[other].T, sr, format) - - def save_audio(self, path, data, sr, format): - # input path should be endwith '.wav' - if format in ["wav", "flac"]: - if format == "flac": - path = path[:-3] + "flac" - sf.write(path, data, sr) - else: - sf.write(path, data, sr) - os.system('ffmpeg -i "{}" -vn "{}" -q:a 2 -y'.format(path, path[:-3] + format)) - try: - os.remove(path) - except: - pass - - def __init__(self, model_path, config_path, device, is_half): - self.device = device - self.is_half = is_half - self.model_type = None - self.config = None - - # get model_type, first try: - if "bs_roformer" in model_path.lower() or "bsroformer" in model_path.lower(): - self.model_type = "bs_roformer" - elif "mel_band_roformer" in model_path.lower() or "melbandroformer" in model_path.lower(): - self.model_type = "mel_band_roformer" - - if not os.path.exists(config_path): - if self.model_type is None: - # if model_type is still None, raise an error - raise ValueError( - "Error: Unknown model type. If you are using a model without a configuration file, Ensure that your model name includes 'bs_roformer', 'bsroformer', 'mel_band_roformer', or 'melbandroformer'. Otherwise, you can manually place the model configuration file into 'tools/uvr5/uvr5w_weights' and ensure that the configuration file is named as '.yaml' then try it again." - ) - self.config = self.get_default_config() - else: - # if there is a configuration file - self.config = self.get_config(config_path) - if self.model_type is None: - # if model_type is still None, second try, get model_type from the configuration file - if "freqs_per_bands" in self.config["model"]: - # if freqs_per_bands in config, it's a bs_roformer model - self.model_type = "bs_roformer" - else: - # else it's a mel_band_roformer model - self.model_type = "mel_band_roformer" - - print(i18n("检测到模型类型：%s") % self.model_type) - model = self.get_model_from_config() - state_dict = torch.load(model_path, map_location="cpu") - model.load_state_dict(state_dict) - - if is_half == False: - self.model = model.to(device) - else: - self.model = model.half().to(device) - - def _path_audio_(self, input, others_root, vocal_root, format, is_hp3=False): - self.run_folder(input, vocal_root, others_root, format) diff --git a/tools/uvr5/lib/lib_v5/layers_123821KB.py b/tools/uvr5/lib/lib_v5/layers_123821KB.py deleted file mode 100644 index 2b9101e..0000000 --- a/tools/uvr5/lib/lib_v5/layers_123821KB.py +++ /dev/null @@ -1,106 +0,0 @@ -import torch -import torch.nn.functional as F -from torch import nn - -from . import spec_utils - - -class Conv2DBNActiv(nn.Module): - def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU): - super(Conv2DBNActiv, self).__init__() - self.conv = nn.Sequential( - nn.Conv2d( - nin, - nout, - kernel_size=ksize, - stride=stride, - padding=pad, - dilation=dilation, - bias=False, - ), - nn.BatchNorm2d(nout), - activ(), - ) - - def __call__(self, x): - return self.conv(x) - - -class SeperableConv2DBNActiv(nn.Module): - def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU): - super(SeperableConv2DBNActiv, self).__init__() - self.conv = nn.Sequential( - nn.Conv2d( - nin, - nin, - kernel_size=ksize, - stride=stride, - padding=pad, - dilation=dilation, - groups=nin, - bias=False, - ), - nn.Conv2d(nin, nout, kernel_size=1, bias=False), - nn.BatchNorm2d(nout), - activ(), - ) - - def __call__(self, x): - return self.conv(x) - - -class Encoder(nn.Module): - def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU): - super(Encoder, self).__init__() - self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ) - self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ) - - def __call__(self, x): - skip = self.conv1(x) - h = self.conv2(skip) - - return h, skip - - -class Decoder(nn.Module): - def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False): - super(Decoder, self).__init__() - self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ) - self.dropout = nn.Dropout2d(0.1) if dropout else None - - def __call__(self, x, skip=None): - x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True) - if skip is not None: - skip = spec_utils.crop_center(skip, x) - x = torch.cat([x, skip], dim=1) - h = self.conv(x) - - if self.dropout is not None: - h = self.dropout(h) - - return h - - -class ASPPModule(nn.Module): - def __init__(self, nin, nout, dilations=(4, 8, 16), activ=nn.ReLU): - super(ASPPModule, self).__init__() - self.conv1 = nn.Sequential( - nn.AdaptiveAvgPool2d((1, None)), - Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ), - ) - self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ) - self.conv3 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[0], dilations[0], activ=activ) - self.conv4 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[1], dilations[1], activ=activ) - self.conv5 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ) - self.bottleneck = nn.Sequential(Conv2DBNActiv(nin * 5, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1)) - - def forward(self, x): - _, _, h, w = x.size() - feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True) - feat2 = self.conv2(x) - feat3 = self.conv3(x) - feat4 = self.conv4(x) - feat5 = self.conv5(x) - out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) - bottle = self.bottleneck(out) - return bottle diff --git a/tools/uvr5/lib/lib_v5/layers_new.py b/tools/uvr5/lib/lib_v5/layers_new.py deleted file mode 100644 index 7d7005c..0000000 --- a/tools/uvr5/lib/lib_v5/layers_new.py +++ /dev/null @@ -1,111 +0,0 @@ -import torch -import torch.nn.functional as F -from torch import nn - -from . import spec_utils - - -class Conv2DBNActiv(nn.Module): - def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU): - super(Conv2DBNActiv, self).__init__() - self.conv = nn.Sequential( - nn.Conv2d( - nin, - nout, - kernel_size=ksize, - stride=stride, - padding=pad, - dilation=dilation, - bias=False, - ), - nn.BatchNorm2d(nout), - activ(), - ) - - def __call__(self, x): - return self.conv(x) - - -class Encoder(nn.Module): - def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU): - super(Encoder, self).__init__() - self.conv1 = Conv2DBNActiv(nin, nout, ksize, stride, pad, activ=activ) - self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ) - - def __call__(self, x): - h = self.conv1(x) - h = self.conv2(h) - - return h - - -class Decoder(nn.Module): - def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False): - super(Decoder, self).__init__() - self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ) - # self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ) - self.dropout = nn.Dropout2d(0.1) if dropout else None - - def __call__(self, x, skip=None): - x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True) - - if skip is not None: - skip = spec_utils.crop_center(skip, x) - x = torch.cat([x, skip], dim=1) - - h = self.conv1(x) - # h = self.conv2(h) - - if self.dropout is not None: - h = self.dropout(h) - - return h - - -class ASPPModule(nn.Module): - def __init__(self, nin, nout, dilations=(4, 8, 12), activ=nn.ReLU, dropout=False): - super(ASPPModule, self).__init__() - self.conv1 = nn.Sequential( - nn.AdaptiveAvgPool2d((1, None)), - Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ), - ) - self.conv2 = Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ) - self.conv3 = Conv2DBNActiv(nin, nout, 3, 1, dilations[0], dilations[0], activ=activ) - self.conv4 = Conv2DBNActiv(nin, nout, 3, 1, dilations[1], dilations[1], activ=activ) - self.conv5 = Conv2DBNActiv(nin, nout, 3, 1, dilations[2], dilations[2], activ=activ) - self.bottleneck = Conv2DBNActiv(nout * 5, nout, 1, 1, 0, activ=activ) - self.dropout = nn.Dropout2d(0.1) if dropout else None - - def forward(self, x): - _, _, h, w = x.size() - feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True) - feat2 = self.conv2(x) - feat3 = self.conv3(x) - feat4 = self.conv4(x) - feat5 = self.conv5(x) - out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) - out = self.bottleneck(out) - - if self.dropout is not None: - out = self.dropout(out) - - return out - - -class LSTMModule(nn.Module): - def __init__(self, nin_conv, nin_lstm, nout_lstm): - super(LSTMModule, self).__init__() - self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0) - self.lstm = nn.LSTM(input_size=nin_lstm, hidden_size=nout_lstm // 2, bidirectional=True) - self.dense = nn.Sequential(nn.Linear(nout_lstm, nin_lstm), nn.BatchNorm1d(nin_lstm), nn.ReLU()) - - def forward(self, x): - N, _, nbins, nframes = x.size() - h = self.conv(x)[:, 0] # N, nbins, nframes - h = h.permute(2, 0, 1) # nframes, N, nbins - h, _ = self.lstm(h) - h = self.dense(h.reshape(-1, h.size()[-1])) # nframes * N, nbins - h = h.reshape(nframes, N, 1, nbins) - h = h.permute(1, 2, 3, 0) - - return h diff --git a/tools/uvr5/lib/lib_v5/model_param_init.py b/tools/uvr5/lib/lib_v5/model_param_init.py deleted file mode 100644 index 0708f1f..0000000 --- a/tools/uvr5/lib/lib_v5/model_param_init.py +++ /dev/null @@ -1,68 +0,0 @@ -import json -import pathlib -from tools.file_io import read_text - -default_param = {} -default_param["bins"] = 768 -default_param["unstable_bins"] = 9 # training only -default_param["reduction_bins"] = 762 # training only -default_param["sr"] = 44100 -default_param["pre_filter_start"] = 757 -default_param["pre_filter_stop"] = 768 -default_param["band"] = {} - - -default_param["band"][1] = { - "sr": 11025, - "hl": 128, - "n_fft": 960, - "crop_start": 0, - "crop_stop": 245, - "lpf_start": 61, # inference only - "res_type": "polyphase", -} - -default_param["band"][2] = { - "sr": 44100, - "hl": 512, - "n_fft": 1536, - "crop_start": 24, - "crop_stop": 547, - "hpf_start": 81, # inference only - "res_type": "sinc_best", -} - - -def int_keys(d): - r = {} - for k, v in d: - if k.isdigit(): - k = int(k) - r[k] = v - return r - - -class ModelParameters(object): - def __init__(self, config_path=""): - if ".pth" == pathlib.Path(config_path).suffix: - import zipfile - - with zipfile.ZipFile(config_path, "r") as zip: - self.param = json.loads(zip.read("param.json"), object_pairs_hook=int_keys) - elif ".json" == pathlib.Path(config_path).suffix: - self.param = json.loads( - read_text(config_path), object_pairs_hook=int_keys - ) - else: - self.param = default_param - - for k in [ - "mid_side", - "mid_side_b", - "mid_side_b2", - "stereo_w", - "stereo_n", - "reverse", - ]: - if k not in self.param: - self.param[k] = False diff --git a/tools/uvr5/lib/lib_v5/modelparams/4band_v2.json b/tools/uvr5/lib/lib_v5/modelparams/4band_v2.json deleted file mode 100644 index 33281a0..0000000 --- a/tools/uvr5/lib/lib_v5/modelparams/4band_v2.json +++ /dev/null @@ -1,54 +0,0 @@ -{ - "bins": 672, - "unstable_bins": 8, - "reduction_bins": 637, - "band": { - "1": { - "sr": 7350, - "hl": 80, - "n_fft": 640, - "crop_start": 0, - "crop_stop": 85, - "lpf_start": 25, - "lpf_stop": 53, - "res_type": "polyphase" - }, - "2": { - "sr": 7350, - "hl": 80, - "n_fft": 320, - "crop_start": 4, - "crop_stop": 87, - "hpf_start": 25, - "hpf_stop": 12, - "lpf_start": 31, - "lpf_stop": 62, - "res_type": "polyphase" - }, - "3": { - "sr": 14700, - "hl": 160, - "n_fft": 512, - "crop_start": 17, - "crop_stop": 216, - "hpf_start": 48, - "hpf_stop": 24, - "lpf_start": 139, - "lpf_stop": 210, - "res_type": "polyphase" - }, - "4": { - "sr": 44100, - "hl": 480, - "n_fft": 960, - "crop_start": 78, - "crop_stop": 383, - "hpf_start": 130, - "hpf_stop": 86, - "res_type": "kaiser_fast" - } - }, - "sr": 44100, - "pre_filter_start": 668, - "pre_filter_stop": 672 -} \ No newline at end of file diff --git a/tools/uvr5/lib/lib_v5/modelparams/4band_v3.json b/tools/uvr5/lib/lib_v5/modelparams/4band_v3.json deleted file mode 100644 index 2a73bc9..0000000 --- a/tools/uvr5/lib/lib_v5/modelparams/4band_v3.json +++ /dev/null @@ -1,54 +0,0 @@ -{ - "bins": 672, - "unstable_bins": 8, - "reduction_bins": 530, - "band": { - "1": { - "sr": 7350, - "hl": 80, - "n_fft": 640, - "crop_start": 0, - "crop_stop": 85, - "lpf_start": 25, - "lpf_stop": 53, - "res_type": "polyphase" - }, - "2": { - "sr": 7350, - "hl": 80, - "n_fft": 320, - "crop_start": 4, - "crop_stop": 87, - "hpf_start": 25, - "hpf_stop": 12, - "lpf_start": 31, - "lpf_stop": 62, - "res_type": "polyphase" - }, - "3": { - "sr": 14700, - "hl": 160, - "n_fft": 512, - "crop_start": 17, - "crop_stop": 216, - "hpf_start": 48, - "hpf_stop": 24, - "lpf_start": 139, - "lpf_stop": 210, - "res_type": "polyphase" - }, - "4": { - "sr": 44100, - "hl": 480, - "n_fft": 960, - "crop_start": 78, - "crop_stop": 383, - "hpf_start": 130, - "hpf_stop": 86, - "res_type": "kaiser_fast" - } - }, - "sr": 44100, - "pre_filter_start": 668, - "pre_filter_stop": 672 -} \ No newline at end of file diff --git a/tools/uvr5/lib/lib_v5/nets_61968KB.py b/tools/uvr5/lib/lib_v5/nets_61968KB.py deleted file mode 100644 index 167d4cb..0000000 --- a/tools/uvr5/lib/lib_v5/nets_61968KB.py +++ /dev/null @@ -1,122 +0,0 @@ -import torch -import torch.nn.functional as F -from torch import nn - -from . import layers_123821KB as layers - - -class BaseASPPNet(nn.Module): - def __init__(self, nin, ch, dilations=(4, 8, 16)): - super(BaseASPPNet, self).__init__() - self.enc1 = layers.Encoder(nin, ch, 3, 2, 1) - self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1) - self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1) - self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1) - - self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations) - - self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1) - self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1) - self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1) - self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1) - - def __call__(self, x): - h, e1 = self.enc1(x) - h, e2 = self.enc2(h) - h, e3 = self.enc3(h) - h, e4 = self.enc4(h) - - h = self.aspp(h) - - h = self.dec4(h, e4) - h = self.dec3(h, e3) - h = self.dec2(h, e2) - h = self.dec1(h, e1) - - return h - - -class CascadedASPPNet(nn.Module): - def __init__(self, n_fft): - super(CascadedASPPNet, self).__init__() - self.stg1_low_band_net = BaseASPPNet(2, 32) - self.stg1_high_band_net = BaseASPPNet(2, 32) - - self.stg2_bridge = layers.Conv2DBNActiv(34, 16, 1, 1, 0) - self.stg2_full_band_net = BaseASPPNet(16, 32) - - self.stg3_bridge = layers.Conv2DBNActiv(66, 32, 1, 1, 0) - self.stg3_full_band_net = BaseASPPNet(32, 64) - - self.out = nn.Conv2d(64, 2, 1, bias=False) - self.aux1_out = nn.Conv2d(32, 2, 1, bias=False) - self.aux2_out = nn.Conv2d(32, 2, 1, bias=False) - - self.max_bin = n_fft // 2 - self.output_bin = n_fft // 2 + 1 - - self.offset = 128 - - def forward(self, x, aggressiveness=None): - mix = x.detach() - x = x.clone() - - x = x[:, :, : self.max_bin] - - bandw = x.size()[2] // 2 - aux1 = torch.cat( - [ - self.stg1_low_band_net(x[:, :, :bandw]), - self.stg1_high_band_net(x[:, :, bandw:]), - ], - dim=2, - ) - - h = torch.cat([x, aux1], dim=1) - aux2 = self.stg2_full_band_net(self.stg2_bridge(h)) - - h = torch.cat([x, aux1, aux2], dim=1) - h = self.stg3_full_band_net(self.stg3_bridge(h)) - - mask = torch.sigmoid(self.out(h)) - mask = F.pad( - input=mask, - pad=(0, 0, 0, self.output_bin - mask.size()[2]), - mode="replicate", - ) - - if self.training: - aux1 = torch.sigmoid(self.aux1_out(aux1)) - aux1 = F.pad( - input=aux1, - pad=(0, 0, 0, self.output_bin - aux1.size()[2]), - mode="replicate", - ) - aux2 = torch.sigmoid(self.aux2_out(aux2)) - aux2 = F.pad( - input=aux2, - pad=(0, 0, 0, self.output_bin - aux2.size()[2]), - mode="replicate", - ) - return mask * mix, aux1 * mix, aux2 * mix - else: - if aggressiveness: - mask[:, :, : aggressiveness["split_bin"]] = torch.pow( - mask[:, :, : aggressiveness["split_bin"]], - 1 + aggressiveness["value"] / 3, - ) - mask[:, :, aggressiveness["split_bin"] :] = torch.pow( - mask[:, :, aggressiveness["split_bin"] :], - 1 + aggressiveness["value"], - ) - - return mask * mix - - def predict(self, x_mag, aggressiveness=None): - h = self.forward(x_mag, aggressiveness) - - if self.offset > 0: - h = h[:, :, :, self.offset : -self.offset] - assert h.size()[3] > 0 - - return h diff --git a/tools/uvr5/lib/lib_v5/nets_new.py b/tools/uvr5/lib/lib_v5/nets_new.py deleted file mode 100644 index ba1a559..0000000 --- a/tools/uvr5/lib/lib_v5/nets_new.py +++ /dev/null @@ -1,125 +0,0 @@ -import torch -import torch.nn.functional as F -from torch import nn - -from . import layers_new - - -class BaseNet(nn.Module): - def __init__(self, nin, nout, nin_lstm, nout_lstm, dilations=((4, 2), (8, 4), (12, 6))): - super(BaseNet, self).__init__() - self.enc1 = layers_new.Conv2DBNActiv(nin, nout, 3, 1, 1) - self.enc2 = layers_new.Encoder(nout, nout * 2, 3, 2, 1) - self.enc3 = layers_new.Encoder(nout * 2, nout * 4, 3, 2, 1) - self.enc4 = layers_new.Encoder(nout * 4, nout * 6, 3, 2, 1) - self.enc5 = layers_new.Encoder(nout * 6, nout * 8, 3, 2, 1) - - self.aspp = layers_new.ASPPModule(nout * 8, nout * 8, dilations, dropout=True) - - self.dec4 = layers_new.Decoder(nout * (6 + 8), nout * 6, 3, 1, 1) - self.dec3 = layers_new.Decoder(nout * (4 + 6), nout * 4, 3, 1, 1) - self.dec2 = layers_new.Decoder(nout * (2 + 4), nout * 2, 3, 1, 1) - self.lstm_dec2 = layers_new.LSTMModule(nout * 2, nin_lstm, nout_lstm) - self.dec1 = layers_new.Decoder(nout * (1 + 2) + 1, nout * 1, 3, 1, 1) - - def __call__(self, x): - e1 = self.enc1(x) - e2 = self.enc2(e1) - e3 = self.enc3(e2) - e4 = self.enc4(e3) - e5 = self.enc5(e4) - - h = self.aspp(e5) - - h = self.dec4(h, e4) - h = self.dec3(h, e3) - h = self.dec2(h, e2) - h = torch.cat([h, self.lstm_dec2(h)], dim=1) - h = self.dec1(h, e1) - - return h - - -class CascadedNet(nn.Module): - def __init__(self, n_fft, nout=32, nout_lstm=128): - super(CascadedNet, self).__init__() - - self.max_bin = n_fft // 2 - self.output_bin = n_fft // 2 + 1 - self.nin_lstm = self.max_bin // 2 - self.offset = 64 - - self.stg1_low_band_net = nn.Sequential( - BaseNet(2, nout // 2, self.nin_lstm // 2, nout_lstm), - layers_new.Conv2DBNActiv(nout // 2, nout // 4, 1, 1, 0), - ) - - self.stg1_high_band_net = BaseNet(2, nout // 4, self.nin_lstm // 2, nout_lstm // 2) - - self.stg2_low_band_net = nn.Sequential( - BaseNet(nout // 4 + 2, nout, self.nin_lstm // 2, nout_lstm), - layers_new.Conv2DBNActiv(nout, nout // 2, 1, 1, 0), - ) - self.stg2_high_band_net = BaseNet(nout // 4 + 2, nout // 2, self.nin_lstm // 2, nout_lstm // 2) - - self.stg3_full_band_net = BaseNet(3 * nout // 4 + 2, nout, self.nin_lstm, nout_lstm) - - self.out = nn.Conv2d(nout, 2, 1, bias=False) - self.aux_out = nn.Conv2d(3 * nout // 4, 2, 1, bias=False) - - def forward(self, x): - x = x[:, :, : self.max_bin] - - bandw = x.size()[2] // 2 - l1_in = x[:, :, :bandw] - h1_in = x[:, :, bandw:] - l1 = self.stg1_low_band_net(l1_in) - h1 = self.stg1_high_band_net(h1_in) - aux1 = torch.cat([l1, h1], dim=2) - - l2_in = torch.cat([l1_in, l1], dim=1) - h2_in = torch.cat([h1_in, h1], dim=1) - l2 = self.stg2_low_band_net(l2_in) - h2 = self.stg2_high_band_net(h2_in) - aux2 = torch.cat([l2, h2], dim=2) - - f3_in = torch.cat([x, aux1, aux2], dim=1) - f3 = self.stg3_full_band_net(f3_in) - - mask = torch.sigmoid(self.out(f3)) - mask = F.pad( - input=mask, - pad=(0, 0, 0, self.output_bin - mask.size()[2]), - mode="replicate", - ) - - if self.training: - aux = torch.cat([aux1, aux2], dim=1) - aux = torch.sigmoid(self.aux_out(aux)) - aux = F.pad( - input=aux, - pad=(0, 0, 0, self.output_bin - aux.size()[2]), - mode="replicate", - ) - return mask, aux - else: - return mask - - def predict_mask(self, x): - mask = self.forward(x) - - if self.offset > 0: - mask = mask[:, :, :, self.offset : -self.offset] - assert mask.size()[3] > 0 - - return mask - - def predict(self, x, aggressiveness=None): - mask = self.forward(x) - pred_mag = x * mask - - if self.offset > 0: - pred_mag = pred_mag[:, :, :, self.offset : -self.offset] - assert pred_mag.size()[3] > 0 - - return pred_mag diff --git a/tools/uvr5/lib/lib_v5/spec_utils.py b/tools/uvr5/lib/lib_v5/spec_utils.py deleted file mode 100644 index 7bc2dd2..0000000 --- a/tools/uvr5/lib/lib_v5/spec_utils.py +++ /dev/null @@ -1,445 +0,0 @@ -import math - -import librosa -import numpy as np -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): - h1_shape = h1.size() - h2_shape = h2.size() - - if h1_shape[3] == h2_shape[3]: - return h1 - elif h1_shape[3] < h2_shape[3]: - raise ValueError("h1_shape[3] must be greater than h2_shape[3]") - - # s_freq = (h2_shape[2] - h1_shape[2]) // 2 - # e_freq = s_freq + h1_shape[2] - s_time = (h1_shape[3] - h2_shape[3]) // 2 - e_time = s_time + h2_shape[3] - h1 = h1[:, :, :, s_time:e_time] - - return h1 - - -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: - 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]) - - def run_thread(**kwargs): - global spec_left - spec_left = librosa.stft(**kwargs) - - thread = threading.Thread( - target=run_thread, - kwargs={"y": wave_left, "n_fft": n_fft, "hop_length": hop_length}, - ) - thread.start() - spec_right = librosa.stft(wave_right, n_fft=n_fft, hop_length=hop_length) - thread.join() - - spec = np.asfortranarray([spec_left, spec_right]) - - return spec - - -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"]) - - for d in range(1, bands_n + 1): - h = mp.param["band"][d]["crop_stop"] - mp.param["band"][d]["crop_start"] - spec_c[:, offset : offset + h, :l] = specs[d][ - :, mp.param["band"][d]["crop_start"] : mp.param["band"][d]["crop_stop"], :l - ] - offset += h - - if offset > mp.param["bins"]: - raise ValueError("Too much bins") - - # lowpass fiter - if mp.param["pre_filter_start"] > 0: # and mp.param['band'][bands_n]['res_type'] in ['scipy', 'polyphase']: - if bands_n == 1: - spec_c = fft_lp_filter(spec_c, mp.param["pre_filter_start"], mp.param["pre_filter_stop"]) - else: - gp = 1 - for b in range(mp.param["pre_filter_start"] + 1, mp.param["pre_filter_stop"]): - g = math.pow(10, -(b - mp.param["pre_filter_start"]) * (3.5 - gp) / 20.0) - gp = g - spec_c[:, b, :] *= g - - if torch.is_tensor(spec_c): - return spec_c.contiguous() - return np.asfortranarray(spec_c) - - -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] - starts = np.insert(idx[np.where(np.diff(idx) != 1)[0] + 1], 0, idx[0]) - ends = np.append(idx[np.where(np.diff(idx) != 1)[0]], idx[-1]) - uninformative = np.where(ends - starts > min_range)[0] - if len(uninformative) > 0: - starts = starts[uninformative] - ends = ends[uninformative] - old_e = None - for s, e in zip(starts, ends): - if old_e is not None and s - old_e < fade_size: - s = old_e - fade_size * 2 - - if s != 0: - weight = np.linspace(0, 1, fade_size) - mag[:, :, s : s + fade_size] += weight * ref[:, :, s : s + fade_size] - else: - s -= fade_size - - if e != mag.shape[2]: - weight = np.linspace(1, 0, fade_size) - 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 - - -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]) - - wave_left = librosa.istft(spec_left, hop_length=hop_length) - wave_right = librosa.istft(spec_right, hop_length=hop_length) - - 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"]) - offset = 0 - - for d in range(1, bands_n + 1): - bp = mp.param["band"][d] - 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, :] - - offset += h - if d == bands_n: # higher - if extra_bins_h: # if --high_end_process bypass - max_bin = bp["n_fft"] // 2 - spec_s[:, max_bin - extra_bins_h : max_bin, :] = extra_bins[:, :extra_bins_h, :] - if bp["hpf_start"] > 0: - spec_s = fft_hp_filter(spec_s, bp["hpf_start"], bp["hpf_stop"] - 1) - if bands_n == 1: - wave = spectrogram_to_wave( - spec_s, - bp["hl"], - mp.param["mid_side"], - mp.param["mid_side_b2"], - mp.param["reverse"], - ) - else: - 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"]) - band_wave = spectrogram_to_wave( - spec_s, - bp["hl"], - mp.param["mid_side"], - mp.param["mid_side_b2"], - mp.param["reverse"], - ) - 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 = wave + spectrogram_to_wave( - spec_s, - bp["hl"], - mp.param["mid_side"], - mp.param["mid_side_b2"], - mp.param["reverse"], - ) - 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.transpose(0, 1) if torch.is_tensor(wave) else wave.T - - -def fft_lp_filter(spec, bin_start, bin_stop): - g = 1.0 - for b in range(bin_start, bin_stop): - g -= 1 / (bin_stop - bin_start) - spec[:, b, :] = g * spec[:, b, :] - - spec[:, bin_stop:, :] *= 0 - - return spec - - -def fft_hp_filter(spec, bin_start, bin_stop): - g = 1.0 - for b in range(bin_start, bin_stop, -1): - g -= 1 / (bin_start - bin_stop) - spec[:, b, :] = g * spec[:, b, :] - - spec[:, 0 : bin_stop + 1, :] *= 0 - - return spec - - -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( - spec_m[ - :, - mp.param["pre_filter_start"] - 10 - input_high_end.shape[1] : mp.param["pre_filter_start"] - 10, - :, - ] - ), - 1, - ) - mirror = mirror * np.exp(1.0j * np.angle(input_high_end)) - - return np.where(np.abs(input_high_end) <= np.abs(mirror), input_high_end, mirror) - - if "mirroring2" == a: - mirror = np.flip( - np.abs( - spec_m[ - :, - mp.param["pre_filter_start"] - 10 - input_high_end.shape[1] : mp.param["pre_filter_start"] - 10, - :, - ] - ), - 1, - ) - mi = np.multiply(mirror, input_high_end * 1.7) - - return np.where(np.abs(input_high_end) <= np.abs(mi), input_high_end, mi) diff --git a/tools/uvr5/lib/utils.py b/tools/uvr5/lib/utils.py deleted file mode 100644 index ba8c458..0000000 --- a/tools/uvr5/lib/utils.py +++ /dev/null @@ -1,198 +0,0 @@ -import numpy as np -import torch -import torch.nn.functional as F -from tools.cuda_graph import clear_cuda_graph_cache, run_cuda_graph -from tqdm import tqdm - - -def make_padding(width, cropsize, offset): - left = offset - roi_size = cropsize - left * 2 - if roi_size == 0: - roi_size = cropsize - right = roi_size - (width % roi_size) + left - - 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(): - preds = [] - - iterations = [n_window] - - total_iterations = sum(iterations) - for i in tqdm(range(n_window)): - start = i * roi_size - X_mag_window = X_mag_pad[None, :, :, start : start + data["window_size"]] - X_mag_window = torch.from_numpy(X_mag_window) - if is_half: - X_mag_window = X_mag_window.half() - X_mag_window = X_mag_window.to(device) - - pred = run_cuda_graph( - model, - "uvr-vr-%s" % repr(aggressiveness), - lambda window: model.predict(window, aggressiveness), - X_mag_window, - ) - - pred = pred.detach().cpu().numpy() - preds.append(pred[0]) - - pred = np.concatenate(preds, axis=2) - return pred - - def preprocess(X_spec): - X_mag = np.abs(X_spec) - X_phase = np.angle(X_spec) - - return X_mag, X_phase - - X_mag, X_phase = preprocess(X_spec) - - coef = X_mag.max() - 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)) - - X_mag_pad = np.pad(X_mag_pre, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant") - - if list(model.state_dict().values())[0].dtype == torch.float16: - is_half = True - else: - is_half = False - pred = _execute(X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half) - pred = pred[:, :, :n_frame] - - if data["tta"]: - pad_l += roi_size // 2 - pad_r += roi_size // 2 - n_window += 1 - - X_mag_pad = np.pad(X_mag_pre, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant") - - pred_tta = _execute(X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half) - pred_tta = pred_tta[:, :, roi_size // 2 :] - pred_tta = pred_tta[:, :, :n_frame] - - return (pred + pred_tta) * 0.5 * coef, X_mag, np.exp(1.0j * X_phase) - else: - return pred * coef, X_mag, np.exp(1.0j * X_phase) diff --git a/tools/uvr5/mdxnet.py b/tools/uvr5/mdxnet.py deleted file mode 100644 index 4ac7301..0000000 --- a/tools/uvr5/mdxnet.py +++ /dev/null @@ -1,446 +0,0 @@ -import os -import logging -import sysconfig - -logger = logging.getLogger(__name__) - -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 = [] - - -def _configure_ort_cuda_dll_paths(): - """Expose pip-installed CUDA 11/cuDNN 8 DLLs to ONNX Runtime on Windows.""" - if os.name != "nt": - return - - site_packages = os.path.normpath(sysconfig.get_paths()["purelib"]) - nvidia_root = os.path.join(site_packages, "nvidia") - dll_dirs = [ - os.path.join(nvidia_root, "cuda_runtime", "bin"), - os.path.join(nvidia_root, "cublas", "bin"), - os.path.join(nvidia_root, "cufft", "bin"), - os.path.join(nvidia_root, "cudnn", "bin"), - os.path.join(nvidia_root, "cuda_nvrtc", "bin"), - os.path.join(os.path.dirname(torch.__file__), "lib"), - ] - dll_dirs = [path for path in dll_dirs if os.path.isdir(path)] - if not dll_dirs: - return - - current_path = os.environ.get("PATH", "") - current_dirs = [path for path in current_path.split(os.pathsep) if path] - known_dirs = {os.path.normcase(os.path.normpath(path)) for path in current_dirs} - prepend_dirs = [] - for path in dll_dirs: - normalized = os.path.normcase(os.path.normpath(path)) - if normalized not in known_dirs: - prepend_dirs.append(path) - known_dirs.add(normalized) - if prepend_dirs: - os.environ["PATH"] = os.pathsep.join(prepend_dirs + current_dirs) - - # Python 3.8+ restricts DLL lookup for extension modules. Keep the handles - # alive for the process lifetime in addition to updating PATH. - if hasattr(os, "add_dll_directory"): - for path in dll_dirs: - try: - _ORT_CUDA_DLL_HANDLES.append(os.add_dll_directory(path)) - except OSError: - logger.warning("Unable to add ONNX Runtime DLL directory: %s", path) - - -_configure_ort_cuda_dll_paths() - -cpu = torch.device("cpu") - - -class ConvTDFNetTrim: - 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 - self.hop = hop - 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.dim_c = 4 - 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]) - x = torch.stft( - x, - n_fft=self.n_fft, - hop_length=self.hop, - window=self.window, - center=True, - return_complex=True, - ) - x = torch.view_as_real(x) - x = x.permute([0, 3, 1, 2]) - 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 = self.freq_pad.expand(x.shape[0], -1, -1, -1) - x = torch.cat([x, freq_pad], -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() - x = torch.view_as_complex(x) - x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True) - return x.reshape([-1, c, self.chunk_size]) - - -def get_models(device, dim_f, dim_t, n_fft): - return ConvTDFNetTrim( - device=device, - dim_f=dim_f, - dim_t=dim_t, - n_fft=n_fft, - ) - - -class Predictor: - def __init__(self, args): - 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 requested_providers - and "CUDAExecutionProvider" not in available_providers - ): - raise RuntimeError( - "CUDAExecutionProvider is required for the FoxJoy ONNX model, " - "but the installed ONNX Runtime does not provide it. Install " - "the matching CUDA ONNX Runtime dependencies with this " - "project's runtime Python." - ) - if ( - "DmlExecutionProvider" in requested_providers - and "DmlExecutionProvider" not in available_providers - ): - raise RuntimeError( - "DmlExecutionProvider is required for the FoxJoy ONNX model, " - "but the installed ONNX Runtime does not provide it. Install " - "requirments_cpu_py312.txt with this project's runtime Python." - ) - self.args = args - 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, "vocals.onnx"), - providers=args.providers, - ) - active_providers = self.model.get_providers() - logger.info("ONNX Runtime active providers: %s", active_providers) - if ( - "CUDAExecutionProvider" in requested_providers - and ( - not active_providers - or active_providers[0] != "CUDAExecutionProvider" - ) - ): - raise RuntimeError( - "The FoxJoy ONNX model did not activate CUDAExecutionProvider; " - "check the CUDA 11/cuDNN 8 DLL installation." - ) - if ( - "DmlExecutionProvider" in requested_providers - and ( - not active_providers - or active_providers[0] != "DmlExecutionProvider" - ) - ): - raise RuntimeError( - "The FoxJoy ONNX model did not activate DmlExecutionProvider; " - "check the ONNX Runtime DirectML installation." - ) - 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] - margin = self.args.margin - chunk_size = self.args.chunks * 44100 - assert not margin == 0, "margin cannot be zero!" - if margin > chunk_size: - margin = chunk_size - - segmented_mix = {} - - if self.args.chunks == 0 or samples < chunk_size: - chunk_size = samples - - counter = -1 - for skip in range(0, samples, chunk_size): - counter += 1 - - s_margin = 0 if counter == 0 else margin - end = min(skip + chunk_size + margin, samples) - - start = skip - s_margin - - 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 - - sources = self.demix_base(segmented_mix, margin_size=margin) - """ - mix:(2,big_sample) - segmented_mix:offset->(2,small_sample) - sources:(1,2,big_sample) - """ - return sources - - def demix_base(self, mixes, margin_size): - chunked_sources = [] - progress_bar = tqdm(total=len(mixes)) - progress_bar.set_description("Processing") - for mix in mixes: - cmix = mixes[mix] - sources = [] - n_sample = cmix.shape[1] - model = self.model_ - trim = model.n_fft // 2 - gen_size = model.chunk_size - 2 * trim - pad = gen_size - n_sample % gen_size - 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 = mix_p[:, i : i + model.chunk_size] - if not torch.is_tensor(waves): - waves = np.array(waves) - mix_waves.append(waves) - i += gen_size - 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_numpy})[0] * 0.5 - + _ort.run(None, {"input": spek_numpy})[0] * 0.5 - ) - tar_waves = model.istft(torch.from_numpy(spec_pred)) - else: - 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 - sources.append(tar_signal[:, start:end]) - - progress_bar.update(1) - - chunked_sources.append(sources) - _sources = np.concatenate(chunked_sources, axis=-1) - # del self.model - progress_bar.close() - return _sources - - def prediction(self, m, vocal_root, others_root, format): - os.makedirs(vocal_root, exist_ok=True) - os.makedirs(others_root, exist_ok=True) - basename = os.path.basename(m) - 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 = 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) - else: - path_vocal = "%s/%s_main_vocal.wav" % (vocal_root, basename) - path_other = "%s/%s_others.wav" % (others_root, basename) - sf.write(path_vocal, mix - opt, rate) - sf.write(path_other, opt, rate) - opt_path_vocal = path_vocal[:-4] + ".%s" % format - opt_path_other = path_other[:-4] + ".%s" % format - if os.path.exists(path_vocal): - os.system('ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path_vocal, opt_path_vocal)) - if os.path.exists(opt_path_vocal): - try: - os.remove(path_vocal) - except: - pass - if os.path.exists(path_other): - os.system('ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path_other, opt_path_other)) - if os.path.exists(opt_path_other): - try: - os.remove(path_other) - except: - pass - - -class MDXNetDereverb: - 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.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) - - def _path_audio_(self, input, others_root, vocal_root, format, is_hp3=False): - self.pred.prediction(input, vocal_root, others_root, format) diff --git a/tools/uvr5/rotary_embedding_torch/__init__.py b/tools/uvr5/rotary_embedding_torch/__init__.py deleted file mode 100644 index b1fbdc2..0000000 --- a/tools/uvr5/rotary_embedding_torch/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .rotary_embedding_torch import ( - apply_rotary_emb, - RotaryEmbedding, - apply_learned_rotations, - broadcat -) diff --git a/tools/uvr5/rotary_embedding_torch/rotary_embedding_torch.py b/tools/uvr5/rotary_embedding_torch/rotary_embedding_torch.py deleted file mode 100644 index ec67505..0000000 --- a/tools/uvr5/rotary_embedding_torch/rotary_embedding_torch.py +++ /dev/null @@ -1,186 +0,0 @@ -from __future__ import annotations -from math import pi, log -import warnings - -warnings.filterwarnings( - "ignore", - message="`torch.cuda.amp.autocast.*is deprecated.*", - category=FutureWarning, -) - -import torch -from torch.nn import Module, ModuleList -from torch.cuda.amp import autocast -from torch import nn, einsum, broadcast_tensors, Tensor -from einops import rearrange, repeat -from typing import Literal - -def exists(val): - return val is not None - -def default(val, d): - return val if exists(val) else d - -def broadcat(tensors, dim=-1): - broadcasted_tensors = broadcast_tensors(*tensors) - return torch.cat(broadcasted_tensors, dim=dim) - -def rotate_half(x): - x = rearrange(x, '... (d r) -> ... d r', r=2) - (x1, x2) = x.unbind(dim=-1) - x = torch.stack((-x2, x1), dim=-1) - return rearrange(x, '... d r -> ... (d r)') - -@autocast(enabled=False) -def apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2): - dtype = t.dtype - if t.ndim == 3: - seq_len = t.shape[seq_dim] - freqs = freqs[-seq_len:] - rot_dim = freqs.shape[-1] - end_index = start_index + rot_dim - assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}' - (t_left, t, t_right) = (t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]) - t = t * freqs.cos() * scale + rotate_half(t) * freqs.sin() * scale - if t.device.type == 'privateuseone': - # DirectML rejects concatenation when one of the slices has a zero - # length. Rotary embeddings normally cover the complete head, so both - # edge slices are empty; omitting them is mathematically identical. - parts = tuple(part for part in (t_left, t, t_right) if part.shape[-1] > 0) - out = parts[0] if len(parts) == 1 else torch.cat(parts, dim=-1) - else: - out = torch.cat((t_left, t, t_right), dim=-1) - return out.type(dtype) - -def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None): - if exists(freq_ranges): - rotations = einsum('..., f -> ... f', rotations, freq_ranges) - rotations = rearrange(rotations, '... r f -> ... (r f)') - rotations = repeat(rotations, '... n -> ... (n r)', r=2) - return apply_rotary_emb(rotations, t, start_index=start_index) - -class RotaryEmbedding(Module): - - def __init__(self, dim, custom_freqs=None, freqs_for='lang', theta=10000, max_freq=10, num_freqs=1, learned_freq=False, use_xpos=False, xpos_scale_base=512, interpolate_factor=1.0, theta_rescale_factor=1.0, seq_before_head_dim=False, cache_if_possible=True): - super().__init__() - theta *= theta_rescale_factor ** (dim / (dim - 2)) - self.freqs_for = freqs_for - if exists(custom_freqs): - freqs = custom_freqs - elif freqs_for == 'lang': - freqs = 1.0 / theta ** (torch.arange(0, dim, 2)[:dim // 2].float() / dim) - elif freqs_for == 'pixel': - freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi - elif freqs_for == 'constant': - freqs = torch.ones(num_freqs).float() - self.cache_if_possible = cache_if_possible - self.tmp_store('cached_freqs', None) - self.tmp_store('cached_scales', None) - self.freqs = nn.Parameter(freqs, requires_grad=learned_freq) - self.learned_freq = learned_freq - self.tmp_store('dummy', torch.tensor(0)) - self.seq_before_head_dim = seq_before_head_dim - self.default_seq_dim = -3 if seq_before_head_dim else -2 - assert interpolate_factor >= 1.0 - self.interpolate_factor = interpolate_factor - self.use_xpos = use_xpos - if not use_xpos: - self.tmp_store('scale', None) - return - scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) - self.scale_base = xpos_scale_base - self.tmp_store('scale', scale) - self.apply_rotary_emb = staticmethod(apply_rotary_emb) - - @property - def device(self): - return self.dummy.device - - def tmp_store(self, key, value): - self.register_buffer(key, value, persistent=False) - - def get_seq_pos(self, seq_len, device, dtype, offset=0): - return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor - - def rotate_queries_or_keys(self, t, seq_dim=None, offset=0, scale=None): - seq_dim = default(seq_dim, self.default_seq_dim) - assert not self.use_xpos or exists(scale), 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings' - (device, dtype, seq_len) = (t.device, t.dtype, t.shape[seq_dim]) - seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset) - freqs = self.forward(seq, seq_len=seq_len, offset=offset) - if seq_dim == -3: - freqs = rearrange(freqs, 'n d -> n 1 d') - return apply_rotary_emb(freqs, t, scale=default(scale, 1.0), seq_dim=seq_dim) - - def rotate_queries_with_cached_keys(self, q, k, seq_dim=None, offset=0): - (dtype, device, seq_dim) = (q.dtype, q.device, default(seq_dim, self.default_seq_dim)) - (q_len, k_len) = (q.shape[seq_dim], k.shape[seq_dim]) - assert q_len <= k_len - q_scale = k_scale = 1.0 - if self.use_xpos: - seq = self.get_seq_pos(k_len, dtype=dtype, device=device) - q_scale = self.get_scale(seq[-q_len:]).type(dtype) - k_scale = self.get_scale(seq).type(dtype) - rotated_q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset) - rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale ** (-1)) - rotated_q = rotated_q.type(q.dtype) - rotated_k = rotated_k.type(k.dtype) - return (rotated_q, rotated_k) - - def rotate_queries_and_keys(self, q, k, seq_dim=None): - seq_dim = default(seq_dim, self.default_seq_dim) - assert self.use_xpos - (device, dtype, seq_len) = (q.device, q.dtype, q.shape[seq_dim]) - seq = self.get_seq_pos(seq_len, dtype=dtype, device=device) - freqs = self.forward(seq, seq_len=seq_len) - scale = self.get_scale(seq, seq_len=seq_len).to(dtype) - if seq_dim == -3: - freqs = rearrange(freqs, 'n d -> n 1 d') - scale = rearrange(scale, 'n d -> n 1 d') - rotated_q = apply_rotary_emb(freqs, q, scale=scale, seq_dim=seq_dim) - rotated_k = apply_rotary_emb(freqs, k, scale=scale ** (-1), seq_dim=seq_dim) - rotated_q = rotated_q.type(q.dtype) - rotated_k = rotated_k.type(k.dtype) - return (rotated_q, rotated_k) - - def get_scale(self, t, seq_len=None, offset=0): - assert self.use_xpos - should_cache = self.cache_if_possible and exists(seq_len) - if should_cache and exists(self.cached_scales) and (seq_len + offset <= self.cached_scales.shape[0]): - return self.cached_scales[offset:offset + seq_len] - scale = 1.0 - if self.use_xpos: - power = (t - len(t) // 2) / self.scale_base - scale = self.scale ** rearrange(power, 'n -> n 1') - scale = torch.cat((scale, scale), dim=-1) - if should_cache: - self.tmp_store('cached_scales', scale) - return scale - - def get_axial_freqs(self, *dims): - Colon = slice(None) - all_freqs = [] - for (ind, dim) in enumerate(dims): - if self.freqs_for == 'pixel': - pos = torch.linspace(-1, 1, steps=dim, device=self.device) - else: - pos = torch.arange(dim, device=self.device) - freqs = self.forward(pos, seq_len=dim) - all_axis = [None] * len(dims) - all_axis[ind] = Colon - new_axis_slice = (Ellipsis, *all_axis, Colon) - all_freqs.append(freqs[new_axis_slice]) - all_freqs = broadcast_tensors(*all_freqs) - return torch.cat(all_freqs, dim=-1) - - @autocast(enabled=False) - def forward(self, t, seq_len=None, offset=0): - should_cache = self.cache_if_possible and (not self.learned_freq) and exists(seq_len) and (self.freqs_for != 'pixel') - if should_cache and exists(self.cached_freqs) and (offset + seq_len <= self.cached_freqs.shape[0]): - return self.cached_freqs[offset:offset + seq_len].detach() - freqs = self.freqs - freqs = einsum('..., f -> ... f', t.type(freqs.dtype), freqs) - freqs = repeat(freqs, '... n -> ... (n r)', r=2) - if should_cache: - self.tmp_store('cached_freqs', freqs.detach()) - return freqs diff --git a/tools/uvr5/vr.py b/tools/uvr5/vr.py deleted file mode 100644 index 31ccc18..0000000 --- a/tools/uvr5/vr.py +++ /dev/null @@ -1,456 +0,0 @@ -import os - -parent_directory = os.path.dirname(os.path.abspath(__file__)) -import logging - -logger = logging.getLogger(__name__) - -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 -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 - self.device = device - self.data = { - # Processing Options - "postprocess": False, - "tta": tta, - # Constants - "window_size": 512, - "agg": agg, - "high_end_process": "mirroring", - } - mp = ModelParameters("%s/lib/lib_v5/modelparams/4band_v2.json" % parent_directory) - model = Nets.CascadedASPPNet(mp.param["bins"] * 2) - cpk = torch.load(model_path, map_location="cpu") - model.load_state_dict(cpk) - model.eval() - if is_half: - model = model.half().to(device) - else: - model = model.to(device) - - self.mp = mp - self.model = model - - def _path_audio_(self, music_file, ins_root=None, vocal_root=None, format="flac", is_hp3=False): - if ins_root is None and vocal_root is None: - return "No save root." - name = os.path.basename(music_file) - if ins_root is not None: - os.makedirs(ins_root, exist_ok=True) - if vocal_root is not None: - os.makedirs(vocal_root, exist_ok=True) - aggresive_set = float(self.data["agg"] / 100) - aggressiveness = { - "value": aggresive_set, - "split_bin": self.mp.param["band"][1]["crop_stop"], - } - 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 - - if ins_root is not None: - if self.data["high_end_process"].startswith("mirroring"): - input_high_end_ = spec_utils.mirroring(self.data["high_end_process"], y_spec_m, input_high_end, self.mp) - wav_instrument = spec_utils.cmb_spectrogram_to_wave( - y_spec_m, self.mp, input_high_end_h, input_high_end_ - ) - else: - wav_instrument = spec_utils.cmb_spectrogram_to_wave(y_spec_m, self.mp) - logger.info("%s instruments done" % name) - if is_hp3 == True: - head = "vocal_" - else: - head = "instrument_" - if format in ["wav", "flac"]: - sf.write( - os.path.join( - ins_root, - head + "{}_{}.{}".format(name, self.data["agg"], format), - ), - (_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, - (_wave_for_write(wav_instrument) * 32768).astype("int16"), - self.mp.param["sr"], - ) - if os.path.exists(path): - opt_format_path = path[:-4] + ".%s" % format - cmd = 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path, opt_format_path) - print(cmd) - os.system(cmd) - if os.path.exists(opt_format_path): - try: - os.remove(path) - 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: - head = "vocal_" - 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_) - else: - wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp) - logger.info("%s vocals done" % name) - if format in ["wav", "flac"]: - sf.write( - os.path.join( - vocal_root, - head + "{}_{}.{}".format(name, self.data["agg"], format), - ), - (_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, - (_wave_for_write(wav_vocals) * 32768).astype("int16"), - self.mp.param["sr"], - ) - if os.path.exists(path): - opt_format_path = path[:-4] + ".%s" % format - cmd = 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path, opt_format_path) - print(cmd) - os.system(cmd) - if os.path.exists(opt_format_path): - try: - os.remove(path) - except: - pass - - -class AudioPreDeEcho: - def __init__(self, agg, model_path, device, is_half, tta=False): - self.model_path = model_path - self.device = device - self.data = { - # Processing Options - "postprocess": False, - "tta": tta, - # Constants - "window_size": 512, - "agg": agg, - "high_end_process": "mirroring", - } - mp = ModelParameters("%s/lib/lib_v5/modelparams/4band_v3.json" % parent_directory) - nout = 64 if "DeReverb" in model_path else 48 - model = CascadedNet(mp.param["bins"] * 2, nout) - cpk = torch.load(model_path, map_location="cpu") - model.load_state_dict(cpk) - model.eval() - if is_half: - model = model.half().to(device) - else: - model = model.to(device) - - self.mp = mp - self.model = model - - def _path_audio_( - self, music_file, vocal_root=None, ins_root=None, format="flac", is_hp3=False - ): # 3个VR模型vocal和ins是反的 - if ins_root is None and vocal_root is None: - return "No save root." - name = os.path.basename(music_file) - if ins_root is not None: - os.makedirs(ins_root, exist_ok=True) - if vocal_root is not None: - os.makedirs(vocal_root, exist_ok=True) - aggresive_set = float(self.data["agg"] / 100) - aggressiveness = { - "value": aggresive_set, - "split_bin": self.mp.param["band"][1]["crop_stop"], - } - 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"): - input_high_end_ = spec_utils.mirroring(self.data["high_end_process"], y_spec_m, input_high_end, self.mp) - wav_instrument = spec_utils.cmb_spectrogram_to_wave( - y_spec_m, self.mp, input_high_end_h, input_high_end_ - ) - else: - wav_instrument = spec_utils.cmb_spectrogram_to_wave(y_spec_m, self.mp) - logger.info("%s instruments done" % name) - if format in ["wav", "flac"]: - sf.write( - os.path.join( - ins_root, - "vocal_{}_{}.{}".format(name, self.data["agg"], format), - ), - (_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, - (_wave_for_write(wav_instrument) * 32768).astype("int16"), - self.mp.param["sr"], - ) - if os.path.exists(path): - opt_format_path = path[:-4] + ".%s" % format - cmd = 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path, opt_format_path) - print(cmd) - os.system(cmd) - if os.path.exists(opt_format_path): - try: - os.remove(path) - 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_) - else: - wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp) - logger.info("%s vocals done" % name) - if format in ["wav", "flac"]: - sf.write( - os.path.join( - vocal_root, - "instrument_{}_{}.{}".format(name, self.data["agg"], format), - ), - (_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, - (_wave_for_write(wav_vocals) * 32768).astype("int16"), - self.mp.param["sr"], - ) - if os.path.exists(path): - opt_format_path = path[:-4] + ".%s" % format - cmd = 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path, opt_format_path) - print(cmd) - os.system(cmd) - if os.path.exists(opt_format_path): - try: - os.remove(path) - except: - pass diff --git a/tools/uvr5/webui.py b/tools/uvr5/webui.py deleted file mode 100644 index cf16a22..0000000 --- a/tools/uvr5/webui.py +++ /dev/null @@ -1,111 +0,0 @@ -import logging -import os -import traceback - -import torch - -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 -from i18n.i18n import I18nAuto - - -logger = logging.getLogger(__name__) -i18n = I18nAuto() -config = Config() -weight_uvr5_root = os.getenv("weight_uvr5_root", "assets/uvr5_weights") - - -def clean_path(path): - path = path or "" - if path.endswith(("\\", "/")): - path = path[:-1] - return path.replace("/", os.sep).replace("\\", os.sep).strip(" '\n\"\u202a") - - -def uvr(model_name, inp_root, save_root_vocal, paths, save_root_ins, agg, format0): - infos = [] - try: - inp_root = clean_path(inp_root) - save_root_vocal = clean_path(save_root_vocal) - save_root_ins = clean_path(save_root_ins) - is_hp3 = "HP3" in model_name - if model_name == "onnx_dereverb_By_FoxJoy": - if config.dml: - providers = ["DmlExecutionProvider", "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, config.device) - elif "roformer" in model_name.lower(): - pre_fun = Roformer_Loader( - model_path=os.path.join(weight_uvr5_root, model_name + ".ckpt"), - config_path=os.path.join(weight_uvr5_root, model_name + ".yaml"), - device=config.device, - is_half=config.is_half, - ) - if not os.path.exists( - os.path.join(weight_uvr5_root, model_name + ".yaml") - ): - infos.append(i18n("未找到Roformer模型配置文件，正在使用内置默认配置")) - yield "\n".join(infos) - else: - func = AudioPre if "DeEcho" not in model_name else AudioPreDeEcho - pre_fun = func( - agg=int(agg), - model_path=os.path.join(weight_uvr5_root, model_name + ".pth"), - device=config.device, - is_half=config.is_half, - ) - if inp_root: - paths = [os.path.join(inp_root, name) for name in os.listdir(inp_root)] - else: - paths = [path.name for path in (paths or [])] - for path in paths: - inp_path = os.path.join(inp_root, path) - if not os.path.isfile(inp_path): - continue - try: - # 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, - save_root_vocal, - format0, - is_hp3, - ) - infos.append(i18n("%s → 成功") % os.path.basename(inp_path)) - yield "\n".join(infos) - except Exception: - infos.append( - "%s → %s\n%s" - % (os.path.basename(inp_path), i18n("失败"), traceback.format_exc()) - ) - yield "\n".join(infos) - except Exception: - infos.append("%s\n%s" % (i18n("失败"), traceback.format_exc())) - yield "\n".join(infos) - finally: - try: - if model_name == "onnx_dereverb_By_FoxJoy": - del pre_fun.pred.model - del pre_fun.pred.model_ - else: - del pre_fun.model - del pre_fun - except: - traceback.print_exc() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - logger.info("Executed torch.cuda.empty_cache()") - yield "\n".join(infos) diff --git a/webui.py b/webui.py index c5fc74a..8ebfc54 100644 --- a/webui.py +++ b/webui.py @@ -10,7 +10,7 @@ os.environ["RVC_CUDA_GRAPH"] = "1" if _offline_cuda_graph else "0" os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") os.environ.setdefault("no_proxy", "localhost, 127.0.0.1, ::1") os.environ.setdefault("weight_root", "assets/weights") -os.environ.setdefault("weight_uvr5_root", "assets/uvr5_weights") +os.environ.setdefault("weight_pymss_root", "assets/pymss_weights") os.environ.setdefault("index_root", "logs") os.environ.setdefault("outside_index_root", "assets/indices") os.environ.setdefault("rmvpe_root", "assets/rmvpe") @@ -33,7 +33,7 @@ for name in os.listdir(tmp): from configs.config import Config, GPU_INDEX, GPU_INFOS, GPU_MEMORY, IS_GPU from infer.vc.modules import VC -from tools.uvr5.webui import uvr +from tools.pymss_webui import PYMSS_MODEL_CHOICES, get_model_info, pymss_separate from tools.file_io import read_text from train.process_ckpt import ( change_info, @@ -125,7 +125,7 @@ def launch_webui_with_port_fallback(app, config): runtime_dirs = ( os.path.join(now_dir, "logs"), os.environ["weight_root"], - os.environ["weight_uvr5_root"], + os.environ["weight_pymss_root"], os.environ["index_root"], os.environ["outside_index_root"], os.environ["rmvpe_root"], @@ -173,7 +173,7 @@ class ToolButton(gr.Button, gr.components.FormComponent): weight_root = os.getenv("weight_root") -weight_uvr5_root = os.getenv("weight_uvr5_root") +weight_pymss_root = os.getenv("weight_pymss_root") outside_index_root = os.getenv("outside_index_root") def weight_names(): @@ -190,11 +190,7 @@ def refresh_weight_choices(previous_names=None, force=False): names = weight_names() -uvr5_names = [] -for name in os.listdir(weight_uvr5_root): - if name.endswith((".pth", ".ckpt")) or "onnx" in name: - uvr5_names.append(name.replace(".pth", "").replace(".ckpt", "")) -uvr5_names.sort() +pymss_names = PYMSS_MODEL_CHOICES def change_choices(): @@ -1514,11 +1510,11 @@ with gr.Blocks(title="RVC WebUI") as app: outputs=[spk_item, protect0, protect1, file_index1, file_index3], api_name="infer_change_voice", ) - with gr.TabItem(i18n("伴奏人声分离&去混响&去回声")): + with gr.TabItem(i18n("人声伴奏分离&去混响")): with gr.Group(): gr.Markdown( value=i18n( - "人声伴奏分离批量处理，使用UVR5模型。
可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。" + "人声、伴奏与混响批量处理，使用pymss/MSST模型。" ) ) with gr.Row(): @@ -1533,22 +1529,27 @@ with gr.Blocks(title="RVC WebUI") as app: ) with gr.Column(): model_choose = gr.Dropdown( - label=i18n("模型"), choices=uvr5_names - ) - agg = gr.Slider( - minimum=0, - maximum=20, - step=1, - label=i18n("人声提取激进程度"), - value=10, + label=i18n("处理方式"), + choices=pymss_names, + value=pymss_names[0], interactive=True, - visible=False, # 先不开放调整 + ) + model_info = gr.Textbox( + label=i18n("底层模型"), + value=get_model_info(pymss_names[0]), + interactive=False, + ) + model_choose.change( + get_model_info, + [model_choose], + [model_info], + queue=False, ) opt_vocal_root = gr.Textbox( - label=i18n("指定输出主人声文件夹"), value="opt" + label=i18n("主结果文件夹"), value="opt" ) opt_ins_root = gr.Textbox( - label=i18n("指定输出非主人声文件夹"), value="opt" + label=i18n("分离残余文件夹"), value="opt" ) format0 = gr.Radio( label=i18n("导出文件格式"), @@ -1559,14 +1560,13 @@ with gr.Blocks(title="RVC WebUI") as app: but2 = gr.Button(i18n("转换"), variant="primary") vc_output4 = gr.Textbox(label=i18n("输出信息")) but2.click( - uvr, + pymss_separate, [ model_choose, dir_wav_input, opt_vocal_root, wav_inputs, opt_ins_root, - agg, format0, ], [vc_output4],