initial refactor of rvc studio

.gitignore

@@ -14,7 +14,6 @@ dist/
 downloads/
 eggs/
 .eggs/
-lib/
 lib64/
 parts/
 sdist/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 RVC Studio
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

__init__.py

@@ -0,0 +1,48 @@
+from .custom_nodes.stt import AudioTranscriptionNode
+from .custom_nodes.uvr import UVR5Node
+from .custom_nodes.rvc import RVCNode
+from .custom_nodes.loaders import LoadAudio, LoadWhisperModelNode, LoadRVCModelNode, LoadHubertModel, LoadPitchExtractionParams
+from .custom_nodes.output import PreviewAudio
+from .custom_nodes.utils import AudioBatchValueNode, MergeImageBatches, MergeLatentBatches, ImageRepeatInterleavedNode, LatentRepeatInterleavedNode, MergeAudioNode
+
+# Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension
+WEB_DIRECTORY = "./web"
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "UVR5Node": UVR5Node,
+    "LoadAudio": LoadAudio,
+    "PreviewAudio": PreviewAudio,
+    "MergeAudioNode": MergeAudioNode,
+    "AudioTranscriptionNode": AudioTranscriptionNode,
+    "LoadWhisperModelNode": LoadWhisperModelNode,
+    "LoadRVCModelNode": LoadRVCModelNode,
+    "RVCNode": RVCNode,
+    "LoadHubertModel": LoadHubertModel,
+    "LoadPitchExtractionParams": LoadPitchExtractionParams,
+    "AudioBatchValueNode": AudioBatchValueNode,
+    "MergeImageBatches": MergeImageBatches,
+    "MergeLatentBatches": MergeLatentBatches,
+    "ImageRepeatInterleavedNode": ImageRepeatInterleavedNode,
+    "LatentRepeatInterleavedNode": LatentRepeatInterleavedNode
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "UVR5Node": "🌺Vocal Removal",
+    "LoadAudio": "🌺Load Audio",
+    "PreviewAudio": "🌺Preview Audio",
+    "AudioTranscriptionNode": "🌺Transcribe Audio",
+    "LoadWhisperModelNode": "🌺Load Whisper Model",
+    "LoadRVCModelNode": "🌺Load RVC Model",
+    "RVCNode": "🌺Voice Changer",
+    "LoadHubertModel": "🌺Load Hubert Model",
+    "LoadPitchExtractionParams": "🌺Load Pitch Extraction Params",
+    "MergeAudioNode": "🌺Merge Audio",
+    "AudioBatchValueNode": "🌺Audio RMS Batch Values",
+    "MergeImageBatches": "🌺Merge Image Batches",
+    "MergeLatentBatches": "🌺Merge Latent Batches",
+    "ImageRepeatInterleavedNode": "🌺Image Repeat Interleaved",
+    "LatentRepeatInterleavedNode": "🌺Latent Repeat Interleaved"
+}

config.py

@@ -0,0 +1,170 @@
+import os
+import argparse
+import sys
+import torch
+from multiprocessing import cpu_count
+
+
+def use_fp32_config():
+    for config_file in [
+        "32k.json",
+        "40k.json",
+        "48k.json",
+        "48k_v2.json",
+        "32k_v2.json",
+    ]:
+        with open(f"configs/{config_file}", "r") as f:
+            strr = f.read().replace("true", "false")
+        with open(f"configs/{config_file}", "w") as f:
+            f.write(strr)
+
+
+class Config:
+    def __init__(self):
+        self.device = "cuda:0"
+        self.is_half = True
+        self.n_cpu = 0
+        self.gpu_name = None
+        self.gpu_mem = None
+        self.has_gpu = torch.cuda.is_available()
+        (
+            self.python_cmd,
+            self.listen_port,
+            self.iscolab,
+            self.noparallel,
+            self.noautoopen,
+            self.dml,
+        ) = self.arg_parse()
+        self.instead = ""
+        self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()
+
+    @staticmethod
+    def arg_parse() -> tuple:
+        exe = sys.executable or "python"
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--port", type=int, default=7865, help="Listen port")
+        parser.add_argument("--pycmd", type=str, default=exe, help="Python command")
+        parser.add_argument("--colab", action="store_true", help="Launch in colab")
+        parser.add_argument(
+            "--noparallel", action="store_true", help="Disable parallel processing"
+        )
+        parser.add_argument(
+            "--noautoopen",
+            action="store_true",
+            help="Do not open in browser automatically",
+        )
+        parser.add_argument(
+            "--dml",
+            action="store_true",
+            help="torch_dml",
+        )
+        cmd_opts, unknown = parser.parse_known_args() # allows import to jupyter notebook
+        print(f"unknown args: {unknown}")
+
+        # cmd_opts.port = cmd_opts.port if 0 <= cmd_opts.port <= 65535 else 7865
+
+        return (
+            cmd_opts.pycmd,
+            cmd_opts.port,
+            cmd_opts.colab,
+            cmd_opts.noparallel,
+            cmd_opts.noautoopen,
+            cmd_opts.dml,
+        )
+
+    # has_mps is only available in nightly pytorch (for now) and MasOS 12.3+.
+    # check `getattr` and try it for compatibility
+    @staticmethod
+    def has_mps() -> bool:
+        if not torch.backends.mps.is_available():
+            return False
+        try:
+            torch.zeros(1).to(torch.device("mps"))
+            return True
+        except Exception:
+            return False
+
+    def device_config(self) -> tuple:
+        if torch.cuda.is_available():
+            i_device = int(self.device.split(":")[-1])
+            self.gpu_name = torch.cuda.get_device_name(i_device)
+            if (
+                ("16" in self.gpu_name and "V100" not in self.gpu_name.upper())
+                or "P40" in self.gpu_name.upper()
+                or "1060" in self.gpu_name
+                or "1070" in self.gpu_name
+                or "1080" in self.gpu_name
+            ):
+                print("Found GPU", self.gpu_name, ", force to fp32")
+                self.is_half = False
+                use_fp32_config()
+            else:
+                print("Found GPU", self.gpu_name)
+            self.gpu_mem = int(
+                torch.cuda.get_device_properties(i_device).total_memory
+                / 1024
+                / 1024
+                / 1024
+                # + 0.4
+            )
+        elif self.has_mps():
+            print("No supported Nvidia GPU found")
+            self.device = self.instead = "mps"
+            self.is_half = False
+            use_fp32_config()
+        else:
+            print("No supported Nvidia GPU found")
+            self.device = self.instead = "cpu"
+            self.is_half = False
+            use_fp32_config()
+
+        if self.n_cpu == 0:
+            self.n_cpu = cpu_count()
+
+        if self.is_half:
+            # 6G显存配置
+            x_pad = 3
+            x_query = 10
+            x_center = 60
+            x_max = 64
+        else:
+            # 5G显存配置
+            x_pad = 1
+            x_query = 6
+            x_center = 38
+            x_max = 41
+
+        if self.gpu_mem is not None and self.gpu_mem <= 4:
+            x_pad = 1
+            x_query = 5
+            x_center = 30
+            x_max = 32
+        if self.dml:
+            print("use DirectML instead")
+            if(os.path.exists("runtime\Lib\site-packages\onnxruntime\capi\DirectML.dll")==False):
+                try:
+                    os.rename("runtime\Lib\site-packages\onnxruntime", "runtime\Lib\site-packages\onnxruntime-cuda")
+                except:
+                    pass
+                try:
+                    os.rename("runtime\Lib\site-packages\onnxruntime-dml", "runtime\Lib\site-packages\onnxruntime")
+                except:
+                    pass
+            import torch_directml
+            self.device = torch_directml.device(torch_directml.default_device())
+            self.is_half = False
+        else:
+            if self.instead:
+                print(f"use {self.instead} instead")
+            if(os.path.exists("runtime\Lib\site-packages\onnxruntime\capi\onnxruntime_providers_cuda.dll")==False):
+                try:
+                    os.rename("runtime\Lib\site-packages\onnxruntime", "runtime\Lib\site-packages\onnxruntime-dml")
+                except:
+                    pass
+                try:
+                    os.rename("runtime\Lib\site-packages\onnxruntime-cuda", "runtime\Lib\site-packages\onnxruntime")
+                except:
+                    pass
+        return x_pad, x_query, x_center, x_max
+
+config = Config()

configs/.gitignore

@@ -0,0 +1,3 @@
+*
+!.gitignore
+!*.json

configs/32k.json

@@ -0,0 +1,46 @@
+{
+  "train": {
+    "log_interval": 200,
+    "seed": 1234,
+    "epochs": 20000,
+    "learning_rate": 1e-4,
+    "betas": [0.8, 0.99],
+    "eps": 1e-9,
+    "batch_size": 4,
+    "fp16_run": true,
+    "lr_decay": 0.999875,
+    "segment_size": 12800,
+    "init_lr_ratio": 1,
+    "warmup_epochs": 0,
+    "c_mel": 45,
+    "c_kl": 1.0
+  },
+  "data": {
+    "max_wav_value": 32768.0,
+    "sampling_rate": 32000,
+    "filter_length": 1024,
+    "hop_length": 320,
+    "win_length": 1024,
+    "n_mel_channels": 80,
+    "mel_fmin": 0.0,
+    "mel_fmax": null
+  },
+  "model": {
+    "inter_channels": 192,
+    "hidden_channels": 192,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 6,
+    "kernel_size": 3,
+    "p_dropout": 0,
+    "resblock": "1",
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+    "upsample_rates": [10,4,2,2,2],
+    "upsample_initial_channel": 512,
+    "upsample_kernel_sizes": [16,16,4,4,4],
+    "use_spectral_norm": false,
+    "gin_channels": 256,
+    "spk_embed_dim": 109
+  }
+}

configs/32k_v2.json

@@ -0,0 +1,46 @@
+{
+  "train": {
+    "log_interval": 200,
+    "seed": 1234,
+    "epochs": 20000,
+    "learning_rate": 1e-4,
+    "betas": [0.8, 0.99],
+    "eps": 1e-9,
+    "batch_size": 4,
+    "fp16_run": true,
+    "lr_decay": 0.999875,
+    "segment_size": 12800,
+    "init_lr_ratio": 1,
+    "warmup_epochs": 0,
+    "c_mel": 45,
+    "c_kl": 1.0
+  },
+  "data": {
+    "max_wav_value": 32768.0,
+    "sampling_rate": 32000,
+    "filter_length": 1024,
+    "hop_length": 320,
+    "win_length": 1024,
+    "n_mel_channels": 80,
+    "mel_fmin": 0.0,
+    "mel_fmax": null
+  },
+  "model": {
+    "inter_channels": 192,
+    "hidden_channels": 192,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 6,
+    "kernel_size": 3,
+    "p_dropout": 0,
+    "resblock": "1",
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+    "upsample_rates": [10,8,2,2],
+    "upsample_initial_channel": 512,
+    "upsample_kernel_sizes": [20,16,4,4],
+    "use_spectral_norm": false,
+    "gin_channels": 256,
+    "spk_embed_dim": 109
+  }
+}

configs/40k.json

@@ -0,0 +1,46 @@
+{
+  "train": {
+    "log_interval": 200,
+    "seed": 1234,
+    "epochs": 20000,
+    "learning_rate": 1e-4,
+    "betas": [0.8, 0.99],
+    "eps": 1e-9,
+    "batch_size": 4,
+    "fp16_run": true,
+    "lr_decay": 0.999875,
+    "segment_size": 12800,
+    "init_lr_ratio": 1,
+    "warmup_epochs": 0,
+    "c_mel": 45,
+    "c_kl": 1.0
+  },
+  "data": {
+    "max_wav_value": 32768.0,
+    "sampling_rate": 40000,
+    "filter_length": 2048,
+    "hop_length": 400,
+    "win_length": 2048,
+    "n_mel_channels": 125,
+    "mel_fmin": 0.0,
+    "mel_fmax": null
+  },
+  "model": {
+    "inter_channels": 192,
+    "hidden_channels": 192,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 6,
+    "kernel_size": 3,
+    "p_dropout": 0,
+    "resblock": "1",
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+    "upsample_rates": [10,10,2,2],
+    "upsample_initial_channel": 512,
+    "upsample_kernel_sizes": [16,16,4,4],
+    "use_spectral_norm": false,
+    "gin_channels": 256,
+    "spk_embed_dim": 109
+  }
+}

configs/48k.json

@@ -0,0 +1,46 @@
+{
+  "train": {
+    "log_interval": 200,
+    "seed": 1234,
+    "epochs": 20000,
+    "learning_rate": 1e-4,
+    "betas": [0.8, 0.99],
+    "eps": 1e-9,
+    "batch_size": 4,
+    "fp16_run": true,
+    "lr_decay": 0.999875,
+    "segment_size": 11520,
+    "init_lr_ratio": 1,
+    "warmup_epochs": 0,
+    "c_mel": 45,
+    "c_kl": 1.0
+  },
+  "data": {
+    "max_wav_value": 32768.0,
+    "sampling_rate": 48000,
+    "filter_length": 2048,
+    "hop_length": 480,
+    "win_length": 2048,
+    "n_mel_channels": 128,
+    "mel_fmin": 0.0,
+    "mel_fmax": null
+  },
+  "model": {
+    "inter_channels": 192,
+    "hidden_channels": 192,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 6,
+    "kernel_size": 3,
+    "p_dropout": 0,
+    "resblock": "1",
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+    "upsample_rates": [10,6,2,2,2],
+    "upsample_initial_channel": 512,
+    "upsample_kernel_sizes": [16,16,4,4,4],
+    "use_spectral_norm": false,
+    "gin_channels": 256,
+    "spk_embed_dim": 109
+  }
+}

configs/48k_v2.json

@@ -0,0 +1,46 @@
+{
+  "train": {
+    "log_interval": 200,
+    "seed": 1234,
+    "epochs": 20000,
+    "learning_rate": 1e-4,
+    "betas": [0.8, 0.99],
+    "eps": 1e-9,
+    "batch_size": 4,
+    "fp16_run": true,
+    "lr_decay": 0.999875,
+    "segment_size": 17280,
+    "init_lr_ratio": 1,
+    "warmup_epochs": 0,
+    "c_mel": 45,
+    "c_kl": 1.0
+  },
+  "data": {
+    "max_wav_value": 32768.0,
+    "sampling_rate": 48000,
+    "filter_length": 2048,
+    "hop_length": 480,
+    "win_length": 2048,
+    "n_mel_channels": 128,
+    "mel_fmin": 0.0,
+    "mel_fmax": null
+  },
+  "model": {
+    "inter_channels": 192,
+    "hidden_channels": 192,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 6,
+    "kernel_size": 3,
+    "p_dropout": 0,
+    "resblock": "1",
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+    "upsample_rates": [12,10,2,2],
+    "upsample_initial_channel": 512,
+    "upsample_kernel_sizes": [24,20,4,4],
+    "use_spectral_norm": false,
+    "gin_channels": 256,
+    "spk_embed_dim": 109
+  }
+}

custom_nodes/loaders.py

@@ -0,0 +1,254 @@
+import os
+import subprocess
+import sys
+
+from .settings import PITCH_EXTRACTION_OPTIONS
+from ..lib import BASE_MODELS_DIR
+from ..lib.model_utils import load_hubert
+from ..lib.utils import get_file_hash, get_filenames, get_hash, get_optimal_torch_device
+from .utils import model_downloader
+import torch
+import folder_paths
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from ..vc_infer_pipeline import get_vc
+from .settings.downloader import RVC_DOWNLOAD_LINK, RVC_MODELS, download_file
+from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, load_input_audio
+from ..config import config
+
+input_path = folder_paths.get_input_directory()
+temp_path = folder_paths.get_temp_directory()
+
+model_ids = [
+    'openai/whisper-large-v3',
+    'openai/whisper-large-v2',
+    'openai/whisper-large',
+    'openai/whisper-medium',
+    'openai/whisper-small',
+    'openai/whisper-base',
+    'openai/whisper-tiny',
+    'openai/whisper-medium.en',
+    'openai/whisper-small.en',
+    'openai/whisper-base.en',
+    'openai/whisper-tiny.en',
+]
+
+languages = ['en', 'fr', 'es']
+
+spacy_model_map = {
+    "en": "en_core_web_md",
+    "fr": "fr_core_news_md",
+    "es": "es_core_news_md"
+}
+
+class LoadPitchExtractionParams:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            'required': {
+                'f0_method': (PITCH_EXTRACTION_OPTIONS,{"default": "rmvpe"}),
+                "f0_autotune": ("BOOLEAN",),
+                "index_rate": ("FLOAT",{
+                    "default": .75, 
+                    "min": 0., #Minimum value
+                    "max": 1., #Maximum value
+                    "step": .01, #Slider's step
+                    "display": "slider"
+                }),
+                "resample_sr": ([0,16000,32000,40000,44100,48000],{"default": 0}),
+                "rms_mix_rate": ("FLOAT",{
+                    "default": 0.25, 
+                    "min": 0., #Minimum value
+                    "max": 1., #Maximum value
+                    "step": .01, #Slider's step
+                    "display": "slider"
+                }),
+                "protect": ("FLOAT",{
+                    "default": 0.25, 
+                    "min": 0., #Minimum value
+                    "max": .5, #Maximum value
+                    "step": .01, #Slider's step
+                    "display": "slider"
+                })
+            },
+        }
+
+    RETURN_TYPES = ('PITCH_EXTRACTION', )
+    RETURN_NAMES = ('pitch_extraction_params', )
+
+    CATEGORY = "🌺RVC-Studio/loaders"
+
+    FUNCTION = 'load_params'
+
+    def load_params(self, **params):
+        if "rmvpe" in params.get("f0_method",""): model_downloader("rmvpe.pt")
+        return (params,)
+
+class LoadHubertModel:
+    @classmethod
+    def INPUT_TYPES(cls):
+        model_list = ["hubert_base.pt"] + get_filenames(root=BASE_MODELS_DIR,folder="*",exts=["pt"],format_func=os.path.basename)
+        return {
+            'required': {
+                'model': (model_list,{"default": "hubert_base.pt"}),
+            },
+        }
+
+    RETURN_TYPES = ('HUBERT_MODEL', )
+    RETURN_NAMES = ('hubert_model', )
+
+    CATEGORY = "🌺RVC-Studio/loaders"
+
+    FUNCTION = 'load_model'
+
+    def load_model(self, model):
+        model_path = model_downloader(model)
+        hubert_model = lambda:load_hubert(model_path,config=config)
+        return (hubert_model,)
+    
+    @classmethod
+    def IS_CHANGED(cls, model):
+        return get_hash(model)
+
+class LoadRVCModelNode:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        model_list = RVC_MODELS + get_filenames(root=BASE_MODELS_DIR,folder="RVC",exts=["pth"],format_func=lambda x: f"RVC/{os.path.basename(x)}")
+        model_list = list(set(model_list)) # dedupe
+        return {
+            'required': {
+                'model': (model_list,{"default": "RVC/Sayano.pth"}),
+            },
+        }
+
+    RETURN_TYPES = ('RVC_MODEL', 'STRING')
+    RETURN_NAMES = ('model', 'model_name')
+
+    CATEGORY = "🌺RVC-Studio/loaders"
+
+    FUNCTION = 'load_model'
+
+    @classmethod
+    def IS_CHANGED(cls, model):
+        return get_hash(model)
+
+    def load_model(self, model):
+        try:
+            filename = os.path.basename(model)
+            subfolder = os.path.dirname(model)
+            model_path = os.path.join(BASE_MODELS_DIR,subfolder,filename)
+
+            if not os.path.isfile(model_path):
+                download_link = f"{RVC_DOWNLOAD_LINK}{model}"
+                params = model_path, download_link
+                if download_file(params): print(f"successfully downloaded: {model_path}")
+            model = lambda:get_vc(model_path)
+            return (model,filename.split(".")[0])
+        except Exception as e:
+            print(f"Error in {self.__class__.__name__}: {e}")
+            raise e
+    
+class LoadWhisperModelNode:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            'required': {
+                'model_id': (model_ids,{"default": "openai/whisper-base.en"}),
+            },
+            "optional": {
+                'language': (languages,{"default": "en"}),
+                "max_new_tokens": ("INT", {"default": 128, "min": 16, "max": 1024, "display": "slider"}),
+                "chunk_length_s": ("INT", {"default": 30, "min": 15, "max": 60, "display": "slider"}),
+                "batch_size": ("INT", {"default": 16, "min": 1, "max": 128, "display": "slider"}),
+            }
+        }
+
+    RETURN_TYPES = ('TRANSCRIPTION_MODEL', )
+    RETURN_NAMES = ('model', )
+
+    CATEGORY = "🌺RVC-Studio/loaders"
+
+    FUNCTION = 'load_model'
+
+    @classmethod
+    def IS_CHANGED(cls, model):
+        return get_hash(model)
+
+    def load_model(self, model_id, language="en", max_new_tokens=128, chunk_length_s=12, batch_size=16):
+        device = get_optimal_torch_device()
+        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+        model = AutoModelForSpeechSeq2Seq.from_pretrained(
+            model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+        )
+        processor = AutoProcessor.from_pretrained(model_id)
+        model.to(device)
+
+        generate_kwargs = {}
+        if model_id.endswith('.en'):
+            if language != 'en':
+                raise ValueError(f'Model {model_id} only supports English language')
+        else:
+            generate_kwargs['language'] = language
+
+        def pipe():
+            
+            def get_spacy_model():
+                import spacy
+                model_name = spacy_model_map[language]
+                model_path = os.path.join(BASE_MODELS_DIR,model_name)
+                
+                if not os.path.exists(model_path):
+                    subprocess.call([sys.executable, "-m", "spacy", "download", model_name])
+                    model = spacy.load(model_name)
+                    model.to_disk(model_path)
+
+                return spacy.load(model_path)
+                
+            return pipeline(
+                'automatic-speech-recognition',
+                model=model,
+                tokenizer=processor.tokenizer,
+                feature_extractor=processor.feature_extractor,
+                max_new_tokens=max_new_tokens,
+                chunk_length_s=chunk_length_s,
+                batch_size=batch_size,
+                return_timestamps=True,
+                torch_dtype=torch_dtype,
+                device=device,
+                generate_kwargs=generate_kwargs,
+            ),get_spacy_model
+        return (pipe, )
+
+class LoadAudio:
+    @classmethod
+    def INPUT_TYPES(cls):
+        input_dir = input_path
+        files = get_filenames(root=input_dir,exts=SUPPORTED_AUDIO,format_func=os.path.basename)
+        
+        return {
+            "required": {
+                "audio": (files,),
+                "sr": (["None",16000,44100,48000],{"default": "None"}),
+            }}
+
+    CATEGORY = "🌺RVC-Studio/loaders"
+
+    RETURN_TYPES = ("STRING","VHS_AUDIO")
+    RETURN_NAMES = ("audio_name","vhs_audio")
+    FUNCTION = "load_audio"
+
+    def load_audio(self, audio, sr):
+        audio_path = folder_paths.get_annotated_filepath(audio)
+        print(f"{audio_path=} {audio=}")
+        audio_name = os.path.basename(audio).split(".")[0]
+        sr = None if sr=="None" else int(sr)
+        audio = load_input_audio(audio_path,sr=sr)
+        return (audio_name,lambda:audio_to_bytes(*audio))
+
+    @classmethod
+    def IS_CHANGED(cls, audio):
+        audio_path = folder_paths.get_annotated_filepath(audio)
+        print(f"{audio_path=}")
+        return get_file_hash(audio_path)

custom_nodes/output.py

@@ -0,0 +1,54 @@
+import os
+
+from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, bytes_to_audio, load_input_audio, save_input_audio
+from ..lib.utils import get_hash
+from .utils import increment_filename_no_overwrite
+import folder_paths
+  
+class PreviewAudio:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("VHS_AUDIO",),
+                "filename": ("STRING",{"default": "test"}),
+                "save_format": (SUPPORTED_AUDIO,{"default": "flac"},),
+                "save_channels": ([1,2],{"default": 1}),
+                "overwrite_existing": ("BOOLEAN", {"default": True}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING","VHS_AUDIO")
+    RETURN_NAMES = ("output_path","vhs_audio")
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "🌺RVC-Studio/output"
+
+    FUNCTION = "save_audio"
+
+    def save_audio(self, audio, filename, save_format, save_channels, overwrite_existing):
+
+        filename = filename.strip()
+        assert filename, "Filename cannot be empty"
+
+        base_output_dir = folder_paths.get_output_directory()
+        assert os.path.exists(base_output_dir), f"Output directory {base_output_dir} does not exist"
+        output_dir = os.path.join(base_output_dir, 'audio')
+        os.makedirs(output_dir, exist_ok=True)
+        output_path = os.path.join(output_dir, f"{filename}.{save_format}")
+
+        if os.path.isfile(output_path) and not overwrite_existing:
+            output_path = increment_filename_no_overwrite(output_path)
+        
+        input_audio = bytes_to_audio(audio())
+        print(save_input_audio(output_path,input_audio,to_int16=True,to_stereo=save_channels==2))
+
+        return output_path, lambda:audio_to_bytes(*input_audio)
+    
+    @classmethod
+    def IS_CHANGED(cls, audio):
+        return get_hash(audio())

custom_nodes/rvc.py

@@ -0,0 +1,76 @@
+import os
+
+from ..lib.audio import audio_to_bytes, bytes_to_audio, load_input_audio, save_input_audio
+
+from ..vc_infer_pipeline import vc_single
+import folder_paths
+from ..lib.utils import get_hash, get_optimal_torch_device
+from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR
+
+input_path = folder_paths.get_input_directory()
+temp_path = folder_paths.get_temp_directory()
+# output_path = folder_paths.get_output_directory()
+# base_path = os.path.dirname(input_path)
+node_path = os.path.join(BASE_MODELS_DIR,"custom_nodes/ComfyUI-UVR5")
+weights_path = os.path.join(BASE_MODELS_DIR, "uvr5")
+device = get_optimal_torch_device()
+is_half = True
+
+class RVCNode:
+    
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+
+
+        return {
+            "required": {
+                "audio": ("VHS_AUDIO",),
+                "model": ("RVC_MODEL",),
+                "hubert_model": ("HUBERT_MODEL",),
+                "pitch_extraction_params": ("PITCH_EXTRACTION",),
+                "f0_up_key": ("INT",{
+                    "default": 0, 
+                    "min": -14, #Minimum value
+                    "max": 14, #Maximum value
+                    "step": 1, #Slider's step
+                    "display": "slider"
+                }),
+                "format":(["wav", "flac", "mp3"],{
+                    "default": "flac"
+                }),
+            },
+            "optional": {
+                "use_cache": ("BOOLEAN",{"default": True})
+            }
+        }
+    
+    OUTPUT_NODE = True
+
+    RETURN_TYPES = ("VHS_AUDIO",)
+
+    FUNCTION = "convert"
+
+    CATEGORY = "🌺RVC-Studio/rvc"
+
+    def convert(self, audio, model, hubert_model, pitch_extraction_params, format, use_cache=True, **kwargs):
+        
+        cache_name = os.path.join(BASE_CACHE_DIR,"rvc",f"{get_hash(audio(), model, *pitch_extraction_params.items())}.{format}")
+
+        if use_cache and os.path.isfile(cache_name): output_audio = load_input_audio(cache_name)
+        else:
+            input_audio = bytes_to_audio(audio())
+            output_audio = vc_single(hubert_model=hubert_model(),input_audio=input_audio,**model(),**pitch_extraction_params,**kwargs)
+            print(save_input_audio(cache_name, output_audio))
+        
+        return (lambda:audio_to_bytes(*output_audio),)
+
+
+    @classmethod
+    def IS_CHANGED(cls, *args, **kwargs):
+        print(f"{args=} {kwargs=}")
+        return get_hash(*args, *kwargs.items())
+    
+

custom_nodes/settings/__init__.py

@@ -0,0 +1,30 @@
+from functools import lru_cache
+import os
+
+from ...lib import BASE_CACHE_DIR, PersistedDict
+
+MENU_ITEMS = {
+    "Get help": "https://github.com/SayanoAI/RVC-Studio/discussions",
+    "Report a Bug": "https://github.com/SayanoAI/RVC-Studio/issues",
+    "About": """This project provides a comprehensive platform for training RVC models and generating AI voice covers.
+    Check out this github for more info: https://github.com/SayanoAI/RVC-Studio
+    """
+}
+
+DEVICE_OPTIONS = ["cpu","cuda"]
+PITCH_EXTRACTION_OPTIONS = ["rmvpe","rmvpe+"]
+MERGE_OPTIONS=["median","mean","min","max"]
+TTS_MODELS = ["edge","speecht5"]
+N_THREADS_OPTIONS=[1,2,4,8,12,16]
+SR_MAP = {"32k": 32000,"40k": 40000, "48k": 48000}
+
+@lru_cache
+def get_servers():
+    os.makedirs(BASE_CACHE_DIR,exist_ok=True)
+    fname = os.path.join(BASE_CACHE_DIR,"servers.shelve")
+    servers = PersistedDict(fname)
+    return servers
+
+SERVERS = get_servers()
+RVC_INFERENCE_URL = SERVERS.RVC_INFERENCE_URL
+UVR_INFERENCE_URL = SERVERS.UVR_INFERENCE_URL

custom_nodes/settings/downloader.py

@@ -0,0 +1,121 @@
+import re
+from typing import IO, List, Tuple
+import unicodedata
+import requests
+import os
+import zipfile
+
+from ...lib import BASE_CACHE_DIR, BASE_MODELS_DIR
+
+RVC_DOWNLOAD_LINK = 'https://huggingface.co/datasets/SayanoAI/RVC-Studio/resolve/main/'
+MDX_MODELS = ["MDXNET/UVR-MDX-NET-vocal_FT.onnx"]
+KARAFAN_MODELS = ["karafan/MDX23C-8KFFT-InstVoc_HQ.ckpt"]
+VR_MODELS = ["UVR/UVR-DeEcho-DeReverb.pth","UVR/HP5-vocals+instrumentals.pth"]
+RVC_MODELS = [
+    "RVC/Claire.pth",
+    "RVC/Sayano.pth",
+    "RVC/Mae_v2.pth",
+    "RVC/Fuji.pth",
+    "RVC/Monika.pth"]
+BASE_MODELS = ["hubert_base.pt", "rmvpe.pt"]
+VITS_MODELS = ["VITS/pretrained_ljs.pth"]
+PRETRAINED_MODELS = [
+    "pretrained_v2/D48k.pth",
+    "pretrained_v2/G48k.pth",
+    "pretrained_v2/D32k.pth",
+    "pretrained_v2/G32k.pth",
+    "pretrained_v2/D40k.pth",
+    "pretrained_v2/G40k.pth",
+    "pretrained_v2/f0D48k.pth",
+    "pretrained_v2/f0G48k.pth",
+    "pretrained_v2/f0D40k.pth",
+    "pretrained_v2/f0G40k.pth",
+    "pretrained_v2/f0D32k.pth",
+    "pretrained_v2/f0G32k.pth",
+
+    "pretrained_v2/f0_RIN_E3_40k_D.pth",
+    "pretrained_v2/f0_RIN_E3_40k_G.pth",
+    "pretrained_v2/f0Ov2Super32kD.pth",
+    "pretrained_v2/f0Ov2Super32kG.pth",
+    "pretrained_v2/f0Ov2Super40kD.pth",
+    "pretrained_v2/f0Ov2Super40kG.pth",
+]
+LLM_MODELS = [
+    "https://huggingface.co/TheBloke/Airoboros-L2-7B-2.1-GGUF/resolve/main/airoboros-l2-7b-2.1.Q4_K_M.gguf",
+    "https://huggingface.co/TheBloke/Pygmalion-2-7B-GGUF/resolve/main/pygmalion-2-7b.Q4_K_M.gguf",
+    "https://huggingface.co/TheBloke/Zarablend-MX-L2-7B-GGUF/resolve/main/zarablend-mx-l2-7b.Q4_K_M.gguf",
+    "https://huggingface.co/TheBloke/MythoMax-L2-Kimiko-v2-13B-GGUF/resolve/main/mythomax-l2-kimiko-v2-13b.Q4_K_M.gguf"
+]
+STT_MODELS = [
+    "https://alphacephei.com/vosk/models/vosk-model-en-us-0.22-lgraph.zip"
+]
+
+def download_file(params: Tuple[str, str]):
+    model_path, download_link = params
+    if os.path.isfile(model_path): raise FileExistsError(f"{model_path} already exists!")
+    os.makedirs(os.path.dirname(model_path), exist_ok=True)
+    
+    with requests.get(download_link,stream=True) as r:
+        r.raise_for_status()
+        with open(model_path, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+        return True
+
+def download_link_generator(download_link: str,model_list: List[str]):
+    for model in model_list:
+        model_path = os.path.join(BASE_MODELS_DIR,model)
+        yield (model_path, f"{download_link}{model}")
+
+def save_file(params: Tuple[str, any]):
+    (data_path, datum) = params
+    if "zip" in os.path.splitext(data_path)[-1]: return save_zipped_files(params) # unzip
+    else:
+        try:
+            with open(data_path,"wb") as f:
+                f.write(datum)
+            return f"Successfully saved file to: {data_path}"
+        except Exception as e:
+            return f"Failed to save file: {e}"
+
+def save_file_generator(save_dir: str, data: List[IO]):
+    for datum in data:
+        data_path = os.path.join(save_dir,datum.name)
+        yield (data_path, datum.read())
+
+def save_zipped_files(params: Tuple[str, any]):
+    (data_path, datum) = params
+
+    try:
+        print(f"saving zip file: {data_path}")
+        temp_dir = os.path.join(BASE_CACHE_DIR,"zips")
+        os.makedirs(temp_dir,exist_ok=True)
+        name = os.path.basename(data_path)
+        zip_path = os.path.join(temp_dir,name)
+
+        with open(zip_path,"wb") as f:
+            f.write(datum)
+
+        print(f"extracting zip file: {zip_path}")
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            zip_ref.extractall(os.path.dirname(data_path))
+        print(f"finished extracting zip file")
+        
+        os.remove(zip_path) # cleanup
+        return f"Successfully saved files to: {data_path}"
+    except Exception as e:
+        return f"Failed to save files: {e}"
+    
+def slugify_filepath(filepath):
+    # Split the path into directory and filename
+    directory, filename = os.path.split(filepath)
+    # Normalize the filename
+    filename = unicodedata.normalize('NFKD', filename)
+    # Encode the filename as ASCII and ignore errors
+    filename = filename.encode('ascii', 'ignore').decode()
+    # Convert the filename to lowercase
+    filename = filename.lower()
+    # Replace spaces and other unwanted characters with dashes
+    filename = re.sub(r'[^a-z0-9.-]+', '-', filename)
+    # Join the directory and the slugified filename
+    return os.path.join(directory, filename)

custom_nodes/stt.py

@@ -0,0 +1,204 @@
+import json
+import os
+from pprint import pprint
+
+import numpy as np
+
+from ..lib import BASE_CACHE_DIR
+from .utils import increment_filename_no_overwrite
+from ..lib.utils import get_hash
+import folder_paths
+from ..lib.audio import bytes_to_audio, remix_audio
+import textacy
+from textacy.extract import keyterms as kt
+
+temp_path = folder_paths.get_temp_directory()
+
+def extract_keywords(text: str, max_words: int, spacy_model, **kwargs):
+    text = text.strip()
+    doc = textacy.make_spacy_doc(text, lang=spacy_model)
+    topn = int(max_words) if max_words>0 else len(text)
+    include_pos = ["NOUN","ADJ","PROPN","VERB","NUM","ADP"]
+    ngrams = [1,2]
+    tags = []
+    try:
+        terms = kt.sgrank(doc, ngrams=ngrams, normalize="lower", topn=topn, include_pos=include_pos)
+        tags = list(map(lambda v:v[0],sorted(terms,key=lambda v:v[1],reverse=True)[:topn]))
+    except Exception as error:
+        print(f"{text=} {error=}")
+
+    return ", ".join(tags).strip()
+
+def limit_sentence(text: str, max_words: int, **kwargs):
+    text = text.strip()
+    topn = int(max_words) if max_words>0 else len(text)
+    if topn>0: return " ".join(text.split()[:topn])
+    else: return text
+
+class AudioTranscriptionNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "pipeline": ('TRANSCRIPTION_MODEL',),
+                "audio": ('VHS_AUDIO',),
+                "format_newlines_on_punctuation": ('BOOLEAN', {'default': True}),
+                "loop": ('BOOLEAN', {'default': True}),
+                "use_tags": ('BOOLEAN', {'default': True}),
+                "max_words": ('INT', {'default': 16, "min": 0, "max": 32, "display": "slider"}),
+                "frame_interpolation": ("INT", {"default": 0, "min": 0, "max": 120})
+            },
+            "optional": {
+                "save_transcription": ('BOOLEAN', {'default': False}),
+                "save_chunks": ('BOOLEAN', {'default': False}),
+                "save_filename": ('STRING', {'default': 'transcription'}),
+                "overwrite_existing": ('BOOLEAN', {'default': True}),
+                "print_output": ('BOOLEAN', {'default': True}),
+                "use_cache": ("BOOLEAN",{"default": True})
+            }
+        }
+
+    OUTPUT_NODE = True
+    RETURN_TYPES = ("STRING", "STRING", "STRING", "INT", "INT")
+    RETURN_NAMES = ("transcription", "batch_prompt_text", "batch_values_text", "num_chunks","num_frames")
+
+    FUNCTION = "transcribe"
+
+    CATEGORY = "🌺RVC-Studio/stt"
+
+    @staticmethod
+    def save_output(filename,text,overwrite_existing=False):
+        base_output_dir = folder_paths.get_output_directory()
+        assert os.path.exists(base_output_dir), f"Output directory {base_output_dir} does not exist"
+        output_dir = os.path.join(base_output_dir, 'transcriptions')
+        
+        file_path = os.path.join(output_dir, filename)
+        dir_name = os.path.dirname(file_path)
+        os.makedirs(dir_name, exist_ok=True)
+        if os.path.exists(file_path) and not overwrite_existing:
+            file_path = increment_filename_no_overwrite(file_path)
+        with open(file_path, 'w') as f:
+            f.write(text)
+
+    @staticmethod
+    def process_text_chunks(ichunk,text_processor,frame_interpolation,**kwargs):
+        i,chunk=ichunk
+        frame_interpolation=int(frame_interpolation)
+        text = text_processor(chunk["text"],**kwargs)
+        timestamp = np.nan_to_num(np.array(chunk["timestamp"],dtype=float),nan=i)
+        index = int(timestamp[0]*frame_interpolation) if frame_interpolation>0 else i
+        prompt = f'"{index}":"{text}"' if len(text) else f'"{index}":""'
+        duration = max(abs(int(np.round(timestamp[1]-timestamp[0]))),1) # at least 1 sec duration
+        if frame_interpolation>1: duration*=frame_interpolation
+        value = f'"{index}":"({duration})"'
+        return prompt,value
+    
+    @staticmethod
+    def load_cache(*args):
+        results = None
+        cache_name = os.path.join(BASE_CACHE_DIR,"stt",f"{get_hash(*args)}.json")
+        os.makedirs(os.path.dirname(cache_name), exist_ok=True)
+        if os.path.isfile(cache_name): results = json.load(cache_name)
+        return results, cache_name
+
+    def transcribe(
+        self, pipeline, audio,
+        format_newlines_on_punctuation,loop,
+        use_tags,max_words,frame_interpolation,
+        save_transcription=False, save_chunks=False,
+        save_filename="transcription", overwrite_existing=True,
+        print_output=True,use_cache=True
+    ):
+        print('Starting Transcription')
+
+        audio_data = bytes_to_audio(audio())
+        if use_cache: result, cache_name = self.load_cache(pipeline,audio(),loop,use_tags,max_words,frame_interpolation)
+        else: result = None
+
+        audio_frames = int(np.ceil(len(audio_data[0])/audio_data[1]))
+        whisper_model,get_spacy_model=pipeline()
+
+        if result is None:
+            audio,_ = remix_audio(audio_data,target_sr=16000,norm=True)
+            result = whisper_model(audio)
+
+        if use_cache and result:
+            with open(cache_name, "w") as f:
+                json.dump(result, f)
+
+        print('Transcription Done')
+        text = result["text"]
+        total_chunks = result['chunks']
+        if format_newlines_on_punctuation:
+            punct = ['.', '?', '!']
+            for p in punct:
+                text = text.replace(f'{p} ', f'{p}\n')
+
+        if loop: # append first frame to chunk stack
+            last_chunk = None
+            for i in range(len(total_chunks)):
+                last_chunk = total_chunks[-1-i]
+
+                if "timestamp" in last_chunk and len(last_chunk["timestamp"])>=1:
+                    last_chunk["timestamp"] = np.nan_to_num(np.array(last_chunk["timestamp"],dtype=float),nan=audio_frames)
+                    break
+            if last_chunk is not None:
+                timestamp = last_chunk["timestamp"]
+                start_time = timestamp[-1]
+                duration = max(1,audio_frames - start_time)
+                total_chunks.append({
+                    "text": total_chunks[0]["text"],
+                    "timestamp": (start_time,start_time+duration)
+                })
+        else: # adds 1s frame to prevent audio from stopping early
+            num_frames+=1
+
+        if use_tags:
+            spacy_model,text_processor = get_spacy_model(),extract_keywords
+        else:
+            spacy_model,text_processor = None,limit_sentence
+
+        # split transcript into prompt based on timestamp
+        # find length of each frame using timestamp
+        prompts = []
+        values = []
+        for ichunk in enumerate(total_chunks):
+            p, v = self.process_text_chunks(ichunk,text_processor,frame_interpolation,spacy_model=spacy_model,max_words=max_words)
+            prompts.append(p)
+            values.append(v)
+        
+        num_chunks = len(total_chunks)
+        batch_prompt = ',\n'.join(prompts)
+        batch_values = ',\n'.join(values)
+        num_frames = int(audio_frames*frame_interpolation if frame_interpolation>0 else audio_frames)
+
+        if print_output:
+            pprint(total_chunks)
+            print(batch_prompt)
+            print(batch_values)
+            print(f"{num_chunks=}, {audio_frames=}s, {num_frames=}")
+            
+        if save_transcription:
+            self.save_output(save_filename + '.txt', text, overwrite_existing=overwrite_existing)
+
+        if save_chunks:
+            self.save_output(
+                save_filename + '_prompt.json',
+                json.dumps(json.loads("{"+batch_prompt+"}"),indent=2), # formats json with indent
+                overwrite_existing=overwrite_existing)
+            
+            self.save_output(
+                save_filename + '_values.json',
+                json.dumps(json.loads("{"+batch_values+"}"),indent=2), # formats json with indent
+                overwrite_existing=overwrite_existing)
+
+        text = text_processor(text,max_words=max_words,spacy_model=spacy_model)
+        return (text, batch_prompt, batch_values, num_chunks, num_frames)
+    
+    @classmethod
+    def IS_CHANGED(cls, *args, **kwargs):
+        print(f"{args=} {kwargs=}")
+        return get_hash(*args, *kwargs.items())

custom_nodes/utils.py

@@ -0,0 +1,368 @@
+import itertools
+import os
+from pprint import pprint
+import re
+import tempfile
+import torch
+import folder_paths
+from .settings import MERGE_OPTIONS
+from ..lib.utils import gc_collect, get_hash
+from .settings.downloader import RVC_DOWNLOAD_LINK, download_file
+from ..lib import BASE_MODELS_DIR
+import numpy as np
+from ..lib.audio import MAX_INT16, audio_to_bytes, bytes_to_audio, merge_audio, remix_audio
+
+CATEGORY = "🌺RVC-Studio/utils"
+temp_path = folder_paths.get_temp_directory()
+
+def model_downloader(model):
+    filename = os.path.basename(model)
+    subfolder = os.path.dirname(model)
+    model_path = os.path.join(BASE_MODELS_DIR,subfolder,filename)
+    if not os.path.isfile(model_path):
+        download_link = f"{RVC_DOWNLOAD_LINK}{model}"
+        params = model_path, download_link
+        if download_file(params): print(f"successfully downloaded: {model_path}")
+    return model_path
+
+def increment_filename_no_overwrite(proposed_path):
+    output_dir, filename_ext = os.path.split(proposed_path)
+    filename, file_format = os.splitext(filename_ext)
+    files = os.listdir(output_dir)
+    files = filter(lambda f: os.path.isfile(os.path.join(output_dir, f)), files)
+    files = filter(lambda f: f.startswith(filename), files)
+    files = filter(lambda f: f.endswith('.' + file_format), files)
+    file_numbers = [re.search(r'_(\d+)\.', f) for f in files]
+    file_numbers = [int(f.group(1)) for f in file_numbers if f]
+    this_file_number = max(file_numbers) + 1 if file_numbers else 1
+    output_path = os.path.join(output_dir, filename + f'_{this_file_number}.' + file_format)
+    return output_path
+
+class MergeImageBatches:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "iterate": ("BOOLEAN",{"default": False})
+            },
+        }
+
+    INPUT_IS_LIST = (True, False)
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "merge"
+
+    CATEGORY = CATEGORY
+
+    __archive__ = None
+
+    @staticmethod
+    def iter_images(fname,images):
+        size = len(images)
+        if size==0: return images
+        shape = images[0].shape
+        if len(shape)>3: shape=shape[1:]
+
+        memmap = np.memmap(fname, mode='w+', dtype=np.float32, shape=(size,*shape))
+        
+        for i,img in enumerate(itertools.chain(*images)):
+            if i >= len(memmap): #need to expand memmap
+                print(f"index {i} is larger than {len(memmap)}. expanding memmap to {size+i}...")
+                memmap.flush()
+                memmap = np.memmap(fname, dtype=np.float32, mode='r+',shape=(size+i,*shape), order='F')
+            memmap[i]=img
+            # del img
+            # gc_collect()
+        memmap.flush()
+
+        print(f"{fname=} size={i} shape={memmap.shape}")
+
+        # del images
+        images = torch.from_numpy(memmap[:i])
+        
+        return images
+    
+    def __del__(self):
+        try:
+            if os.path.exists(self.__archive__): os.remove(self.__archive__)
+            gc_collect()
+        except Exception as error:
+            print(f"Failed to delete {self.__archive__}: {error=}")
+
+    def merge(self, images, iterate):
+        
+        if len(images) <= 1:
+            return (images[0],)
+        else:
+            if hasattr(iterate,"pop"): iterate=iterate.pop()
+            print(f"{iterate=}")
+
+            if iterate:
+                if self.__archive__ is None:
+                    with tempfile.NamedTemporaryFile(delete=False) as ntf:
+                        self.__archive__ = ntf.name
+                images = self.iter_images(self.__archive__,images)
+            else:
+                images = torch.cat(images)
+            
+            gc_collect()
+            print(f"Merged images: {len(images)=} {images.shape=}")
+            
+            return (images,)
+        
+        
+class MergeLatentBatches:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "latents": ("LATENT",),
+        }
+        }
+
+    INPUT_IS_LIST = True
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "merge"
+
+    CATEGORY = CATEGORY
+
+    @staticmethod
+    def merge_tensor(tensors):
+        return torch.cat(tensors, dim=0)
+
+    def merge(self, latents):
+        merged_latents = {}
+        gc_collect()
+        merged_latents["samples"] = self.merge_tensor([latent["samples"] for latent in latents])
+        merged_latents["noise_mask"] = self.merge_tensor([latent["noise_mask"] for latent in latents])
+        merged_latents["batch_index"] = self.merge_tensor([latent["batch_index"] for latent in latents])
+        
+        print(f'Merged latents: {len(merged_latents["samples"])=} {latents["batch_index"]=}')
+        return (merged_latents,)
+        
+class AudioBatchValueNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ('VHS_AUDIO',),
+                "num_frames": ('INT', {"min": 1, "max": 128, "step": 1}),
+                "output_min": ('FLOAT', {'default': 0., "min": -1000., "max": 1000., "step": .01}),
+                "output_max": ('FLOAT', {'default': 1., "min": 0., "max": 1000., "step": .01}),
+                "norm": (["scale","tanh","sigmoid"], {"default": "scale"}),
+            },
+            "optional": {
+                "silence_threshold": ("INT", {"default": 1000, "min": 1, "max": MAX_INT16, "step": 1, "display": "slider"}),
+                "frame_multiplier": ("INT", {"default": 1, "min": 1, "max": 120, "step": 1}),
+                "print_output": ("BOOLEAN", {"default": False}),
+                "inverse": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING","INT")
+    RETURN_NAMES = ("batch_value_text","num_values")
+
+    FUNCTION = "get_frame_weights"
+
+    CATEGORY = CATEGORY
+
+    @staticmethod
+    def get_rms(audio): # root mean squared of audio segment
+        return np.sqrt(np.nanmean(audio**2))
+
+    def get_frame_weights(self, audio, num_frames, output_min, output_max, norm,
+                          silence_threshold=1000, frame_multiplier=1, print_output=False, inverse=False):
+        assert output_max>=output_min, f"{output_max=} must be greater or equal to {output_min=}!"
+
+        audio_data = bytes_to_audio(audio())
+        audio,_ = remix_audio(audio_data,norm=True,to_int16=True)
+        num_values = int(num_frames*frame_multiplier)
+        audio_rms = np.nan_to_num(list(map(self.get_rms,np.array_split(audio.flatten()/silence_threshold, num_values))),nan=0)
+        audio_zscore = (audio_rms-audio_rms.mean())/audio_rms.std()
+        output_range = output_max-output_min
+
+        print(f"{audio_rms.min()=} {audio_rms.max()=} {audio_rms.mean()=}")
+        if norm=="tanh":
+            x_norm = np.tanh(audio_zscore) #tanh=-1 to 1
+            if inverse: x_norm*=-1
+            x_norm = (x_norm * output_range + output_max + output_min)/2
+        elif norm=="sigmoid":
+            x_norm = 1. / (1. + np.exp(-audio_zscore)) #sigmoid=0 to 1
+            if inverse: x_norm=1-x_norm
+            x_norm = x_norm * output_range + output_min 
+        else:
+            x_min = audio_zscore.min() 
+            x_norm = (audio_zscore - x_min) / (audio_zscore.max() - x_min) #scale=0 to 1
+            if inverse: x_norm=1-x_norm
+            x_norm = x_norm * output_range + output_min 
+        print(f"{x_norm.min()=} {x_norm.max()=} {x_norm.mean()=}")
+
+        batch_value = ",\n".join([f'{n}:({v})' for n,v in enumerate(x_norm)])
+        if print_output: pprint(batch_value)
+        return (batch_value,num_values)
+    
+    @classmethod
+    def IS_CHANGED(cls, *args, **kwargs):
+        print(f"{args=} {kwargs=}")
+        return get_hash(*args, *kwargs.items())
+    
+class ImageRepeatInterleavedNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "repeats": ("INT", {"default": 1, "min": 1, "max": 4096}),
+            },
+            "optional": {
+                "fps": ("INT", {"default": 1, "min": 1})
+            }
+        }
+    RETURN_TYPES = ("IMAGE",)
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True, )
+
+    FUNCTION = "rebatch"
+
+    CATEGORY = CATEGORY
+
+    def rebatch(self, images, repeats, fps=1):
+        repeats = np.array(repeats).flatten()*fps
+        print(f"{repeats=}")
+        all_images = []
+        for img in images:
+            for i in range(img.shape[0]):
+                all_images.append(img[i:i+1])
+
+        # expand repeat
+        if len(repeats)==1:
+            repeats = repeats.repeat(len(all_images),axis=0)
+        elif len(repeats)<len(all_images):
+            pad_width = len(all_images) - len(repeats)
+            repeats = np.pad(repeats, (0, pad_width), mode='constant', constant_values=1)
+
+        for i,img in enumerate(all_images):
+            if repeats[i]>1: all_images[i] = img.expand(repeats[i], *img.shape).flatten(0,1)
+            else: all_images[i] = img
+            print(f"{img.shape=} {all_images[i].shape=}")
+
+        print(f"{len(images)=} {len(all_images)=}")
+        gc_collect()
+
+        return (all_images,)
+    
+class LatentRepeatInterleavedNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "latents": ("LATENT",),
+                "repeats": ("INT", {"default": 1, "min": 1, "max": 4096}),
+            },
+            "optional": {
+                "fps": ("INT", {"default": 1, "min": 1})
+            }
+        }
+    RETURN_TYPES = ("LATENT",)
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True, )
+
+    FUNCTION = "merge"
+
+    CATEGORY = CATEGORY
+
+    @staticmethod
+    def repeat(latents, repeat):
+        if latents.ndim>3: #batched
+            expanded_latents=(latents.expand(repeat, *latents.shape).flatten(0,1))
+        else:
+            expanded_latents=(latents.expand(repeat, *latents.shape))
+
+        print(f"{len(latents)}*{repeat}={len(expanded_latents)} ({latents.shape=} => {expanded_latents.shape=})")
+        return expanded_latents
+    
+    @staticmethod
+    def rebatch(samples, repeats):
+        offset = 0
+        for i in range(len(samples)):
+
+            samples[i]["batch_index"] = [i+offset for i in range(repeats[i])]
+            offset += repeats[i]
+            
+            samples[i]["samples"] = LatentRepeatInterleavedNode.repeat(samples[i]["samples"],repeats[i])
+
+            if "noise_mask" in samples[i]:
+                samples[i]["noise_mask"] = LatentRepeatInterleavedNode.repeat(samples[i]["noise_mask"],repeats[i])
+
+        return samples
+
+    def merge(self, latents, repeats, fps=1):
+        repeats = np.array(repeats).flatten()*fps
+        print(f"{repeats=}")
+
+        # combine all latents
+        new_latents = []
+        for latent in latents:
+            if latent["samples"].ndim==3: #unbatched
+                new_latents.append(latent)
+            else:
+                for i in range(len(latent["samples"])):
+                    new_latent = {
+                        "samples": latent["samples"][i],
+                    }
+                    if "noise_mask" in latent: new_latent["noise_mask"] = latent["noise_mask"][i]
+                    new_latents.append(new_latent)
+
+        # expand repeat
+        if len(repeats)==1:
+            repeats = repeats.repeat(len(new_latents),axis=0)
+        elif len(repeats)<len(new_latents):
+            pad_width = len(new_latents) - len(repeats)
+            repeats = np.pad(repeats, (0, pad_width), mode='constant', constant_values=1)
+
+        new_latents = self.rebatch(new_latents,repeats.astype(int))
+        gc_collect()
+
+        return (new_latents,)
+    
+class MergeAudioNode:
+   
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio1": ("VHS_AUDIO",),
+                "audio2": ("VHS_AUDIO",),
+                "sr": (["None",32000,40000,44100,48000],{
+                    "default": "None"
+                }),
+                "merge_type": (MERGE_OPTIONS,{"default": "mean"}),
+                "normalize": ("BOOLEAN",{"default": True})
+            }
+        }
+
+    RETURN_TYPES = ("VHS_AUDIO",)
+    RETURN_NAMES = ("vhs_audio",)
+
+    FUNCTION = "merge"
+
+    CATEGORY = CATEGORY
+
+    def merge(self, audio1, audio2, sr, merge_type, normalize):
+        if sr=="None": sr=None
+        input_audio1 = bytes_to_audio(audio1())
+        input_audio2 = bytes_to_audio(audio2())
+        merged_audio = merge_audio(input_audio1,input_audio2,sr=sr,to_int16=True,norm=normalize,merge_type=merge_type)
+        return (lambda: audio_to_bytes(*merged_audio),)
+
+    @classmethod
+    def IS_CHANGED(cls, *args, **kwargs):
+        print(f"{args=} {kwargs=}")
+        return get_hash(*args, *kwargs.items())

custom_nodes/uvr.py

@@ -0,0 +1,86 @@
+import os
+
+from ..lib.audio import audio_to_bytes, bytes_to_audio, save_input_audio
+import folder_paths
+from ..lib.utils import get_filenames, get_hash, get_optimal_torch_device
+from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR, karafan
+from ..uvr5_cli import Separator
+from .settings.downloader import KARAFAN_MODELS, MDX_MODELS, RVC_DOWNLOAD_LINK, VR_MODELS, download_file
+
+temp_path = folder_paths.get_temp_directory()
+cache_dir = os.path.join(BASE_CACHE_DIR,"uvr")
+device = get_optimal_torch_device()
+is_half = True
+
+class UVR5Node:
+ 
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+
+        model_list = MDX_MODELS + VR_MODELS + KARAFAN_MODELS + get_filenames(root=BASE_MODELS_DIR,format_func=os.path.basename)
+
+        return {
+            "required": {
+                "audio": ("VHS_AUDIO",),
+                "model": (model_list,{
+                    "default": "UVR/HP5-vocals+instrumentals.pth"
+                }),
+                "agg":("INT",{
+                    "default": 10, 
+                    "min": 0, #Minimum value
+                    "max": 20, #Maximum value
+                    "step": 1, #Slider's step
+                    "display": "slider"
+                }),
+                "format":(["wav", "flac", "mp3"],{
+                    "default": "flac"
+                }),
+            },
+            "optional": {
+                "use_cache": ("BOOLEAN",{"default": True})
+            }
+        }
+
+    RETURN_TYPES = ("VHS_AUDIO","VHS_AUDIO","VHS_AUDIO")
+    RETURN_NAMES = ("primary_stem","secondary_stem","audio_passthrough")
+
+    FUNCTION = "split"
+
+    CATEGORY = "🌺RVC-Studio/uvr"
+
+    def split(self, audio, model, **kwargs):
+        filename = os.path.basename(model)
+        subfolder = os.path.dirname(model)
+        model_path = os.path.join(BASE_MODELS_DIR,subfolder,filename)
+        if not os.path.isfile(model_path):
+            download_link = f"{RVC_DOWNLOAD_LINK}{model}"
+            params = model_path, download_link
+            if download_file(params): print(f"successfully downloaded: {model_path}")
+        
+        input_audio = bytes_to_audio(audio())
+        audio_path = os.path.join(temp_path,"uvr",f"{get_hash(audio(),model, *kwargs.items())}.wav")
+        if not os.path.isfile(audio_path):
+            os.makedirs(os.path.dirname(audio_path),exist_ok=True)
+            print(save_input_audio(audio_path,input_audio))
+
+        if "karafan" in model_path:
+            vocals, instrumental, input_audio = karafan.inference.Process(audio_path,cache_dir=cache_dir,**kwargs)
+        else:
+            model = Separator(
+                    model_path=model_path,
+                    device=device,
+                    is_half="cuda" in str(device),
+                    cache_dir=cache_dir,              
+                    **kwargs
+                    )
+            vocals, instrumental, input_audio = model.run_inference(audio_path,format=kwargs.get("format"))
+        
+        return (lambda:audio_to_bytes(*vocals), lambda:audio_to_bytes(*instrumental), audio)
+
+    @classmethod
+    def IS_CHANGED(cls, *args, **kwargs):
+        print(f"{args=} {kwargs=}")
+        return get_hash(*args, *kwargs.items())

i18n/__init__.py

@@ -0,0 +1,29 @@
+import json
+import locale
+import os
+
+class I18nAuto:
+    def __init__(self, language=None):
+        if language in ["Auto", None]:
+            language = locale.getdefaultlocale()[
+                0
+            ]  # getlocale can't identify the system's language ((None, None))
+        if not os.path.exists(f"./{language}.json"):
+            language = "en_US"
+        self.language = language
+        self.language_map = self.load_language_list(language)
+        self.print()
+
+    def __call__(self, key):
+        return self.language_map.get(key, key)
+
+    def print(self):
+        print("Use Language:", self.language)
+
+    @staticmethod
+    def load_language_list(language):
+        with open(f"./i18n/{language}.json", "r", encoding="utf-8") as f:
+            language_list = json.load(f)
+        return language_list
+    
+i18n = I18nAuto()

i18n/en_US.json

@@ -0,0 +1,77 @@
+{
+    "inference.save.button": "Save Options",
+    "inference.refresh_data.button": "Refresh Data",
+    "inference.clear_data.button": "Unload Model",
+    "inference.load_model.button": "Load Model",
+    "inference.one_click.button": "1-Click VC",
+    "inference.song.selectbox": "Select a song file",
+    "inference.voice.selectbox": "Select a voice model",
+
+    "inference.split_vocals": "Extract Vocals",
+    "inference.split_vocals.expander": "Extract Vocals Options",
+    "inference.preprocess_model": "Choose one or more preprocessing models to run in sequence (optional)",
+    "inference.uvr_models": "Choose one or more vocal extraction models to run in parallel (required)",
+    "inference.postprocess_model": "Choose one or more postprocessing model to run in sequence (optional)",
+    "inference.device": "Device",
+    "inference.format": "Audio Format",
+    "inference.channels": "Output Channels",
+    "inference.merge_type": "Merge Type",
+    "inference.agg": "Aggressiveness in isolating vocals",
+    "inference.use_cache": "Reduce future processing time by saving results to disk",
+
+    "inference.convert_vocals": "Change Vocals",
+    "inference.convert_vocals.expander": "Change Vocals Options",
+    "inference.convert_vocals.expander.form_submit_button": "Save Options",
+    
+    "inference.f0_up_key": "Recommended +7 or +12 key for male to female conversion, and -12 or -5 key for female to male conversion. If the sound range goes too far and the voice is distorted, you can also adjust it to the appropriate range by yourself.",
+    "inference.f0_method": "Select the pitch extraction algorithm ('crepe': better quality but GPU intensive, 'rmvpe': fast and best quality, and little GPU requirement, others: use them if you really want to...)",
+    "inference.f0_autotune": "Enable autotuning of extracted pitch",
+    "inference.filter_radius": "If >=3: apply median filtering to the harvested pitch results. The value represents the filter radius and can reduce breathiness.",
+    "inference.index_rate": "Search feature ratio (controls accent strength, too high has artifacting):",
+    "inference.resample_sr": "Resample the output audio in post-processing to the final sample rate. Set to 0 for no resampling:",
+    "inference.rms_mix_rate": "Adjust the volume envelope scaling. Closer to 0, the more it mimicks the volume of the original vocals. Can help mask noise and make volume sound more natural when set relatively low. Closer to 1 will be more of a consistently loud volume:",
+    "inference.protect": "Protect voiceless consonants and breath sounds to prevent artifacts such as tearing in electronic music. Set to 0.5 to disable. Decrease the value to increase protection, but it may reduce indexing accuracy:",
+    
+    "inference.download.button": "Download Song",
+
+    "training.preprocess_data.title": "Process Data",
+    "training.preprocess_data.text": "Step 1: Fill in the experimental configuration. Experimental data is stored in the 'logs' folder, with each experiment having a separate folder. Manually enter the experiment name path, which contains the experimental configuration, logs, and trained model files.",
+    "training.exp_dir": "Enter the voice model name:",
+    "training.sr": "Target sample rate:",
+    "training.if_f0": "Whether the model has pitch guidance (required for singing, optional for speech):",
+    "training.version": "Version:",
+    "training.device": "Device:",
+    "training.n_threads": "Number of CPU processes used for pitch extraction and data processing:",
+    "training.extract_features.title": "Extract Features",
+    "training.extract_features.text": "Step 2: Automatically traverse all files in the training folder that can be decoded into audio and perform slice normalization. Generates 2 wav folders in the experiment directory. Currently, only single-singer/speaker training is supported.",
+    "training.preprocess_data.trainset_dir": "Enter the path of the training folder:",
+    "training.sid": "Please specify the speaker/singer ID:",
+    "training.preprocess_data.submit": "1. Process data",
+    "training.train_model.gpus": "Enter the GPU index(es) separated by '-', e.g., 0-1-2 to use GPU 0, 1, and 2:",
+    "training.f0method": "Select the pitch extraction algorithm ('crepe': better quality but GPU intensive, 'rmvpe': best quality, and little GPU requirement)",
+    "training.extract_features.submit": "2. Feature extraction",
+    "training.train_model.title": "Train Model",
+    "training.train_model.text": "Step 3: Fill in the training settings and start training the model and index",
+    "training.save_every_epoch": "Save frequency (save_every_epoch):",
+    "training.total_epoch": "Total training epochs (total_epoch):",
+    "training.batch_size": "Batch size per GPU:",
+    "training.if_save_latest": "Save only the latest '.ckpt' file to save disk space",
+    "training.if_save_best": "Save only the best model (lowest overall loss) to the 'logs' folder",
+    "training.if_cache_gpu": "Cache all training sets to GPU memory. Caching small datasets (less than 10 minutes) can speed up training, but caching large datasets will consume a lot of GPU memory and may not provide much speed improvement",
+    "training.if_save_every_weights": "Save a small final model to the 'logs' folder at each save point",
+    "training.pretrained_G": "Load pre-trained base model G path:",
+    "training.pretrained_D": "Load pre-trained base model D path:",
+    "training.train_model.submit": "3. Train model",
+    "training.train_index.submit": "4. Train feature index",
+    "training.train_speaker.submit": "Train Optional Speaker Embedding",
+    "training.one_click": "One-click training",
+    "training.gpus": "Which CUDA GPU to use for training (leave empty to use all of them)",
+
+    "process.pids": "Active Processes",
+    "process.kill_all_pids": "Close All Processes",
+    "process.kill_one_pid": "Stop",
+
+        "tts.inference": "Text to Speech",
+        "tts.options": "TTS Options",
+        "tts.model.selectbox": "TTS Models"
+}

i18n/es_ES.json

@@ -0,0 +1,124 @@
+{
+  "inference.save.button": "Guardar opciones",
+  "inference.refresh_data.button": "Actualizar datos",
+  "inference.clear_data.button": "Descargar modelo",
+  "inference.load_model.button": "Cargar modelo",
+  "inference.one_click.button": "1-Click VC",
+  "inference.song.selectbox": "Elige un archivo",
+  "inference.voice.selectbox": "Selecciona un modelo de voz",
+  "inference.split_vocals": "Extraer voces",
+  "inference.split_vocals.expander": "Opciones de extracción de voces",
+  "inference.preprocess_model": "Elige un modelo de preprocesamiento",
+  "inference.uvr5_name": "Elige un modelo de extracción vocal",
+  "inference.device": "Dispositivo",
+  "inference.merge_type": "Tipo de unión",
+  "inference.agg": "Agresividad en el aislamiento vocal",
+  "inference.use_cache": "Reduzca el tiempo de procesamiento futuro guardando los resultados en el disco",
+  "inference.convert_vocals": "Cambiar de voz.",
+  "inference.convert_vocals.expander": "Cambiar opciones de voz",
+  "inference.convert_vocals.expander.form_submit_button": "Guardar opciones",
+  "inference.f0_up_key": "Se recomienda la tecla +12 para la conversión de hombre a mujer y la tecla -12 para la conversión de mujer a hombre. Si el rango de sonido va demasiado lejos y la voz está distorsionada, también puede ajustarlo al rango apropiado usted mismo.",
+  "inference.f0_method": "Seleccione el algoritmo de extracción de tono ('crepe': mejor calidad pero intensivo en GPU, 'rmvpe': mejor calidad y poco requisito de GPU)",
+  "inference.filter_radius": "Si >=3: aplicar filtrado de mediana a los resultados de brea cosechada. El valor representa el radio del filtro y puede reducir la transpiración.",
+  "inference.index_rate": "Relación de funciones de búsqueda (controla la intensidad del acento, demasiado alto tiene artefactos):",
+  "inference.resample_sr": "Vuelva a muestrear el audio de salida en el posprocesamiento a la frecuencia de muestreo final. Establecer en 0 para no volver a muestrear:",
+  "inference.rms_mix_rate": "Ajuste la escala de la envolvente de volumen. Más cerca de 0, más imita el volumen de las voces originales. Puede ayudar a enmascarar el ruido y hacer que el volumen suene más natural cuando se configura relativamente bajo. Más cerca de 1 será más un volumen consistentemente alto:",
+  "inference.protect": "Proteja las consonantes sordas y los sonidos de la respiración para evitar artefactos como el desgarro en la música electrónica. Establecer en 0.5 para desactivar. Disminuya el valor para aumentar la protección, pero puede reducir la precisión de la indexación:",
+  "inference.download.button": "Descargar canción",
+  "training.preprocess_data.title": "Procesar datos",
+  "training.preprocess_data.text": "Paso 1: Complete la configuración experimental. Los datos experimentales se almacenan en la carpeta 'logs', y cada experimento tiene una carpeta separada. Introduzca manualmente la ruta del nombre del experimento, que contiene la configuración experimental, los registros y los archivos de modelo entrenados.",
+  "training.exp_dir": "Introducir el nombre de modelo:",
+  "training.sr": "Tasa de muestreo objetivo:",
+  "training.if_f0": "Si el modelo tiene guía de tono (requerido para cantar, opcional para el habla):",
+  "training.version": "Versión:",
+  "training.device": "Producto:",
+  "training.n_threads": "Número de procesos de CPU utilizados para la extracción de pitch y el procesamiento de datos:",
+  "training.extract_features.title": "Extraer características",
+  "training.extract_features.text": "Paso 2: Recorra automáticamente todos los archivos de la carpeta de capacitación que se pueden decodificar en audio y realice la normalización del segmento. Genera 2 carpetas wav en el directorio del experimento. Actualmente, solo se admite la capacitación de un solo cantante/orador.",
+  "training.preprocess_data.trainset_dir": "Introduzca la ruta de la carpeta de formación:",
+  "training.sid": "Por favor, especifique la identificación del orador/cantante:",
+  "training.preprocess_data.submit": "1. Procesar datos",
+  "training.train_model.gpus": "Introduzca los índices de la GPU separados por '-', por ejemplo, 0-1-2 para usar la GPU 0, 1 y 2:",
+  "training.f0method": "Seleccione el algoritmo de extracción de tono ('crepe': mejor calidad pero intensivo en GPU, 'rmvpe': mejor calidad y poco requisito de GPU)",
+  "training.extract_features.submit": "2. Extracción de característicos",
+  "training.train_model.title": "Modelo de tren",
+  "training.train_model.text": "Paso 3: Complete la configuración de capacitación y comience a capacitar el modelo y el índice",
+  "training.save_every_epoch": "Guardar frecuencia (save_every_epoch):",
+  "training.total_epoch": "Total de épocas de entrenamiento (total_Epoch):",
+  "training.batch_size": "Tamaño de lote por GPU:",
+  "training.if_save_latest": "Guarde solo el último archivo '.ckpt' para ahorrar espacio en disco",
+  "training.if_save_best": "Guarde solo el mejor modelo (pérdida general más baja) en la carpeta 'logs'",
+  "training.if_cache_gpu": "Almacene en caché todos los conjuntos de entrenamiento en la memoria de la GPU. El almacenamiento en caché de pequeños conjuntos de datos (menos de 10 minutos) puede acelerar la capacitación, pero el almacenamiento en caché de grandes conjuntos de datos consumirá mucha memoria de la GPU y puede no proporcionar mucha mejora de la velocidad",
+  "training.if_save_every_weights": "Guarde un pequeño modelo final en la carpeta 'logs' en cada punto de guardado",
+  "training.pretrained_G": "Cargar ruta base modelo G pre-entrenada:",
+  "training.pretrained_D": "Ruta del modelo D base de carga pre-entrenada:",
+  "training.train_model.submit": "5. Entrenar modelo",
+  "training.train_index.submit": "4. Índice de características",
+  "training.train_speaker.submit": "Incrustación de altavoces opcionales de Train",
+  "training.one_click": "Capacitación con un solo clic",
+  "process.pids": "Procesos Activos",
+  "process.kill_all_pids": "Cerrar todos los procesos",
+  "process.kill_one_pid": "Parada",
+  "tts.inference": "Texto a voz",
+  "tts.options": "Opciones de TTS",
+  "tts.model.selectbox": "Modelos TTS",
+  "ckpt处理": "procesamiento ckpt",
+  "模型融合, 可用于测试音色融合": "Modelo de fusión, se puede utilizar para probar la fusión de timbre",
+  "A模型路径": "Ruta al Modelo A:",
+  "B模型路径": "Ruta al Modelo B:",
+  "A模型权重": "Peso (w) para el Modelo A:",
+  "模型是否带音高指导": "Si el modelo tiene guía de tono:",
+  "要置入的模型信息": "Información del modelo a colocar:",
+  "保存的模型名不带后缀": "Nombre de modelo guardado (sin extensión):",
+  "模型版本型号": "Versión de la arquitectura del modelo:",
+  "融合": "Fusión",
+  "修改模型信息(仅支持weights文件夹下提取的小模型文件)": "Modificar la información del modelo (solo se admite para archivos de modelos pequeños extraídos de la carpeta 'weights')",
+  "模型路径": "Ruta al modelo:",
+  "要改的模型信息": "Información del modelo a modificar:",
+  "保存的文件名, 默认空为和源文件同名": "Guardar nombre de archivo (predeterminado: igual que el archivo de origen):",
+  "修改": "Modificar",
+  "查看模型信息(仅支持weights文件夹下提取的小模型文件)": "Ver información del modelo (solo se admite para archivos de modelos pequeños extraídos de la carpeta 'weights')",
+  "查看": "Ver",
+  "模型提取(输入logs文件夹下大文件模型路径),适用于训一半不想训了模型没有自动提取保存小文件模型,或者想测试中间模型的情况": "Extracción del modelo (introduzca la ruta del modelo de archivo grande en la carpeta 'logs'). Esto es útil si desea dejar de entrenar a mitad de camino y extraer y guardar manualmente un archivo de modelo pequeño, o si desea probar un modelo intermedio:",
+  "保存名": "Guardar Nombre",
+  "模型是否带音高指导,1是0否": "Si el modelo tiene guía de tono (1: sí, 0: no):",
+  "提取": "Extraer",
+  "Onnx导出": "Exportar Onnx",
+  "RVC模型路径": "Ruta del modelo RVC:",
+  "Onnx输出路径": "Ruta de exportación Onnx:",
+  "导出Onnx模型": "Exportar modelo Onnx",
+  "常见问题解答": "FAQ (Preguntas más frecuentes)",
+  "招募音高曲线前端编辑器": "Reclutamiento de editores front-end para curvas de tono",
+  "加开发群联系我xxxxx": "Únete al grupo de desarrollo y ponte en contacto conmigo en xxxxx",
+  "点击查看交流、问题反馈群号": "Haga clic para ver el número del grupo de comunicación y comentarios sobre el problema",
+  "xxxxx": "xxxxx",
+  "加载模型": "Cargar modelo",
+  "Hubert模型": "Modelo Hubert",
+  "选择.pth文件": "Seleccione un archivo.",
+  "选择.index文件": "Seleccione el archivo .index",
+  "选择.npy文件": "Seleccione un archivo.",
+  "输入设备": "Dispositivo de entrada",
+  "输出设备": "Dispositivo de salida",
+  "音频设备(请使用同种类驱动)": "Dispositivo de audio (utilice el mismo tipo de controlador)",
+  "响应阈值": "Umbral de respuesta",
+  "音调设置": "Ajustes de tono",
+  "Index Rate": "Tasa de Índice",
+  "常规设置": "Configuración general",
+  "采样长度": "Longitud de la muestra",
+  "淡入淡出长度": "Longitud del desvanecimiento",
+  "额外推理时长": "Tiempo de inferencia adicional",
+  "输入降噪": "Reducción de ruido de entrada",
+  "输出降噪": "Reducción de ruido de salida",
+  "性能设置": "Ajustes de rendimiento",
+  "开始音频转换": "Iniciar conversión de audio",
+  "停止音频转换": "Detener conversión de audio",
+  "推理时间(ms):": "Tiempo de inferencia (ms):",
+  "请选择pth文件": "请选择pth文件",
+  "请选择index文件": "请选择index文件",
+  "hubert模型路径不可包含中文": "hubert模型路径不可包含中文",
+  "pth文件路径不可包含中文": "pth文件路径不可包含中文",
+  "index文件路径不可包含中文": "index文件路径不可包含中文",
+  "重载设备列表": "Dispositivo de recarga",
+  "音高算法": "algoritmo de detección de cabeceo",
+  "harvest进程数": "Número de procesos de CPU utilizados para el algoritmo de tono de cosecha"
+}

i18n/zh_CN.json

@@ -0,0 +1,124 @@
+{
+  "inference.save.button": "保存选项",
+  "inference.refresh_data.button": "刷新数据",
+  "inference.clear_data.button": "卸载模型",
+  "inference.load_model.button": "载入模型",
+  "inference.one_click.button": "一键式VC",
+  "inference.song.selectbox": "选择歌曲文件",
+  "inference.voice.selectbox": "选择语音模型",
+  "inference.split_vocals": "提取人声",
+  "inference.split_vocals.expander": "“提取人声”选项",
+  "inference.preprocess_model": "选择预处理模型",
+  "inference.uvr5_name": "选择声音提取模型",
+  "inference.device": "设备",
+  "inference.merge_type": "合并类型",
+  "inference.agg": "隔离声音的攻击性",
+  "inference.use_cache": "通过将结果保存到磁盘来缩短未来的处理时间",
+  "inference.convert_vocals": "更改人声",
+  "inference.convert_vocals.expander": "更改人声选项",
+  "inference.convert_vocals.expander.form_submit_button": "保存选项",
+  "inference.f0_up_key": "建议+12键用于男性到女性的转换， -12键用于女性到男性的转换。如果声音范围太大，声音失真，您也可以自行调整到适当的范围。",
+  "inference.f0_method": "选择音高提取算法（ “crepe” ：质量更好但GPU密集型， “rmvpe” ：质量最好，对GPU的要求很低）",
+  "inference.filter_radius": "如果> = 3 ：对收获的沥青结果应用中值过滤。该值表示过滤器半径，可以减少呼吸。",
+  "inference.index_rate": "搜索特征比率（控制重音强度，太高有伪影） ：",
+  "inference.resample_sr": "将后期处理中的输出音频重新采样到最终采样率。设置为0表示不重采样：",
+  "inference.rms_mix_rate": "调整体积信封缩放。越接近0 ，就越能模仿原声的音量。可以帮助掩盖噪音，并在设置相对较低时使音量听起来更自然。接近1的音量将更接近持续响亮的音量：",
+  "inference.protect": "保护无声辅音和呼吸声，以防止电子音乐中的撕裂等伪影。设置为0.5以禁用。减小值以增加保护，但可能会降低索引精度：",
+  "inference.download.button": "下载歌曲",
+  "training.preprocess_data.title": "过程数据",
+  "training.preprocess_data.text": "第1步：填写实验配置。实验数据存储在“日志”文件夹中，每个实验都有一个单独的文件夹。手动输入实验名称路径，其中包含实验配置、日志和训练好的模型文件。",
+  "training.exp_dir": "输入语音模型名称：",
+  "training.sr": "目标采样率：",
+  "training.if_f0": "模特是否有音调指导（唱歌需要，演讲可选） ：",
+  "training.version": "版本：",
+  "training.device": "器件：",
+  "training.n_threads": "用于音调提取和数据处理的CPU进程数量：",
+  "training.extract_features.title": "提取特征",
+  "training.extract_features.text": "第2步：自动遍历训练文件夹中可以解码为音频的所有文件，并执行切片归一化。在实验目录中生成2个wav文件夹。目前，仅支持单声道/扬声器培训。",
+  "training.preprocess_data.trainset_dir": "输入培训文件夹的路径：",
+  "training.sid": "请指定演讲者/歌手ID ：",
+  "training.preprocess_data.submit": "1. 进程",
+  "training.train_model.gpus": "输入以“-”分隔的GPU索引，例如0-1-2以使用GPU 0、1和2 ：",
+  "training.f0method": "选择音高提取算法（ “crepe” ：质量更好但GPU密集型， “rmvpe” ：质量最好，对GPU的要求很低）",
+  "training.extract_features.submit": "特征抽取",
+  "training.train_model.title": "火车模型",
+  "training.train_model.text": "第3步：填写培训设置，开始培训模型和索引",
+  "training.save_every_epoch": "保存频率（ save_every_epoch ） ：",
+  "training.total_epoch": "总培训纪元（ TOTAL_EPOCH ） ：",
+  "training.batch_size": "每个GPU的批处理大小：",
+  "training.if_save_latest": "仅保存最新的“.ckpt”文件以节省磁盘空间",
+  "training.if_save_best": "仅将最佳模型（最低总体损失）保存到“日志”文件夹",
+  "training.if_cache_gpu": "将所有训练集缓存到GPU内存。缓存小数据集（少于10分钟）可以加快训练速度，但缓存大数据集会占用大量GPU内存，可能无法提供太多速度提升",
+  "training.if_save_every_weights": "在每个保存点将一个小的最终模型保存到“日志”文件夹中",
+  "training.pretrained_G": "加载预先训练的基本模型G路径：",
+  "training.pretrained_D": "加载预训练基础模型D路径：",
+  "training.train_model.submit": "3.火车模型",
+  "training.train_index.submit": "4.列车特征索引",
+  "training.train_speaker.submit": "培训可选音箱嵌入",
+  "training.one_click": "一键培训",
+  "process.pids": "活动进程",
+  "process.kill_all_pids": "关闭所有流程",
+  "process.kill_one_pid": "停止",
+  "tts.inference": "文本到语音",
+  "tts.options": "TTS选项",
+  "tts.model.selectbox": "TTS模型",
+  "ckpt处理": "ckpt处理",
+  "模型融合, 可用于测试音色融合": "模型融合，可用于测试音色融合",
+  "A模型路径": "模型A的路径：",
+  "B模型路径": "通往模型B的路径：",
+  "A模型权重": "型号A的重量（ w ） ：",
+  "模型是否带音高指导": "模型是否具有变桨指导：",
+  "要置入的模型信息": "要放置的型号信息：",
+  "保存的模型名不带后缀": "保存的模型名称（不带扩展名） ：",
+  "模型版本型号": "模型架构版本：",
+  "融合": "Fusion",
+  "修改模型信息(仅支持weights文件夹下提取的小模型文件)": "修改模型信息（仅支持从“权重”文件夹中提取的小型模型文件）",
+  "模型路径": "模型路径：",
+  "要改的模型信息": "要修改的型号信息：",
+  "保存的文件名, 默认空为和源文件同名": "保存文件名（默认：与源文件相同） ：",
+  "修改": "修改",
+  "查看模型信息(仅支持weights文件夹下提取的小模型文件)": "查看模型信息（仅支持从“权重”文件夹中提取的小型模型文件）",
+  "查看": "查看",
+  "模型提取(输入logs文件夹下大文件模型路径),适用于训一半不想训了模型没有自动提取保存小文件模型,或者想测试中间模型的情况": "模型提取（在“日志”文件夹下输入大文件模型的路径）。如果要中途停止训练并手动提取和保存小模型文件，或者如果要测试中间模型，这很有用：",
+  "保存名": "保存名称",
+  "模型是否带音高指导,1是0否": "模型是否有音调指导（ 1 ：是， 0 ：否） ：",
+  "提取": "提取",
+  "Onnx导出": "导出Onnx",
+  "RVC模型路径": "RVC模型路径：",
+  "Onnx输出路径": "Onnx导出路径：",
+  "导出Onnx模型": "导出Onnx模型",
+  "常见问题解答": "FAQ（常见问题）",
+  "招募音高曲线前端编辑器": "为俯仰曲线招募前端编辑器",
+  "加开发群联系我xxxxx": "加入发展小组，并通过xxxxx与我联系",
+  "点击查看交流、问题反馈群号": "点击查看沟通和问题反馈组号",
+  "xxxxx": "xxxxx",
+  "加载模型": "加载模型",
+  "Hubert模型": "休伯特模型",
+  "选择.pth文件": "选择文件",
+  "选择.index文件": "选择.index文件",
+  "选择.npy文件": "选择文件",
+  "输入设备": "输入装置",
+  "输出设备": "輸出裝置",
+  "音频设备(请使用同种类驱动)": "音频设备（请使用相同类型的驱动程序）",
+  "响应阈值": "响应阈值",
+  "音调设置": "音调设置",
+  "Index Rate": "指数利率",
+  "常规设置": "常规设置",
+  "采样长度": "试样长度",
+  "淡入淡出长度": "淡出 长度",
+  "额外推理时长": "额外推理时间",
+  "输入降噪": "输入噪音降低",
+  "输出降噪": "输出噪声降低",
+  "性能设置": "性能设置",
+  "开始音频转换": "开始音频转换",
+  "停止音频转换": "停止音频转换",
+  "推理时间(ms):": "推理时间（毫秒） ：",
+  "请选择pth文件": "请选择pth文件",
+  "请选择index文件": "请选择index文件",
+  "hubert模型路径不可包含中文": "休伯特模型路径不可包含中文",
+  "pth文件路径不可包含中文": "pth文件路径不可包含中文",
+  "index文件路径不可包含中文": "index文件路径不可包含中文",
+  "重载设备列表": "重新加载设备列表",
+  "音高算法": "音调检测算法",
+  "harvest进程数": "用于收获音高算法的CPU进程数"
+}

lib/__init__.py

@@ -0,0 +1,118 @@
+from functools import lru_cache
+import os
+import shelve
+from contextlib import contextmanager
+import sys
+import weakref
+
+@lru_cache
+def get_cwd():
+    CWD = os.getcwd()
+    if CWD not in sys.path:
+        sys.path.append(CWD)
+    return CWD
+
+class ObjectNamespace(dict):
+    def __init__(self,**kwargs):
+        for k,v in kwargs.items():
+            if hasattr(v,"items"): self[k]=ObjectNamespace(**v)
+            else: self[k]=v
+
+    def __missing__(self, name: str): return ObjectNamespace()
+    def get(self, name: str, default_value=None): return self.__getitem__(name) if name in self.keys() else default_value
+    def __getattr__(self, name: str): return self.__getitem__(name) if name in self.keys() else ObjectNamespace()
+    def __getitem__(self, name: str):
+        value = super().__getitem__(name) # get the value from the parent class
+        if isinstance(value, weakref.ref): # check if the value is a weak reference
+            value = value() # call the weak reference object to get the referent
+        return value # return the referent or the original value
+    def __setattr__(self, name: str, value): return self.__setitem__(name, value)
+    def __delattr__(self, name: str): return self.__delitem__(name) if name in self.keys() else None
+    def __delitem__(self, name: str): return super().__delitem__(name) if name in self.keys() else None
+    def __setstate__(self, state):
+        for key in state: self[key]=state[key]
+    def __getstate__(self): return dict(**self)
+
+class PersistedDict:
+    # initialize the class with an optional filename and dict arguments
+    def __init__(self, filename=None, **data):
+        # store the filename as an attribute
+        self.filename = filename
+
+        for key, value in data.items():
+            # recursively convert the values to NestedDict
+            self.__setattr__(key, value)
+
+    # define a context manager to open and close the shelve file
+    @contextmanager
+    def open_shelf(self):
+        # if filename is given, open the shelve file
+        shelf = shelve.open(self.filename) if self.filename else {}
+
+        # yield the shelf as the resource
+        yield shelf
+        if hasattr(shelf,"close"):
+            # close the shelf when exiting the context
+            shelf.close()
+
+    # define a method to get the attribute value given a key
+    def __getattr__(self, key: str):
+        is_private = key.startswith("_") and key.endswith("_")
+        
+        # if the key is filename, set it as an attribute
+        if key == "filename" or is_private:
+            if key in self.__dict__: return self.__dict__[key]
+            else: return None
+
+        # use the context manager to open the shelve file
+        with self.open_shelf() as shelf:
+            # if the key exists in the shelve file, return the value
+            # return getattr(shelf, key, None)
+            if key in shelf:
+                return shelf[key]
+            # else, return None
+            else:
+                return None
+
+    # define a method to set the attribute value given a key
+    def __setattr__(self, key, value):
+        # if the key is filename, set it as an attribute
+        if key == "filename":
+            self.__dict__[key] = value
+        # else, use the context manager to open the shelve file
+        else:
+            with self.open_shelf() as shelf:
+                # store the value in the shelve file
+                print(f"{key}={value}")
+                shelf[key] = value
+
+    # define a method to represent the class as a dict
+    def __repr__(self):
+        # initialize an empty dict
+        result = {}
+        # use the context manager to open the shelve file
+        with self.open_shelf() as shelf:
+            # loop through the keys in the shelve file
+            for key in shelf.keys():
+                # add the key and value to the result
+                result[key] = shelf[key]
+        # return the result
+        return str(result)
+    
+    def __setitem__(self, key, value): self.__setattr__(key, value)
+    def __getitem__(self, key): self.__getattr__(key)
+    def __lt__(self,_): return False
+    def __eq__(self,other):
+        if hasattr(other,"filename"): return self.filename==other.filename
+        else: return False
+    def __call__(self,*args,**kwargs):
+        print(f"{args=}, {kwargs=}")
+        return str(self)
+
+BASE_DIR = get_cwd()
+BASE_MODELS_DIR = os.path.join(BASE_DIR,"models")
+SONG_DIR = os.path.join(BASE_DIR,"songs")
+BASE_CACHE_DIR = os.path.join(BASE_DIR,".cache")
+DATASETS_DIR = os.path.join(BASE_DIR,"datasets")
+LOG_DIR = os.path.join(BASE_DIR,"logs")
+OUTPUT_DIR = os.path.join(BASE_DIR,"output")

lib/audio.py

@@ -0,0 +1,210 @@
+import base64
+import io
+import os
+from typing import Union
+import zlib
+from .utils import get_hash, get_merge_func
+import numpy as np
+import librosa
+import soundfile as sf
+import ffmpeg
+
+MAX_INT16 = 32768
+SUPPORTED_AUDIO = ["mp3","flac","wav"] # ogg breaks soundfile
+OUTPUT_CHANNELS = ["mono", "stereo"]
+AUTOTUNE_NOTES = np.array([
+    65.41, 69.30, 73.42, 77.78, 82.41, 87.31,
+    92.50, 98.00, 103.83, 110.00, 116.54, 123.47,
+    130.81, 138.59, 146.83, 155.56, 164.81, 174.61,
+    185.00, 196.00, 207.65, 220.00, 233.08, 246.94,
+    261.63, 277.18, 293.66, 311.13, 329.63, 349.23,
+    369.99, 392.00, 415.30, 440.00, 466.16, 493.88,
+    523.25, 554.37, 587.33, 622.25, 659.25, 698.46,
+    739.99, 783.99, 830.61, 880.00, 932.33, 987.77,
+    1046.50, 1108.73, 1174.66, 1244.51, 1318.51, 1396.91,
+    1479.98, 1567.98, 1661.22, 1760.00, 1864.66, 1975.53,
+    2093.00, 2217.46, 2349.32, 2489.02, 2637.02, 2793.83,
+    2959.96, 3135.96, 3322.44, 3520.00, 3729.31, 3951.07
+])
+
+def load_audio(file, sr, **kwargs):
+    try:
+        # https://github.com/openai/whisper/blob/main/whisper/audio.py#L26
+        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
+        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
+        file = (
+            file.strip(" ").strip('"').strip("\n").strip('"').strip(" ")
+        )  # 防止小白拷路径头尾带了空格和"和回车
+        out, _ = (
+            ffmpeg.input(file, threads=0)
+            .output("-", format="f32le", acodec="pcm_f32le", ac=1, ar=sr)
+            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
+        )
+    except Exception as e:
+        raise RuntimeError(f"Failed to load audio: {e}")
+
+    return remix_audio((np.frombuffer(out, np.float32).flatten(), sr),**kwargs)
+
+def remix_audio(input_audio,target_sr=None,norm=False,to_int16=False,resample=False,axis=0,merge_type=None,**kwargs):
+    audio = np.array(input_audio[0],dtype="float32")
+    if target_sr is None: target_sr=input_audio[1]
+
+    print(f"before remix: shape={audio.shape}, max={audio.max()}, min={audio.min()}, mean={audio.mean()} sr={input_audio[1]}")
+    if resample or input_audio[1]!=target_sr:
+        audio = librosa.resample(np.array(input_audio[0],dtype="float32"),orig_sr=input_audio[1],target_sr=target_sr,**kwargs)
+    
+    if audio.ndim>1:
+        merge_func = get_merge_func(merge_type)
+        audio=merge_func(audio,axis=axis)
+    if norm: audio = librosa.util.normalize(audio,axis=axis)
+
+    audio_max = np.abs(audio).max()/.99
+    if audio_max > 1: audio = audio / audio_max
+        
+    if to_int16: audio = np.clip(audio * MAX_INT16, a_min=-MAX_INT16+1, a_max=MAX_INT16-1).astype("int16")
+    print(f"after remix: shape={audio.shape}, max={audio.max()}, min={audio.min()}, mean={audio.mean()}, sr={target_sr}")
+
+    return audio, target_sr
+
+def load_input_audio(fname,sr=None,**kwargs):
+    if sr is None: sr=44100
+    audio, sr = load_audio(fname, sr, **kwargs)
+    print(f"loading sound {fname=} {audio.ndim=} {audio.max()=} {audio.min()=} {audio.dtype=} {sr=}")
+    return audio, sr
+   
+def save_input_audio(fname,input_audio,sr=None,to_int16=False,to_stereo=False):
+    print(f"saving sound to {fname}")
+    os.makedirs(os.path.dirname(fname),exist_ok=True)
+    audio=np.array(input_audio[0],dtype="float32")
+
+    if to_int16:
+        audio_max = np.abs(audio).max()/.99
+        if audio_max > 1: audio = audio / audio_max
+        audio = np.clip(audio * MAX_INT16, a_min=-MAX_INT16+1, a_max=MAX_INT16-1)
+
+    if to_stereo and audio.ndim<2: audio=np.stack([audio,audio],axis=-1)
+    
+    try:
+        if audio.ndim>1 and audio.shape[0]<audio.shape[1]: audio=audio.T # soundfile expects data in frames x channels
+        print(f"{audio.shape=}")
+        sf.write(fname, audio.astype("int16" if np.abs(audio).max()>1 else "float32"), sr if sr else input_audio[1])
+        return f"File saved to ${fname}"
+    except Exception as e:
+        return f"failed to save audio: {e}"
+    
+def audio_to_bytes(audio,sr,target_sr=None,to_int16=False,to_stereo=False,format="WAV"):
+    
+    with io.BytesIO() as bytes_io:
+        audio=np.array(audio,dtype="float32")
+
+        if to_int16:
+            audio_max = np.abs(audio).max()/.99
+            if audio_max > 1: audio = audio / audio_max
+            audio = np.clip(audio * MAX_INT16, a_min=-MAX_INT16+1, a_max=MAX_INT16-1)
+
+        if to_stereo and audio.ndim<2: audio=np.stack([audio,audio],axis=-1)
+
+        if audio.ndim>1 and audio.shape[0]<audio.shape[1]: audio=audio.T # soundfile expects data in frames x channels
+        print(f"{audio.shape=}")
+        samplerate = sr if target_sr is None else target_sr
+        sf.write(bytes_io, audio.astype("int16" if np.abs(audio).max()>1 else "float32"), samplerate=samplerate,format=format)
+        bytes_io.seek(0)
+        return bytes_io.read()
+
+def bytes_to_audio(data: Union[io.BytesIO,bytes],**kwargs):
+    with io.BytesIO(data) as bytes_io:
+        audio, sr = sf.read(bytes_io,**kwargs)
+        if audio.ndim>1:
+            if audio.shape[-1]<audio.shape[0]: # is channel-last format
+                audio = audio.T # transpose to channels-first
+        return audio, sr
+
+def bytes2audio(data: str):
+    try:
+        # Split the suffixed data by the colon
+        dtype,data,shape,sr = data.split(":")
+
+        # Get the data, the dtype, and the shape from the split data
+        shape = tuple(map(int, shape.split(",")))
+        sr=int(sr)
+
+        # Decode the data using base64
+        decoded_data = base64.b64decode(data)
+
+        # Decompress the decoded data using zlib
+        decompressed_data = zlib.decompress(decoded_data)
+
+        # Convert the decompressed data to a numpy array with the given dtype
+        arr = np.frombuffer(decompressed_data, dtype=dtype)
+
+        # Reshape the array to the original shape
+        arr = arr.reshape(shape)
+        return arr, sr
+    except Exception as e:
+        print(e)
+    return None
+
+def audio2bytes(audio: np.array, sr: int):
+    try:
+        # Get the dtype, the shape, and the data of the array
+        dtype = audio.dtype.name
+        shape = audio.shape
+        data = audio.tobytes()
+
+        # Compress the data using zlib
+        compressed_data = zlib.compress(data)
+
+        # Encode the compressed data using base64
+        encoded_data = base64.b64encode(compressed_data)
+
+        # Add a suffix with the dtype and the shape to the encoded data
+        suffixed_data = ":".join([dtype,encoded_data.decode(),",".join(map(str, shape)),str(sr)])
+        return suffixed_data
+    except Exception as e:
+        print(e)
+    return ""
+
+def pad_audio(*audios,axis=0):
+    maxlen = max(len(a) if a is not None else 0 for a in audios)
+    if maxlen>0:
+        stack = librosa.util.stack([librosa.util.pad_center(data=a,size=maxlen) for a in audios if a is not None],axis=axis)
+        return stack
+    else: return np.stack(audios,axis=axis)
+
+def merge_audio(audio1,audio2,sr=40000,**kwargs):
+    print(f"merging audio audio1={audio1[0].shape,audio1[1]} audio2={audio2[0].shape,audio2[1]} sr={sr}")
+    if sr is None: sr=min(audio1[-1],audio2[-1])
+    m1,_=remix_audio(audio1,target_sr=sr,axis=0,**kwargs)
+    m2,_=remix_audio(audio2,target_sr=sr,axis=0,**kwargs)
+    
+    mixed = pad_audio(m1,m2,axis=0)
+
+    return remix_audio((mixed,sr),axis=0,**kwargs)
+
+def autotune_f0(f0, threshold=0.):
+    # autotuned_f0 = []
+    # for freq in f0:
+        # closest_notes = [x for x in self.note_dict if abs(x - freq) == min(abs(n - freq) for n in self.note_dict)]
+        # autotuned_f0.append(random.choice(closest_notes))
+    # for note in self.note_dict:
+    #     closest_notes = np.where((f0 - note)/note<.05,f0,note)
+    print("autotuning f0 using note_dict...")
+
+    autotuned_f0 = []
+    # Loop through each value in array1
+    for freq in f0:
+        # Find the absolute difference between x and each value in array2
+        diff = np.abs(AUTOTUNE_NOTES - freq)
+        # Find the index of the minimum difference
+        idx = np.argmin(diff)
+        # Find the corresponding value in array2
+        y = AUTOTUNE_NOTES[idx]
+        # Check if the difference is less than threshold
+        if diff[idx] < threshold:
+            # Keep the value in array1
+            autotuned_f0.append(freq)
+        else:
+            # Use the nearest value in array2
+            autotuned_f0.append(y)
+    # Return the result as a numpy array
+    return np.array(autotuned_f0, dtype="float32")

lib/infer_pack/attentions.py

@@ -0,0 +1,413 @@
+import math
+from .commons import convert_pad_shape, subsequent_mask
+import torch
+from torch import nn
+from torch.nn import functional as F
+from .modules import LayerNorm
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        window_size=10,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(LayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                    causal=True,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, h, h_mask):
+        """
+        x: decoder input
+        h: encoder output
+        """
+        self_attn_mask = subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        n_heads,
+        p_dropout=0.0,
+        window_size=None,
+        heads_share=True,
+        block_length=None,
+        proximal_bias=False,
+        proximal_init=False,
+    ):
+        super().__init__()
+        assert channels % n_heads == 0
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
+
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                self.conv_k.bias.copy_(self.conv_q.bias)
+
+    def forward(self, x, c, attn_mask=None):
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask=None):
+        # reshape [b, d, t] -> [b, n_h, t, d_k]
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            )
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)  # [b, n_h, t_t, t_s]
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_v, t_s
+            )
+            output = output + self._matmul_with_relative_values(
+                relative_weights, value_relative_embeddings
+            )
+        output = (
+            output.transpose(2, 3).contiguous().view(b, d, t_t)
+        )  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+        return output, p_attn
+
+    def _matmul_with_relative_values(self, x, y):
+        """
+        x: [b, h, l, m]
+        y: [h or 1, m, d]
+        ret: [b, h, l, d]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
+
+    def _matmul_with_relative_keys(self, x, y):
+        """
+        x: [b, h, l, d]
+        y: [h or 1, m, d]
+        ret: [b, h, l, m]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        return ret
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        max_relative_position = 2 * self.window_size + 1
+        # Pad first before slice to avoid using cond ops.
+        pad_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
+
+    def _relative_position_to_absolute_position(self, x):
+        """
+        x: [b, h, l, 2*l-1]
+        ret: [b, h, l, l]
+        """
+        batch, heads, length, _ = x.size()
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
+
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(
+            x_flat, convert_pad_shape([[0, 0], [0, 0], [0, length - 1]])
+        )
+
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
+            :, :, :length, length - 1 :
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        """
+        x: [b, h, l, l]
+        ret: [b, h, l, 2*l-1]
+        """
+        batch, heads, length, _ = x.size()
+        # padd along column
+        x = F.pad(
+            x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
+        )
+        x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
+        # add 0's in the beginning that will skew the elements after reshape
+        x_flat = F.pad(x_flat, convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length):
+        """Bias for self-attention to encourage attention to close positions.
+        Args:
+          length: an integer scalar.
+        Returns:
+          a Tensor with shape [1, 1, length, length]
+        """
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        filter_channels,
+        kernel_size,
+        p_dropout=0.0,
+        activation=None,
+        causal=False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.activation = activation
+        self.causal = causal
+
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
+
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding(x * x_mask))
+        if self.activation == "gelu":
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
+
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, convert_pad_shape(padding))
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, convert_pad_shape(padding))
+        return x

lib/infer_pack/commons.py

@@ -0,0 +1,170 @@
+import math
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def convert_pad_shape(pad_shape):
+    l = pad_shape[::-1]
+    pad_shape = [item for sublist in l for item in sublist]
+    return pad_shape
+
+def intersperse(lst, item):
+  result = [item] * (len(lst) * 2 + 1)
+  result[1::2] = lst
+  return result
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+    """KL(P||Q)"""
+    kl = (logs_q - logs_p) - 0.5
+    kl += (
+        0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
+    )
+    return kl
+
+
+def rand_gumbel(shape):
+    """Sample from the Gumbel distribution, protect from overflows."""
+    uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+    return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+    g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+    return g
+
+
+def slice_segments(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = ids_str[i]
+        idx_end = idx_str + segment_size
+        ret[i] = x[i, :, idx_str:idx_end]
+    return ret
+
+
+def slice_segments2(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = ids_str[i]
+        idx_end = idx_str + segment_size
+        ret[i] = x[i, idx_str:idx_end]
+    return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
+    position = torch.arange(length, dtype=torch.float)
+    num_timescales = channels // 2
+    log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
+        num_timescales - 1
+    )
+    inv_timescales = min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
+    )
+    scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+    signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+    signal = F.pad(signal, [0, 0, 0, channels % 2])
+    signal = signal.view(1, channels, length)
+    return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def convert_pad_shape(pad_shape):
+    l = pad_shape[::-1]
+    pad_shape = [item for sublist in l for item in sublist]
+    return pad_shape
+
+
+def shift_1d(x):
+    x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+    return x
+
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def generate_path(duration, mask):
+    """
+    duration: [b, 1, t_x]
+    mask: [b, 1, t_y, t_x]
+    """
+    device = duration.device
+
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm

lib/infer_pack/models_onnx.py

@@ -0,0 +1,817 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+from ..infer_pack import modules
+from ..infer_pack import attentions
+from ..infer_pack import commons
+from ..infer_pack.commons import init_weights, get_padding
+from torch.nn import Conv1d, ConvTranspose1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+from ..infer_pack.commons import init_weights
+import numpy as np
+from ..infer_pack import commons
+
+class TextEncoder256(nn.Module):
+    def __init__(
+        self,
+        out_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        f0=True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.emb_phone = nn.Linear(256, hidden_channels)
+        self.lrelu = nn.LeakyReLU(0.1, inplace=True)
+        if f0 == True:
+            self.emb_pitch = nn.Embedding(256, hidden_channels)  # pitch 256
+        self.encoder = attentions.Encoder(
+            hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, phone, pitch, lengths):
+        if pitch == None:
+            x = self.emb_phone(phone)
+        else:
+            x = self.emb_phone(phone) + self.emb_pitch(pitch)
+        x = x * math.sqrt(self.hidden_channels)  # [b, t, h]
+        x = self.lrelu(x)
+        x = torch.transpose(x, 1, -1)  # [b, h, t]
+        x_mask = torch.unsqueeze(commons.sequence_mask(lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.encoder(x * x_mask, x_mask)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return m, logs, x_mask
+
+
+class TextEncoder768(nn.Module):
+    def __init__(
+        self,
+        out_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        f0=True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.emb_phone = nn.Linear(768, hidden_channels)
+        self.lrelu = nn.LeakyReLU(0.1, inplace=True)
+        if f0 == True:
+            self.emb_pitch = nn.Embedding(256, hidden_channels)  # pitch 256
+        self.encoder = attentions.Encoder(
+            hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, phone, pitch, lengths):
+        if pitch == None:
+            x = self.emb_phone(phone)
+        else:
+            x = self.emb_phone(phone) + self.emb_pitch(pitch)
+        x = x * math.sqrt(self.hidden_channels)  # [b, t, h]
+        x = self.lrelu(x)
+        x = torch.transpose(x, 1, -1)  # [b, h, t]
+        x_mask = torch.unsqueeze(commons.sequence_mask(lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.encoder(x * x_mask, x_mask)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return m, logs, x_mask
+
+
+class ResidualCouplingBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        n_flows=4,
+        gin_channels=0,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ResidualCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    dilation_rate,
+                    n_layers,
+                    gin_channels=gin_channels,
+                    mean_only=True,
+                )
+            )
+            self.flows.append(modules.Flip())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+    def remove_weight_norm(self):
+        for i in range(self.n_flows):
+            self.flows[i * 2].remove_weight_norm()
+
+
+class PosteriorEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, g=None):
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+    def remove_weight_norm(self):
+        self.enc.remove_weight_norm()
+
+
+class Generator(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        gin_channels=0,
+    ):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(self, x, g=None):
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+
+
+class SineGen(torch.nn.Module):
+    """Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(np.pi) or cos(0)
+    """
+
+    def __init__(
+        self,
+        samp_rate,
+        harmonic_num=0,
+        sine_amp=0.1,
+        noise_std=0.003,
+        voiced_threshold=0,
+        flag_for_pulse=False,
+    ):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = self.harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = torch.ones_like(f0)
+        uv = uv * (f0 > self.voiced_threshold)
+        return uv
+
+    def forward(self, f0, upp):
+        """sine_tensor, uv = forward(f0)
+        input F0: tensor(batchsize=1, length, dim=1)
+                  f0 for unvoiced steps should be 0
+        output sine_tensor: tensor(batchsize=1, length, dim)
+        output uv: tensor(batchsize=1, length, 1)
+        """
+        with torch.no_grad():
+            f0 = f0[:, None].transpose(1, 2)
+            f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim, device=f0.device)
+            # fundamental component
+            f0_buf[:, :, 0] = f0[:, :, 0]
+            for idx in np.arange(self.harmonic_num):
+                f0_buf[:, :, idx + 1] = f0_buf[:, :, 0] * (
+                    idx + 2
+                )  # idx + 2: the (idx+1)-th overtone, (idx+2)-th harmonic
+            rad_values = (f0_buf / self.sampling_rate) % 1  ###%1意味着n_har的乘积无法后处理优化
+            rand_ini = torch.rand(
+                f0_buf.shape[0], f0_buf.shape[2], device=f0_buf.device
+            )
+            rand_ini[:, 0] = 0
+            rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+            tmp_over_one = torch.cumsum(rad_values, 1)  # % 1  #####%1意味着后面的cumsum无法再优化
+            tmp_over_one *= upp
+            tmp_over_one = F.interpolate(
+                tmp_over_one.transpose(2, 1),
+                scale_factor=upp,
+                mode="linear",
+                align_corners=True,
+            ).transpose(2, 1)
+            rad_values = F.interpolate(
+                rad_values.transpose(2, 1), scale_factor=upp, mode="nearest"
+            ).transpose(
+                2, 1
+            )  #######
+            tmp_over_one %= 1
+            tmp_over_one_idx = (tmp_over_one[:, 1:, :] - tmp_over_one[:, :-1, :]) < 0
+            cumsum_shift = torch.zeros_like(rad_values)
+            cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+            sine_waves = torch.sin(
+                torch.cumsum(rad_values + cumsum_shift, dim=1) * 2 * np.pi
+            )
+            sine_waves = sine_waves * self.sine_amp
+            uv = self._f02uv(f0)
+            uv = F.interpolate(
+                uv.transpose(2, 1), scale_factor=upp, mode="nearest"
+            ).transpose(2, 1)
+            noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+            noise = noise_amp * torch.randn_like(sine_waves)
+            sine_waves = sine_waves * uv + noise
+        return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+    """SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+
+    def __init__(
+        self,
+        sampling_rate,
+        harmonic_num=0,
+        sine_amp=0.1,
+        add_noise_std=0.003,
+        voiced_threshod=0,
+        is_half=True,
+    ):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+        self.is_half = is_half
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(
+            sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshod
+        )
+
+        # to merge source harmonics into a single excitation
+        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = torch.nn.Tanh()
+
+    def forward(self, x, upp=None):
+        sine_wavs, uv, _ = self.l_sin_gen(x, upp)
+        if self.is_half:
+            sine_wavs = sine_wavs.half()
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+        return sine_merge, None, None  # noise, uv
+
+
+class GeneratorNSF(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        gin_channels,
+        sr,
+        is_half=False,
+    ):
+        super(GeneratorNSF, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates))
+        self.m_source = SourceModuleHnNSF(
+            sampling_rate=sr, harmonic_num=0, is_half=is_half
+        )
+        self.noise_convs = nn.ModuleList()
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            c_cur = upsample_initial_channel // (2 ** (i + 1))
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+            if i + 1 < len(upsample_rates):
+                stride_f0 = np.prod(upsample_rates[i + 1 :])
+                self.noise_convs.append(
+                    Conv1d(
+                        1,
+                        c_cur,
+                        kernel_size=stride_f0 * 2,
+                        stride=stride_f0,
+                        padding=stride_f0 // 2,
+                    )
+                )
+            else:
+                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+        self.upp = np.prod(upsample_rates)
+
+    def forward(self, x, f0, g=None):
+        har_source, noi_source, uv = self.m_source(f0, self.upp)
+        har_source = har_source.transpose(1, 2)
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            x_source = self.noise_convs[i](har_source)
+            x = x + x_source
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+
+
+sr2sr = {
+    "32k": 32000,
+    "40k": 40000,
+    "48k": 48000,
+}
+
+
+class SynthesizerTrnMsNSFsidM(nn.Module):
+    def __init__(
+        self,
+        spec_channels,
+        segment_size,
+        inter_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        spk_embed_dim,
+        gin_channels,
+        sr,
+        version,
+        **kwargs
+    ):
+        super().__init__()
+        if type(sr) == type("strr"):
+            sr = sr2sr[sr]
+        self.spec_channels = spec_channels
+        self.inter_channels = inter_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.resblock = resblock
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.upsample_rates = upsample_rates
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.segment_size = segment_size
+        self.gin_channels = gin_channels
+        # self.hop_length = hop_length#
+        self.spk_embed_dim = spk_embed_dim
+        if version == "v1":
+            self.enc_p = TextEncoder256(
+                inter_channels,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+            )
+        else:
+            self.enc_p = TextEncoder768(
+                inter_channels,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+            )
+        self.dec = GeneratorNSF(
+            inter_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+            sr=sr,
+            is_half=kwargs["is_half"],
+        )
+        self.enc_q = PosteriorEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+        self.flow = ResidualCouplingBlock(
+            inter_channels, hidden_channels, 5, 1, 3, gin_channels=gin_channels
+        )
+        self.emb_g = nn.Embedding(self.spk_embed_dim, gin_channels)
+        self.speaker_map = None
+       
+
+    def remove_weight_norm(self):
+        self.dec.remove_weight_norm()
+        self.flow.remove_weight_norm()
+        self.enc_q.remove_weight_norm()
+
+    def construct_spkmixmap(self, n_speaker):
+        self.speaker_map = torch.zeros((n_speaker, 1, 1, self.gin_channels))
+        for i in range(n_speaker):
+            self.speaker_map[i] = self.emb_g(torch.LongTensor([[i]]))
+        self.speaker_map = self.speaker_map.unsqueeze(0)
+
+    def forward(self, phone, phone_lengths, pitch, nsff0, g, rnd, max_len=None):
+        if self.speaker_map is not None:  # [N, S]  *  [S, B, 1, H]
+            g = g.reshape((g.shape[0], g.shape[1], 1, 1, 1))  # [N, S, B, 1, 1]
+            g = g * self.speaker_map  # [N, S, B, 1, H]
+            g = torch.sum(g, dim=1)  # [N, 1, B, 1, H]
+            g = g.transpose(0, -1).transpose(0, -2).squeeze(0)  # [B, H, N]
+        else:
+            g = g.unsqueeze(0)
+            g = self.emb_g(g).transpose(1, 2)
+
+        m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
+        z_p = (m_p + torch.exp(logs_p) * rnd) * x_mask
+        z = self.flow(z_p, x_mask, g=g, reverse=True)
+        o = self.dec((z * x_mask)[:, :, :max_len], nsff0, g=g)
+        return o
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self, use_spectral_norm=False):
+        super(MultiPeriodDiscriminator, self).__init__()
+        periods = [2, 3, 5, 7, 11, 17]
+        # periods = [3, 5, 7, 11, 17, 23, 37]
+
+        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+        discs = discs + [
+            DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods
+        ]
+        self.discriminators = nn.ModuleList(discs)
+
+    def forward(self, y, y_hat):
+        y_d_rs = []  #
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            # for j in range(len(fmap_r)):
+            #     print(i,j,y.shape,y_hat.shape,fmap_r[j].shape,fmap_g[j].shape)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class MultiPeriodDiscriminatorV2(torch.nn.Module):
+    def __init__(self, use_spectral_norm=False):
+        super(MultiPeriodDiscriminatorV2, self).__init__()
+        # periods = [2, 3, 5, 7, 11, 17]
+        periods = [2, 3, 5, 7, 11, 17, 23, 37]
+
+        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+        discs = discs + [
+            DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods
+        ]
+        self.discriminators = nn.ModuleList(discs)
+
+    def forward(self, y, y_hat):
+        y_d_rs = []  #
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            # for j in range(len(fmap_r)):
+            #     print(i,j,y.shape,y_hat.shape,fmap_r[j].shape,fmap_g[j].shape)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorS(torch.nn.Module):
+    def __init__(self, use_spectral_norm=False):
+        super(DiscriminatorS, self).__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(Conv1d(1, 16, 15, 1, padding=7)),
+                norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+                norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+                norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+                norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+                norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x):
+        fmap = []
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        self.use_spectral_norm = use_spectral_norm
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(
+                    Conv2d(
+                        1,
+                        32,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        32,
+                        128,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        128,
+                        512,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        512,
+                        1024,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        1024,
+                        1024,
+                        (kernel_size, 1),
+                        1,
+                        padding=(get_padding(kernel_size, 1), 0),
+                    )
+                ),
+            ]
+        )
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x):
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap

lib/infer_pack/modules.py

@@ -0,0 +1,518 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from torch.nn import Conv1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from .commons import fused_add_tanh_sigmoid_multiply, init_weights, get_padding
+from .transforms import piecewise_rational_quadratic_transform
+
+
+LRELU_SLOPE = 0.1
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        out_channels,
+        kernel_size,
+        n_layers,
+        p_dropout,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+        assert n_layers > 1, "Number of layers should be larger than 0."
+
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.conv_layers.append(
+            nn.Conv1d(
+                in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
+            )
+        )
+        self.norm_layers.append(LayerNorm(hidden_channels))
+        self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
+        for _ in range(n_layers - 1):
+            self.conv_layers.append(
+                nn.Conv1d(
+                    hidden_channels,
+                    hidden_channels,
+                    kernel_size,
+                    padding=kernel_size // 2,
+                )
+            )
+            self.norm_layers.append(LayerNorm(hidden_channels))
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask):
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
+
+
+class DDSConv(nn.Module):
+    """
+    Dialted and Depth-Separable Convolution
+    """
+
+    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+
+        self.drop = nn.Dropout(p_dropout)
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(n_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(LayerNorm(channels))
+            self.norms_2.append(LayerNorm(channels))
+
+    def forward(self, x, x_mask, g=None):
+        if g is not None:
+            x = x + g
+        for i in range(self.n_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.drop(y)
+            x = x + y
+        return x * x_mask
+
+
+class WN(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+        p_dropout=0,
+    ):
+        super(WN, self).__init__()
+        assert kernel_size % 2 == 1
+        self.hidden_channels = hidden_channels
+        self.kernel_size = (kernel_size,)
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = p_dropout
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
+
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(
+                gin_channels, 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
+
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, x, x_mask, g=None, **kwargs):
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
+            acts = self.drop(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+        self.convs.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class Log(nn.Module):
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+            logdet = torch.sum(-y, [1, 2])
+            return y, logdet
+        else:
+            x = torch.exp(x) * x_mask
+            return x
+
+
+class Flip(nn.Module):
+    def forward(self, x, *args, reverse=False, **kwargs):
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
+
+
+class ElementwiseAffine(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.m = nn.Parameter(torch.zeros(channels, 1))
+        self.logs = nn.Parameter(torch.zeros(channels, 1))
+
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = self.m + torch.exp(self.logs) * x
+            y = y * x_mask
+            logdet = torch.sum(self.logs * x_mask, [1, 2])
+            return y, logdet
+        else:
+            x = (x - self.m) * torch.exp(-self.logs) * x_mask
+            return x
+
+
+class ResidualCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        p_dropout=0,
+        gin_channels=0,
+        mean_only=False,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=p_dropout,
+            gin_channels=gin_channels,
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+    def remove_weight_norm(self):
+        self.enc.remove_weight_norm()
+
+
+class ConvFlow(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        filter_channels,
+        kernel_size,
+        n_layers,
+        num_bins=10,
+        tail_bound=5.0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.half_channels = in_channels // 2
+
+        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+        self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
+        self.proj = nn.Conv1d(
+            filter_channels, self.half_channels * (num_bins * 3 - 1), 1
+        )
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
+
+        b, c, t = x0.shape
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)  # [b, cx?, t] -> [b, c, t, ?]
+
+        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
+            self.filter_channels
+        )
+        unnormalized_derivatives = h[..., 2 * self.num_bins :]
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x

lib/infer_pack/modules/F0Predictor/DioF0Predictor.py

@@ -0,0 +1,90 @@
+from lib.infer_pack.modules.F0Predictor.F0Predictor import F0Predictor
+import pyworld
+import numpy as np
+
+
+class DioF0Predictor(F0Predictor):
+    def __init__(self, hop_length=512, f0_min=50, f0_max=1100, sampling_rate=44100):
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.sampling_rate = sampling_rate
+
+    def interpolate_f0(self, f0):
+        """
+        对F0进行插值处理
+        """
+
+        data = np.reshape(f0, (f0.size, 1))
+
+        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
+        vuv_vector[data > 0.0] = 1.0
+        vuv_vector[data <= 0.0] = 0.0
+
+        ip_data = data
+
+        frame_number = data.size
+        last_value = 0.0
+        for i in range(frame_number):
+            if data[i] <= 0.0:
+                j = i + 1
+                for j in range(i + 1, frame_number):
+                    if data[j] > 0.0:
+                        break
+                if j < frame_number - 1:
+                    if last_value > 0.0:
+                        step = (data[j] - data[i - 1]) / float(j - i)
+                        for k in range(i, j):
+                            ip_data[k] = data[i - 1] + step * (k - i + 1)
+                    else:
+                        for k in range(i, j):
+                            ip_data[k] = data[j]
+                else:
+                    for k in range(i, frame_number):
+                        ip_data[k] = last_value
+            else:
+                ip_data[i] = data[i]  # 这里可能存在一个没有必要的拷贝
+                last_value = data[i]
+
+        return ip_data[:, 0], vuv_vector[:, 0]
+
+    def resize_f0(self, x, target_len):
+        source = np.array(x)
+        source[source < 0.001] = np.nan
+        target = np.interp(
+            np.arange(0, len(source) * target_len, len(source)) / target_len,
+            np.arange(0, len(source)),
+            source,
+        )
+        res = np.nan_to_num(target)
+        return res
+
+    def compute_f0(self, wav, p_len=None):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        f0, t = pyworld.dio(
+            wav.astype(np.double),
+            fs=self.sampling_rate,
+            f0_floor=self.f0_min,
+            f0_ceil=self.f0_max,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
+        for index, pitch in enumerate(f0):
+            f0[index] = round(pitch, 1)
+        return self.interpolate_f0(self.resize_f0(f0, p_len))[0]
+
+    def compute_f0_uv(self, wav, p_len=None):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        f0, t = pyworld.dio(
+            wav.astype(np.double),
+            fs=self.sampling_rate,
+            f0_floor=self.f0_min,
+            f0_ceil=self.f0_max,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
+        for index, pitch in enumerate(f0):
+            f0[index] = round(pitch, 1)
+        return self.interpolate_f0(self.resize_f0(f0, p_len))

lib/infer_pack/modules/F0Predictor/F0Predictor.py

@@ -0,0 +1,16 @@
+class F0Predictor(object):
+    def compute_f0(self, wav, p_len):
+        """
+        input: wav:[signal_length]
+               p_len:int
+        output: f0:[signal_length//hop_length]
+        """
+        pass
+
+    def compute_f0_uv(self, wav, p_len):
+        """
+        input: wav:[signal_length]
+               p_len:int
+        output: f0:[signal_length//hop_length],uv:[signal_length//hop_length]
+        """
+        pass

lib/infer_pack/modules/F0Predictor/HarvestF0Predictor.py

@@ -0,0 +1,86 @@
+from lib.infer_pack.modules.F0Predictor.F0Predictor import F0Predictor
+import pyworld
+import numpy as np
+
+
+class HarvestF0Predictor(F0Predictor):
+    def __init__(self, hop_length=512, f0_min=50, f0_max=1100, sampling_rate=44100):
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.sampling_rate = sampling_rate
+
+    def interpolate_f0(self, f0):
+        """
+        对F0进行插值处理
+        """
+
+        data = np.reshape(f0, (f0.size, 1))
+
+        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
+        vuv_vector[data > 0.0] = 1.0
+        vuv_vector[data <= 0.0] = 0.0
+
+        ip_data = data
+
+        frame_number = data.size
+        last_value = 0.0
+        for i in range(frame_number):
+            if data[i] <= 0.0:
+                j = i + 1
+                for j in range(i + 1, frame_number):
+                    if data[j] > 0.0:
+                        break
+                if j < frame_number - 1:
+                    if last_value > 0.0:
+                        step = (data[j] - data[i - 1]) / float(j - i)
+                        for k in range(i, j):
+                            ip_data[k] = data[i - 1] + step * (k - i + 1)
+                    else:
+                        for k in range(i, j):
+                            ip_data[k] = data[j]
+                else:
+                    for k in range(i, frame_number):
+                        ip_data[k] = last_value
+            else:
+                ip_data[i] = data[i]  # 这里可能存在一个没有必要的拷贝
+                last_value = data[i]
+
+        return ip_data[:, 0], vuv_vector[:, 0]
+
+    def resize_f0(self, x, target_len):
+        source = np.array(x)
+        source[source < 0.001] = np.nan
+        target = np.interp(
+            np.arange(0, len(source) * target_len, len(source)) / target_len,
+            np.arange(0, len(source)),
+            source,
+        )
+        res = np.nan_to_num(target)
+        return res
+
+    def compute_f0(self, wav, p_len=None):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        f0, t = pyworld.harvest(
+            wav.astype(np.double),
+            fs=self.hop_length,
+            f0_ceil=self.f0_max,
+            f0_floor=self.f0_min,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.fs)
+        return self.interpolate_f0(self.resize_f0(f0, p_len))[0]
+
+    def compute_f0_uv(self, wav, p_len=None):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        f0, t = pyworld.harvest(
+            wav.astype(np.double),
+            fs=self.sampling_rate,
+            f0_floor=self.f0_min,
+            f0_ceil=self.f0_max,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
+        return self.interpolate_f0(self.resize_f0(f0, p_len))

lib/infer_pack/modules/F0Predictor/PMF0Predictor.py

@@ -0,0 +1,97 @@
+from lib.infer_pack.modules.F0Predictor.F0Predictor import F0Predictor
+import parselmouth
+import numpy as np
+
+
+class PMF0Predictor(F0Predictor):
+    def __init__(self, hop_length=512, f0_min=50, f0_max=1100, sampling_rate=44100):
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.sampling_rate = sampling_rate
+
+    def interpolate_f0(self, f0):
+        """
+        对F0进行插值处理
+        """
+
+        data = np.reshape(f0, (f0.size, 1))
+
+        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
+        vuv_vector[data > 0.0] = 1.0
+        vuv_vector[data <= 0.0] = 0.0
+
+        ip_data = data
+
+        frame_number = data.size
+        last_value = 0.0
+        for i in range(frame_number):
+            if data[i] <= 0.0:
+                j = i + 1
+                for j in range(i + 1, frame_number):
+                    if data[j] > 0.0:
+                        break
+                if j < frame_number - 1:
+                    if last_value > 0.0:
+                        step = (data[j] - data[i - 1]) / float(j - i)
+                        for k in range(i, j):
+                            ip_data[k] = data[i - 1] + step * (k - i + 1)
+                    else:
+                        for k in range(i, j):
+                            ip_data[k] = data[j]
+                else:
+                    for k in range(i, frame_number):
+                        ip_data[k] = last_value
+            else:
+                ip_data[i] = data[i]  # 这里可能存在一个没有必要的拷贝
+                last_value = data[i]
+
+        return ip_data[:, 0], vuv_vector[:, 0]
+
+    def compute_f0(self, wav, p_len=None):
+        x = wav
+        if p_len is None:
+            p_len = x.shape[0] // self.hop_length
+        else:
+            assert abs(p_len - x.shape[0] // self.hop_length) < 4, "pad length error"
+        time_step = self.hop_length / self.sampling_rate * 1000
+        f0 = (
+            parselmouth.Sound(x, self.sampling_rate)
+            .to_pitch_ac(
+                time_step=time_step / 1000,
+                voicing_threshold=0.6,
+                pitch_floor=self.f0_min,
+                pitch_ceiling=self.f0_max,
+            )
+            .selected_array["frequency"]
+        )
+
+        pad_size = (p_len - len(f0) + 1) // 2
+        if pad_size > 0 or p_len - len(f0) - pad_size > 0:
+            f0 = np.pad(f0, [[pad_size, p_len - len(f0) - pad_size]], mode="constant")
+        f0, uv = self.interpolate_f0(f0)
+        return f0
+
+    def compute_f0_uv(self, wav, p_len=None):
+        x = wav
+        if p_len is None:
+            p_len = x.shape[0] // self.hop_length
+        else:
+            assert abs(p_len - x.shape[0] // self.hop_length) < 4, "pad length error"
+        time_step = self.hop_length / self.sampling_rate * 1000
+        f0 = (
+            parselmouth.Sound(x, self.sampling_rate)
+            .to_pitch_ac(
+                time_step=time_step / 1000,
+                voicing_threshold=0.6,
+                pitch_floor=self.f0_min,
+                pitch_ceiling=self.f0_max,
+            )
+            .selected_array["frequency"]
+        )
+
+        pad_size = (p_len - len(f0) + 1) // 2
+        if pad_size > 0 or p_len - len(f0) - pad_size > 0:
+            f0 = np.pad(f0, [[pad_size, p_len - len(f0) - pad_size]], mode="constant")
+        f0, uv = self.interpolate_f0(f0)
+        return f0, uv

lib/infer_pack/onnx_inference.py

@@ -0,0 +1,145 @@
+import onnxruntime
+import librosa
+import numpy as np
+import soundfile
+
+
+class ContentVec:
+    def __init__(self, vec_path="pretrained/vec-768-layer-12.onnx", device=None):
+        print("load model(s) from {}".format(vec_path))
+        if device == "cpu" or device is None:
+            providers = ["CPUExecutionProvider"]
+        elif device == "cuda":
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        elif device == "dml":
+            providers = ["DmlExecutionProvider"]
+        else:
+            raise RuntimeError("Unsportted Device")
+        self.model = onnxruntime.InferenceSession(vec_path, providers=providers)
+
+    def __call__(self, wav):
+        return self.forward(wav)
+
+    def forward(self, wav):
+        feats = wav
+        if feats.ndim == 2:  # double channels
+            feats = feats.mean(-1)
+        assert feats.ndim == 1, feats.ndim
+        feats = np.expand_dims(np.expand_dims(feats, 0), 0)
+        onnx_input = {self.model.get_inputs()[0].name: feats}
+        logits = self.model.run(None, onnx_input)[0]
+        return logits.transpose(0, 2, 1)
+
+
+def get_f0_predictor(f0_predictor, hop_length, sampling_rate, **kargs):
+    if f0_predictor == "pm":
+        from lib.infer_pack.modules.F0Predictor.PMF0Predictor import PMF0Predictor
+
+        f0_predictor_object = PMF0Predictor(
+            hop_length=hop_length, sampling_rate=sampling_rate
+        )
+    elif f0_predictor == "harvest":
+        from lib.infer_pack.modules.F0Predictor.HarvestF0Predictor import (
+            HarvestF0Predictor,
+        )
+
+        f0_predictor_object = HarvestF0Predictor(
+            hop_length=hop_length, sampling_rate=sampling_rate
+        )
+    elif f0_predictor == "dio":
+        from lib.infer_pack.modules.F0Predictor.DioF0Predictor import DioF0Predictor
+
+        f0_predictor_object = DioF0Predictor(
+            hop_length=hop_length, sampling_rate=sampling_rate
+        )
+    else:
+        raise Exception("Unknown f0 predictor")
+    return f0_predictor_object
+
+
+class OnnxRVC:
+    def __init__(
+        self,
+        model_path,
+        sr=40000,
+        hop_size=512,
+        vec_path="vec-768-layer-12",
+        device="cpu",
+    ):
+        vec_path = f"pretrained/{vec_path}.onnx"
+        self.vec_model = ContentVec(vec_path, device)
+        if device == "cpu" or device is None:
+            providers = ["CPUExecutionProvider"]
+        elif device == "cuda":
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        elif device == "dml":
+            providers = ["DmlExecutionProvider"]
+        else:
+            raise RuntimeError("Unsportted Device")
+        self.model = onnxruntime.InferenceSession(model_path, providers=providers)
+        self.sampling_rate = sr
+        self.hop_size = hop_size
+
+    def forward(self, hubert, hubert_length, pitch, pitchf, ds, rnd):
+        onnx_input = {
+            self.model.get_inputs()[0].name: hubert,
+            self.model.get_inputs()[1].name: hubert_length,
+            self.model.get_inputs()[2].name: pitch,
+            self.model.get_inputs()[3].name: pitchf,
+            self.model.get_inputs()[4].name: ds,
+            self.model.get_inputs()[5].name: rnd,
+        }
+        return (self.model.run(None, onnx_input)[0] * 32767).astype(np.int16)
+
+    def inference(
+        self,
+        raw_path,
+        sid,
+        f0_method="dio",
+        f0_up_key=0,
+        pad_time=0.5,
+        cr_threshold=0.02,
+    ):
+        f0_min = 50
+        f0_max = 1100
+        f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+        f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+        f0_predictor = get_f0_predictor(
+            f0_method,
+            hop_length=self.hop_size,
+            sampling_rate=self.sampling_rate,
+            threshold=cr_threshold,
+        )
+        wav, sr = librosa.load(raw_path, sr=self.sampling_rate)
+        org_length = len(wav)
+        if org_length / sr > 50.0:
+            raise RuntimeError("Reached Max Length")
+
+        wav16k = librosa.resample(wav, orig_sr=self.sampling_rate, target_sr=16000)
+        wav16k = wav16k
+
+        hubert = self.vec_model(wav16k)
+        hubert = np.repeat(hubert, 2, axis=2).transpose(0, 2, 1).astype(np.float32)
+        hubert_length = hubert.shape[1]
+
+        pitchf = f0_predictor.compute_f0(wav, hubert_length)
+        pitchf = pitchf * 2 ** (f0_up_key / 12)
+        pitch = pitchf.copy()
+        f0_mel = 1127 * np.log(1 + pitch / 700)
+        f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * 254 / (
+            f0_mel_max - f0_mel_min
+        ) + 1
+        f0_mel[f0_mel <= 1] = 1
+        f0_mel[f0_mel > 255] = 255
+        pitch = np.rint(f0_mel).astype(np.int64)
+
+        pitchf = pitchf.reshape(1, len(pitchf)).astype(np.float32)
+        pitch = pitch.reshape(1, len(pitch))
+        ds = np.array([sid]).astype(np.int64)
+
+        rnd = np.random.randn(1, 192, hubert_length).astype(np.float32)
+        hubert_length = np.array([hubert_length]).astype(np.int64)
+
+        out_wav = self.forward(hubert, hubert_length, pitch, pitchf, ds, rnd).squeeze()
+        out_wav = np.pad(out_wav, (0, 2 * self.hop_size), "constant")
+        return out_wav[0:org_length]

lib/infer_pack/text/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2017 Keith Ito
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

lib/infer_pack/text/__init__.py

@@ -0,0 +1,59 @@
+""" from https://github.com/keithito/tacotron """
+from . import cleaners
+from .symbols import symbols
+import os
+import platform
+
+if platform.system()=="windows": #local windows package
+  os.environ["PHONEMIZER_ESPEAK_PATH"] = os.path.join(os.getcwd(),"dist","eSpeak NG","espeak-ng.exe")
+  print(os.environ["PHONEMIZER_ESPEAK_PATH"])
+
+# Mappings from symbol to numeric ID and vice versa:
+_symbol_to_id = {s: i for i, s in enumerate(symbols)}
+_id_to_symbol = {i: s for i, s in enumerate(symbols)}
+
+
+def text_to_sequence(text, cleaner_names):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+      cleaner_names: names of the cleaner functions to run the text through
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  sequence = []
+
+  clean_text = _clean_text(text, cleaner_names)
+  for symbol in clean_text:
+    symbol_id = _symbol_to_id[symbol]
+    sequence += [symbol_id]
+  return sequence
+
+
+def cleaned_text_to_sequence(cleaned_text):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  sequence = [_symbol_to_id[symbol] for symbol in cleaned_text]
+  return sequence
+
+
+def sequence_to_text(sequence):
+  '''Converts a sequence of IDs back to a string'''
+  result = ''
+  for symbol_id in sequence:
+    s = _id_to_symbol[symbol_id]
+    result += s
+  return result
+
+
+def _clean_text(text, cleaner_names):
+  for name in cleaner_names:
+    cleaner = getattr(cleaners, name)
+    if not cleaner:
+      raise Exception('Unknown cleaner: %s' % name)
+    text = cleaner(text)
+  return text

lib/infer_pack/text/cleaners.py

@@ -0,0 +1,99 @@
+""" from https://github.com/keithito/tacotron """
+
+'''
+Cleaners are transformations that run over the input text at both training and eval time.
+
+Cleaners can be selected by passing a comma-delimited list of cleaner names as the "cleaners"
+hyperparameter. Some cleaners are English-specific. You'll typically want to use:
+  1. "english_cleaners" for English text
+  2. "transliteration_cleaners" for non-English text that can be transliterated to ASCII using
+     the Unidecode library (https://pypi.python.org/pypi/Unidecode)
+  3. "basic_cleaners" if you do not want to transliterate (in this case, you should also update
+     the symbols in symbols.py to match your data).
+'''
+
+import re
+import emoji
+from unidecode import unidecode
+from phonemizer import phonemize
+from num2words import num2words
+
+# Regular expression matching whitespace:
+_whitespace_re = re.compile(r'\s+')
+
+# List of (regular expression, replacement) pairs for abbreviations:
+_abbreviations = [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [
+  ('mrs', 'misess'),
+  ('mr', 'mister'),
+  ('dr', 'doctor'),
+  ('st', 'saint'),
+  ('co', 'company'),
+  ('jr', 'junior'),
+  ('maj', 'major'),
+  ('gen', 'general'),
+  ('drs', 'doctors'),
+  ('rev', 'reverend'),
+  ('lt', 'lieutenant'),
+  ('hon', 'honorable'),
+  ('sgt', 'sergeant'),
+  ('capt', 'captain'),
+  ('esq', 'esquire'),
+  ('ltd', 'limited'),
+  ('col', 'colonel'),
+  ('ft', 'fort'),
+]]
+
+def expand_abbreviations(text):
+  for regex, replacement in _abbreviations:
+    text = re.sub(regex, replacement, text)
+  return text
+
+def expand_numbers(text):
+  # return normalize_numbers(text)
+  words = [num2words(word) if word.isdigit() else word for word in text.split()]
+  return " ".join(words)
+
+def lowercase(text):
+  return text.lower()
+
+
+def collapse_whitespace(text):
+  return re.sub(_whitespace_re, ' ', text)
+
+def convert_to_ascii(text):
+  return unidecode(text)
+
+def strip_narratives(text):
+  return re.sub(r'([A-Za-z]+:|[*#]+.+?[*#]+)', '', text, flags=re.MULTILINE)
+
+def basic_cleaners(text):
+  '''Basic pipeline that lowercases and collapses whitespace without transliteration.'''
+  text = lowercase(text)
+  text = collapse_whitespace(text)
+  return text
+
+def emoji_cleaner(text: str):
+  return emoji.demojize(text)
+
+def transliteration_cleaners(text):
+  '''Pipeline for non-English text that transliterates to ASCII.'''
+  text = convert_to_ascii(text)
+  text = lowercase(text)
+  text = collapse_whitespace(text)
+  return text
+
+def english_cleaners(text,dialog_only=False):
+    text = transliteration_cleaners(text)
+    text = lowercase(text)
+    text = emoji_cleaner(text)
+    text = expand_abbreviations(text)
+    text = expand_numbers(text)
+    if dialog_only: text = strip_narratives(text)
+    text = collapse_whitespace(text)
+    return text
+
+def english_cleaners2(text):
+    text = english_cleaners(text)
+    text = phonemize(text, language='en-us', backend='espeak', strip=True, preserve_punctuation=True, with_stress=True)
+    text = collapse_whitespace(text)
+    return text

lib/infer_pack/text/symbols.py

@@ -0,0 +1,16 @@
+""" from https://github.com/keithito/tacotron """
+
+'''
+Defines the set of symbols used in text input to the model.
+'''
+_pad        = '_'
+_punctuation = ';:,.!?¡¿—…"«»“” '
+_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
+_letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+
+
+# Export all symbols:
+symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+
+# Special symbol ids
+SPACE_ID = symbols.index(" ")

lib/infer_pack/transforms.py

@@ -0,0 +1,209 @@
+import torch
+from torch.nn import functional as F
+
+import numpy as np
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails=None,
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if tails is None:
+        spline_fn = rational_quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unconstrained_rational_quadratic_spline
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+    outputs, logabsdet = spline_fn(
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unconstrained_rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails="linear",
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == "linear":
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError("{} tails are not implemented.".format(tails))
+
+    (
+        outputs[inside_interval_mask],
+        logabsdet[inside_interval_mask],
+    ) = rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+
+    return outputs, logabsdet
+
+
+def rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    left=0.0,
+    right=1.0,
+    bottom=0.0,
+    top=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError("Minimal bin width too large for the number of bins")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError("Minimal bin height too large for the number of bins")
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (
+            input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+        )
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet

lib/karafan/Data/Models.csv

@@ -0,0 +1,16 @@
+Name,Use,Stem,Vol_Comp,Cut_OFF,N_FFT_scale,dim_F_set,dim_T_set,Date,Repo_FileName
+MDX23C 8K FFT,x,BOTH,1.0,0,8192,4096,256,2023-09-26,MDX23C-8KFFT-InstVoc_HQ.ckpt
+MDX23C 8K FFT - v2,x,BOTH,1.0,0,8192,4096,256,2023-09-26,MDX23C-8KFFT-InstVoc_HQ_2.ckpt
+Karaoke,,Choirs,1.035,14700,6144,2048,256,2022-07-21,UVR_MDXNET_KARA.onnx
+Karaoke 2,x,Choirs,1.065,17640,5120,2048,256,2023-03-20,UVR_MDXNET_KARA_2.onnx
+Instrum 3,x,Instrumental,1.028,17640,7680,3072,256,2022-10-31,UVR-MDX-NET-Inst_3.onnx
+Instrum HQ 1,x,Instrumental,1.035,0,6144,3072,256,2023-03-28,UVR-MDX-NET-Inst_HQ_1.onnx
+Instrum HQ 2,x,Instrumental,1.033,0,6144,3072,256,2023-04-10,UVR-MDX-NET-Inst_HQ_2.onnx
+Instrum HQ 3,x,Instrumental,1.022,0,6144,3072,256,2023-06-11,UVR-MDX-NET-Inst_HQ_3.onnx
+Instrum Main,x,Instrumental,1.025,17640,5120,2048,256,2022-12-19,UVR-MDX-NET-Inst_Main.onnx
+Kim Instrum,x,Instrumental,1.030,17640,7680,3072,256,2023-04-26,Kim_Inst.onnx
+(de-)Reverb HQ,,OTHER,1.035,0,6144,3072,512,2023-04-13,Reverb_HQ_By_FoxJoy.onnx
+Kim Vocal 1,x,Vocals,1.043,17640,7680,3072,256,2022-12-27,Kim_Vocal_1.onnx
+Kim Vocal 2,x,Vocals,1.009,17640,7680,3072,256,2023-05-19,Kim_Vocal_2.onnx
+Voc FT,x,Vocals,1.021,17640,7680,3072,256,2023-06-22,UVR-MDX-NET-Voc_FT.onnx
+Vocal Main,x,Vocals,1.043,17640,7680,3072,256,2022-07-21,UVR_MDXNET_Main.onnx

lib/karafan/Data/model_2_stem_full_band_8k.yaml

@@ -0,0 +1,43 @@
+audio:
+  chunk_size: 261120
+  dim_f: 4096
+  dim_t: 256
+  hop_length: 1024
+  n_fft: 8192
+  num_channels: 2
+  sample_rate: 44100
+  min_mean_abs: 0.001
+model:
+  act: gelu
+  bottleneck_factor: 4
+  growth: 128
+  norm: InstanceNorm
+  num_blocks_per_scale: 2
+  num_channels: 128
+  num_scales: 5
+  num_subbands: 4
+  scale:
+  - 2
+  - 2
+training:
+  batch_size: 6
+  grad_clip: 0
+  instruments:
+  - Vocals
+  - Instrumental
+  lr: 1.0e-05
+  patience: 2
+  reduce_factor: 0.95
+  target_instrument: null
+  num_epochs: 1000
+  num_steps: 1000
+  augmentation: 1
+  augmentation_type: simple1
+  augmentation_mix: true
+  q: 0.95
+  coarse_loss_clip: true
+  ema_momentum: 0.999
+inference:
+  batch_size: 1
+  dim_t: 256
+  num_overlap: 8

lib/karafan/__init__.py

@@ -0,0 +1,2 @@
+# Version 5.4
+from . import inference

lib/karafan/audio_utils.py

@@ -0,0 +1,398 @@
+
+#   MIT License - Copyright (c) 2023 - Captain FLAM & Jarredou
+#
+#   https://github.com/Captain-FLAM/KaraFan
+
+import os, librosa, subprocess, tempfile, soundfile as sf, numpy as np
+from scipy import signal
+
+def Load_Audio(file, sample_rate, ffmpeg = None, output_path = None):
+
+	
+	# Load the audio file from User's input
+	try:
+		print("Load the audio file from User's input")
+		audio, sr = sf.read(file, dtype='float32')
+
+	# Corrupted file : try to correct it with ffmpeg
+	except RuntimeError as e:
+		print('<div style="font-size:18px;color:#ff0040;"><b>Your audio file is Bad encoded ! Trying to correct ...</b></div>')
+		try:
+			corrected = os.path.join(output_path, "X - CORRECTED.flac")
+
+			if not os.path.exists(corrected):
+				subprocess.run(f'"{ffmpeg}" -y -i "{file}" -codec:a flac -compression_level 5 -ch_mode mid_side -lpc_type cholesky -lpc_passes 1 -exact_rice_parameters 1 "{corrected}"', shell=True, text=True, capture_output=True, check=True)
+
+			print("Load corrected audio file")
+			audio, sr = sf.read(corrected, dtype='float32')
+			
+		except RuntimeError as e:
+			print(f"Error : {e}");  return None, 0
+		except Exception as e:
+			print(f"Error : {e}");  return None, 0
+	except Exception as e:
+		print(f"Error : {e}");  return None, 0
+	
+	audio = audio.T
+
+	# Convert mono to stereo (if needed)
+	if audio.ndim == 1:
+		audio = np.asfortranarray([audio, audio])
+
+	# Convert to 44100 Hz if needed (for MDX models)
+	if sr != sample_rate:
+		audio = librosa.resample(audio, orig_sr = sr, target_sr = sample_rate, res_type = 'kaiser_best', axis=-1, fix=True)
+
+	return audio, sample_rate
+
+def Save_Audio(file_path, audio, sample_rate, output_format, cut_off, ffmpeg):
+
+	# if output_format == 'PCM_16' or output_format == 'FLOAT':
+	if output_format == 'WAV':
+		sf.write(file_path + '.wav',  audio.T, sample_rate, format='wav', subtype = "PCM_16" if "int" in audio.dtype.name else "FLOAT")
+	else:
+		# Create a temporary file
+		temp = tempfile.NamedTemporaryFile(suffix='.wav', delete=False)
+
+		if output_format == 'FLAC':
+			sf.write(temp, audio.T, sample_rate, format='wav', subtype='FLOAT')
+			ffmpeg = f'"{ffmpeg}" -y -i "{temp.name}" -codec:a flac -compression_level 5 -ch_mode mid_side -frame_size {sample_rate} -lpc_type cholesky -lpc_passes 1 -exact_rice_parameters 1 "{file_path}.flac"'
+
+		elif output_format == 'MP3':
+			# TODO : Correct the BUG of Lame encoder which modify the length of audio results (~ +30 ms on short song, -30 ms on long song) ?!?!
+
+			# about VBR/CBR/ABR		: https://trac.ffmpeg.org/wiki/Encode/MP3
+			# about ffmpeg wrapper	: http://ffmpeg.org/ffmpeg-codecs.html#libmp3lame-1
+			# recommended settings	: https://wiki.hydrogenaud.io/index.php?title=LAME#Recommended_encoder_settings
+
+			# 320k is mandatory, else there is a weird cutoff @ 16khz with VBR parameters = ['-q','0'] !!
+			# (equivalent to lame "-V0" - 220-260 kbps , 245 kbps average)
+			# And also, parameters = ['-joint_stereo', '0'] (Separated stereo channels)
+			# is WORSE than "Joint Stereo" for High Frequencies !
+			# So let's use it by default for MP3 encoding !!
+
+			sf.write(temp, audio.T, sample_rate, format='wav', subtype='PCM_16')
+			ffmpeg = f'"{ffmpeg}" -y -i "{temp.name}" -codec:a libmp3lame -b:a 320k -q:a 0 -joint_stereo 1 -cutoff {cut_off} "{file_path}.mp3"'
+
+		try:
+			subprocess.run(ffmpeg, shell=True, text=True, capture_output=True, check=True)
+
+		except subprocess.CalledProcessError as e:
+			if e.returncode == 127:
+				print('WARNING : "\nFFmpeg" is not installed on your system !!\n')
+			else:
+				print("Error :\n" + e.stderr + "\n" + e.stdout)
+
+		temp.close()
+		os.remove(temp.name)
+
+def Normalize(audio, threshold_dB = -1.0):
+	"""
+	Normalize audio to -1.0 dB peak amplitude
+	This is mandatory for SOME audio files because every process is based on RMS dB levels.
+	(Models processing, Volumes Compensations & audio Substractions)
+	"""
+	audio = audio.T
+	
+	# Suppress DC shift (center on 0.0 vertically)
+	audio -= np.mean(audio)
+
+	# Normalize audio peak amplitude to -1.0 dB
+	max_peak = np.max(np.abs(audio))
+	if max_peak > 0.0:
+		max_db = 10 ** (threshold_dB / 20)  # Convert -X dB to linear scale
+		audio /= max_peak
+		audio *= max_db
+
+	return audio.T
+
+def Silent(audio_in, sample_rate, threshold_dB = -50):
+	"""
+	Make silent the parts of audio where dynamic range (RMS) goes below threshold.
+	Don't misundertand : this function is NOT a noise reduction !
+	Its behavior is to clean the audio from "silent parts" (below -XX dB) to :
+	- avoid the MLM model to work on "silent parts", and save GPU time
+	- avoid the MLM model to produce artifacts on "silent parts"
+	- clean the final audio files from residues of "silent parts"
+	"""
+
+	min_size		= int(1.000 * sample_rate)  # 1000 ms
+	window_frame	= int(0.500 * sample_rate)  #  500 ms
+	fade_duration	= int(0.300 * sample_rate)  #  450 ms
+	fade_out		= np.linspace(1.0, 0.0, fade_duration)
+	fade_in			= np.linspace(0.0, 1.0, fade_duration)
+
+	start = 0; end = 0
+	audio = audio_in.copy()
+	audio_length = audio_in.shape[1]
+
+	for i in range(0, audio_length, window_frame):
+		
+		# TODO : Maybe use S=audio (Spectrogram) instead of y=audio ??
+		RMS = np.max(librosa.amplitude_to_db(librosa.feature.rms(y=audio[:, i:(i + window_frame)], frame_length=window_frame, hop_length=window_frame)))
+		if RMS < threshold_dB:
+			end = i + window_frame
+			# Last part (in case of silence at the end)
+			if i >= audio_length - window_frame:
+				if end - start > min_size:
+					# Fade out
+					if start > fade_duration:
+						audio[:, start:(start + fade_duration)] *= fade_out
+						start += fade_duration
+
+					# Clean last part
+					audio[:, start:audio_length] = 0.0
+					break
+		else:
+			# Clean the "min_size" samples found
+			if end - start > min_size:
+
+				# print(f"RMS : {RMS} / start : {start} / end : {end}")
+
+				# Fade out
+				if start > fade_duration:
+					audio[:, start:(start + fade_duration)] *= fade_out
+					start += fade_duration
+
+				# Fade in
+				if end < audio_length - fade_duration:
+					audio[:, (end - fade_duration):end] *= fade_in
+					end -= fade_duration
+		
+				# Clean in between
+				audio[:, start:end] = 0.0
+
+			start = i
+			
+	return audio
+
+# Linkwitz-Riley filter
+#
+# Avec cutoff = 17.4khz & -80dB d'atténuation:
+#
+# ordre =  4 => filtre target freq = 10500hz
+# ordre =  6 => filtre target freq = 13200hz
+# ordre =  8 => filtre target freq = 14300hz
+# ordre = 10 => filtre target freq = 15000hz
+# ordre = 12 => filtre target freq = 15500hz
+# ordre = 14 => filtre target freq = 15800hz
+# ordre = 16 => filtre target freq = 16100hz
+#
+# Avec cutoff = 17.4khz & -60dB d'atténuation:
+#
+# ordre =  4 => filtre target freq = 12500hz (-4900)
+# ordre =  6 => filtre target freq = 14400hz
+# ordre =  8 => filtre target freq = 15200hz (-2200)
+# ordre = 10 => filtre target freq = 15700hz
+# ordre = 12 => filtre target freq = 16000hz (-1640)
+# ordre = 14 => filtre target freq = 16200hz
+# ordre = 16 => filtre target freq = 16400hz
+
+def Linkwitz_Riley_filter(type, cutoff, audio, sample_rate, order=8):
+	
+	# cutoff -= 2200
+
+	nyquist = 0.5 * sample_rate
+	normal_cutoff = cutoff / nyquist
+
+	sos = signal.butter(order // 2, normal_cutoff, btype=type, analog=False, output='sos')
+	filtered_audio = signal.sosfiltfilt(sos, audio, padlen=0, axis=1)
+
+	return filtered_audio
+
+# Band Pass filter
+#
+# Vocals -> lowest : 85 - 100 Hz, highest : 20 KHz
+# Music  -> lowest : 30 -  50 Hz, highest : 18-20 KHz
+#
+# Voix masculine :
+#
+# Minimale : 85 Hz
+# Fondamentale : 180 Hz
+# Maximale (y compris les harmoniques) : 14 kHz
+#
+# Voix féminine :
+#
+# Minimale : 165 Hz
+# Fondamentale : 255 Hz
+# Maximale (y compris les harmoniques) : 16 kHz
+#
+# Voix d'enfants :
+#
+# Minimale : 250 Hz
+# Fondamentale : 400 Hz
+# Maximale (y compris les harmoniques) : 20 kHz ou +
+
+def Pass_filter(type, cutoff, audio, sample_rate, order=32):
+
+	if cutoff >= sample_rate / 2:
+		cutoff = (sample_rate / 2) - 1
+
+	sos = signal.butter(order // 2, cutoff, btype=type, fs=sample_rate, output='sos')
+	filtered_audio = signal.sosfiltfilt(sos, audio, padlen=0, axis=1)
+
+	return filtered_audio
+
+# SRS : Sample Rate Scaling
+def Change_sample_rate(audio, way, current_cutoff, target_cutoff):
+
+	if way == "DOWN":
+		current_cutoff, target_cutoff = target_cutoff, current_cutoff
+
+	pitched_audio = librosa.resample(audio, orig_sr = current_cutoff * 2, target_sr = target_cutoff * 2, res_type = 'kaiser_best', axis=1)
+
+	# print(f"SRS input audio shape: {audio.shape}")
+	# print(f"SRS output audio shape: {pitched_audio.shape}")
+	# print (f"ratio : {ratio}")
+
+	return pitched_audio
+
+# def Remove_High_freq_Noise(audio, threshold_freq):
+
+# 	# Calculer la transformée de Fourier
+# 	stft = librosa.stft(audio)
+	
+# 	# Calculer la somme des amplitudes pour chaque fréquence dans le spectre
+# 	amplitude_sum = np.sum(np.abs(stft), axis=0)
+
+# 	# Appliquer un masque pour supprimer les fréquences supérieures lorsque la somme des amplitudes est inférieure au seuil
+# 	stft[:, amplitude_sum > threshold_freq] = 0.0
+
+# 	# Reconstruire l'audio à partir du STFT modifié
+	# 	audio_filtered = librosa.istft(stft)
+
+# 	return audio_filtered
+
+# def Match_Freq_CutOFF(self, audio1, audio2, sample_rate):
+# 	# This option matches the Primary stem frequency cut-off to the Secondary stem frequency cut-off
+# 	# (if the Primary stem frequency cut-off is lower than the Secondary stem frequency cut-off)
+
+# 	# Get the Primary stem frequency cut-off
+# 	freq_cut_off1 = Find_Cut_OFF(audio1, sample_rate)
+# 	freq_cut_off2 = Find_Cut_OFF(audio2, sample_rate)
+
+# 	# Match the Primary stem frequency cut-off to the Secondary stem frequency cut-off
+# 	if freq_cut_off1 < freq_cut_off2:
+# 		audio1 = Resize_Freq_CutOFF(audio1, freq_cut_off2, sample_rate)
+
+# 	return audio1
+
+# # Find the high cut-off frequency of the input audio
+# def Find_Cut_OFF(audio, sample_rate, threshold=0.01):
+
+# 	# Appliquer un filtre passe-bas pour réduire le bruit
+# 	cutoff_frequency = sample_rate / 2.0  # Fréquence de Nyquist (la moitié du taux d'échantillonnage)
+
+# 	# Définir l'ordre du filtre passe-bas
+# 	order = 6
+
+# 	# Calculer les coefficients du filtre passe-bas
+# 	b, a = signal.butter(order, cutoff_frequency - threshold, btype='low', analog=False, fs=sample_rate)
+
+# 	# Appliquer le filtre au signal audio
+# 	filtered_audio = signal.lfilter(b, a, audio, axis=0)
+
+# 	# Calculer la FFT du signal audio filtré
+# 	fft_result = np.fft.fft(filtered_audio, axis=0)
+
+# 	# Calculer les magnitudes du spectre de fréquence
+# 	magnitudes = np.abs(fft_result)
+
+# 	# Calculer les fréquences correspondant aux bins de la FFT
+# 	frequencies = np.fft.fftfreq(len(audio), 1.0 / sample_rate)
+
+# 	# Trouver la fréquence de coupure où la magnitude tombe en dessous du seuil
+# 	cut_off_frequencies = frequencies[np.where(magnitudes > threshold)]
+
+# 	# Trouver la fréquence de coupure maximale parmi toutes les valeurs
+# 	return int(max(cut_off_frequencies))
+
+
+
+# - For the code below :
+#   MIT License
+#
+#   Copyright (c) 2023 Anjok07 & aufr33 - Ultimate Vocal Remover (UVR 5)
+#
+# - https://github.com/Anjok07/ultimatevocalremovergui
+
+MAX_SPEC = 'Max'
+MIN_SPEC = 'Min'
+AVERAGE  = 'Average'
+
+def Make_Ensemble(algorithm, audio_input):
+
+	if len(audio_input) == 1:  return audio_input[0]
+	
+	waves = []
+	
+	if algorithm == AVERAGE:
+
+		waves_shapes = []
+		final_waves = []
+
+		for i in range(len(audio_input)):
+			wave = audio_input[i]
+			waves.append(wave)
+			waves_shapes.append(wave.shape[1])
+
+		wave_shapes_index = waves_shapes.index(max(waves_shapes))
+		target_shape = waves[wave_shapes_index]
+		waves.pop(wave_shapes_index)
+		final_waves.append(target_shape)
+
+		for n_array in waves:
+			wav_target = to_shape(n_array, target_shape.shape)
+			final_waves.append(wav_target)
+
+		waves = sum(final_waves)
+		output = waves / len(audio_input)
+	else:
+		specs = []
+		
+		for i in range(len(audio_input)):  
+			waves.append(audio_input[i])
+			
+			# wave_to_spectrogram_no_mp
+			spec = librosa.stft(audio_input[i], n_fft=6144, hop_length=1024)
+			
+			if spec.ndim == 1:  spec = np.asfortranarray([spec, spec])
+
+			specs.append(spec)
+		
+		waves_shapes = [w.shape[1] for w in waves]
+		target_shape = waves[waves_shapes.index(max(waves_shapes))]
+		
+		# spectrogram_to_wave_no_mp
+		wave = librosa.istft(ensembling(algorithm, specs), n_fft=6144, hop_length=1024)
+	
+		if wave.ndim == 1:  wave = np.asfortranarray([wave, wave])
+
+		output = to_shape(wave, target_shape.shape)
+
+	return output
+
+def ensembling(a, specs):
+	for i in range(1, len(specs)):
+		if i == 1:
+			spec = specs[0]
+
+		ln = min([spec.shape[2], specs[i].shape[2]])
+		spec = spec[:,:,:ln]
+		specs[i] = specs[i][:,:,:ln]
+		
+		if MIN_SPEC == a:
+			spec = np.where(np.abs(specs[i]) <= np.abs(spec), specs[i], spec)
+		elif MAX_SPEC == a:
+			spec = np.where(np.abs(specs[i]) >= np.abs(spec), specs[i], spec)
+
+	return spec
+
+def to_shape(x, target_shape):
+	padding_list = []
+	for x_dim, target_dim in zip(x.shape, target_shape):
+		pad_value = (target_dim - x_dim)
+		pad_tuple = ((0, pad_value))
+		padding_list.append(pad_tuple)
+	
+	return np.pad(x, tuple(padding_list), mode='constant')

lib/karafan/compare.py

@@ -0,0 +1,125 @@
+
+#   MIT License
+#
+#   Copyright (c) 2023 - Captain FLAM
+#   Copyright (c) 2021 - Roman Solovyev (ZFTurbo)
+#
+#   https://github.com/Captain-FLAM/KaraFan
+#   https://github.com/ZFTurbo/Audio-separation-models-checker
+
+import os, glob, datetime, librosa, soundfile as sf, numpy as np
+
+"""
+SDR - Source to distortion ratio   
+SIR - Source to inferences ratio
+SNR - Source to noise ration
+SAR - Source to artifacts ratio
+"""
+
+# compute SDR for one song
+def calculate(reference, estimate):
+	references = np.expand_dims(reference, axis=0)
+	estimates  = np.expand_dims(estimate, axis=0)
+
+	if estimates.shape != references.shape:
+		print('Warning: Different audio lengths : {} != {} !'.format(estimates.shape, references.shape))
+		return [None]
+
+	delta = 1e-7  # avoid numerical errors
+	num = np.sum(np.square(references), axis=(1, 2))
+	den = np.sum(np.square(references - estimates), axis=(1, 2))
+	num += delta
+	den += delta
+	return 10 * np.log10(num / den)
+
+def SDR(song_output_path, output_format, Gdrive, elapsed_time):
+
+	# The "song_output_path" contains the NAME of the song to compare within the "Gdrive > KaraFan_user > Multi-Song" folder
+
+	song_name	= os.path.basename(song_output_path)
+	MultiSong	= os.path.join(Gdrive, "KaraFan_user", "Multi_Song", "Stems")
+	instrum		= os.path.join(MultiSong, song_name[-3:] + '_instrum.flac')  # Get only the number of the song
+	vocals		= os.path.join(MultiSong, song_name[-3:] + '_vocals.flac')
+
+	if not os.path.exists(MultiSong):  print ("No Multi-Song folder found !");  return
+	if not os.path.exists(instrum) or not os.path.exists(vocals):  print (f'No References Stems for "{song_name}" found !');  return
+
+	if output_format=='PCM_16':	ext = 'wav'
+	elif output_format=='FLOAT': ext = 'wav'
+	elif output_format == 'WAV': ext = 'wav'
+	elif output_format== "FLAC": ext = 'flac'
+	elif output_format== 'MP3':	ext = 'mp3'
+
+	Extracted_files = glob.glob(os.path.join(song_output_path, '*.' + ext))
+
+	if len(Extracted_files) == 0:  print('Check output folder. Cant find any files !');  return
+
+	 # The stems we want to process
+	Scores			= {"instrum": [], "vocals": []}
+	References		= {"instrum": None, "vocals": None}
+
+	Results = ""
+	References["instrum"], _ = sf.read(os.path.join(MultiSong, song_name[-3:] + '_instrum.flac'), dtype='float32')  # Get only the number of the song
+	References["vocals"], _  = sf.read(os.path.join(MultiSong, song_name[-3:] + '_vocals.flac'), dtype='float32')
+
+	Extracted_files = sorted(Extracted_files)
+
+	for extract in Extracted_files:
+
+		file_name = os.path.basename(extract)
+		file_name = file_name.replace(os.path.splitext(file_name)[1], "")  # Remove extension
+
+		if "Bleedings" in file_name:  continue  # Skip Bleedings
+
+		if "Vocal" in file_name:	type = "vocals"
+		elif "Music" in file_name:	type = "instrum"
+		else:						type = "others"
+
+		if type == "others":  continue  # Skip Others
+
+		estimate, _	= sf.read(extract, dtype='float32')
+		song_score	= calculate(References[type], estimate)[0]
+
+		pad = 40 - len(file_name)
+		
+		if song_score is None:
+			print("• " + file_name + (" " * pad) + 'SDR : --- [ Skipped ]')
+		else:
+			print("• " + file_name + ("&nbsp;" * pad) + 'SDR : <b>{:9.6f}</b>'.format(song_score))
+			Results += file_name + (" " * pad) + 'SDR : {:9.6f}'.format(song_score) + "\n"
+
+			Scores[type].append(song_score)
+
+	print("\n")
+	
+	# TODO : Use for batch SDR with multiple songs
+	# for type in Scores:
+	# 	if len(Scores[type]) > 0:
+	# 		print('Average SDR {} : <b>{:9.6f}</b>'.format(type, np.array(Scores[type]).mean()))
+	# 	else:
+	# 		print('Average SDR {} : ---'.format(type))
+
+	# Write results on disk
+	if Results != "":
+		with open(os.path.join(song_output_path, "SDR_Results.txt"), 'a', encoding='utf-8') as file:
+			file.write(f"\n► {datetime.datetime.now().strftime('%Y-%m-%d ~ %H:%M:%S')} - {song_name}\n\n")
+			file.write(Results)
+			file.write("\n" + elapsed_time + "\n")
+
+
+# def Spectrograms(audio_file1, audio_file2):
+	
+# 	audio1, _ = sf.read(audio_file1, sr=None, mono=False, dtype='float32')
+# 	audio2, _ = sf.read(audio_file2, sr=None, mono=False, dtype='float32')
+
+# 	# Convertir les signaux audio en spectrogrammes
+# 	spec1 = librosa.stft(audio1.T, n_fft=4096, hop_length=1024)
+# 	spec2 = librosa.stft(audio2.T, n_fft=4096, hop_length=1024)
+
+# 	# Calculer la distance euclidienne entre les spectrogrammes
+# 	distance = np.linalg.norm(spec2 - spec1)
+
+# 	# Normaliser la distance pour obtenir une mesure de similarité (plus proche de zéro est meilleur)
+# 	similarity = 1 / (1 + distance)
+
+# 	return f"{similarity * 100:.6f} %"

lib/karafan/inference.py

@@ -0,0 +1,880 @@
+
+#   MIT License
+#
+#   Copyright (c) 2023 - Captain FLAM - Heavily modified ! (GUI, sequential processing, etc ...)
+#   Copyright (c) 2023 - Jarredou     - Did all the job for Inference & Assists me in audio tricks !!
+#   Copyright (c) 2023 - ZFTurbo      - Started the project MVSEP-MDX23 (music separation model)
+#
+#   https://github.com/Captain-FLAM/KaraFan
+#   https://github.com/jarredou/MVSEP-MDX23-Colab_v2/
+#   https://github.com/ZFTurbo/MVSEP-MDX23-music-separation-model
+
+
+import os, gc, csv, time, platform, regex, requests, torch, librosa
+import numpy as np, onnxruntime as ort
+
+# for MDX23C models
+import yaml
+from .. import BASE_CACHE_DIR, BASE_DIR, BASE_MODELS_DIR, ObjectNamespace as ConfigDict
+from ..audio import remix_audio
+from ..utils import gc_collect
+
+# Maybe not available on PC
+try:
+	from IPython.display import display, HTML
+except:
+	pass
+
+from . import settings, audio_utils, compare, tfc_tdf
+
+CWD = os.path.dirname(os.path.abspath(__file__))
+
+def demix_mdxv3(mix, model, device, config, overlap_MDX23):
+	
+	mix = torch.tensor(mix, dtype=torch.float32)
+	try:
+		S = model.num_target_instruments
+	except Exception as e:
+		S = model.module.num_target_instruments
+	
+	mdx_window_size = config.inference.dim_t
+	
+	# batch_size = config.inference.batch_size
+	batch_size = 1
+	C = config.audio.hop_length * (mdx_window_size - 1)
+	H = C // overlap_MDX23
+	L = mix.shape[1]
+	pad_size = H - (L - C) % H
+	mix = torch.cat([torch.zeros(2, C - H, dtype=torch.float32), mix, torch.zeros(2, pad_size + C - H, dtype=torch.float32)], 1)
+	mix = mix.to(device)
+
+	chunks = mix.unfold(1, C, H).transpose(0, 1)
+
+	batches = [chunks[i : i + batch_size] for i in range(0, len(chunks), batch_size)]
+
+	X = torch.zeros(S, *mix.shape, dtype=torch.float32).to(device) if S > 1 else torch.zeros_like(mix, dtype=torch.float32)
+
+	with torch.cuda.amp.autocast(dtype=torch.float32):  # BUG fix : float16 by default !!
+		with torch.no_grad():
+			cnt = 0
+			for batch in batches:
+
+				x = model(batch)
+				x[torch.isnan(x)] = 0.0  # Replace "NaN" by zeros, just by security ...
+
+				for w in x:
+					X[..., cnt * H : (cnt * H) + C] += w
+					cnt += 1
+
+	estimated_sources = X[..., C - H:-(pad_size + C - H)] / overlap_MDX23
+	
+	if S > 1:
+		return {k: v for k, v in zip(config.training.instruments, estimated_sources.cpu().numpy())}
+	else:
+		return estimated_sources.cpu().numpy()
+
+def demix_base(mix, device, model, infer_session):
+	sources = []
+	n_sample = mix.shape[1]
+	
+	trim = model.n_fft // 2
+	gen_size = model.chunk_size - 2 * trim
+	pad = gen_size - n_sample % gen_size
+	mix_p = np.concatenate(
+		(
+			np.zeros((2, trim)),
+			mix,
+			np.zeros((2, pad)),
+			np.zeros((2, trim))
+		), 1
+	)
+
+	mix_waves = []
+	i = 0
+	while i < n_sample + pad:
+		waves = np.array(mix_p[:, i:i + model.chunk_size])
+		mix_waves.append(waves)
+		i += gen_size
+	mix_waves = np.array(mix_waves)
+	mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(device)
+
+	try:
+		with torch.no_grad():
+			_ort = infer_session
+			stft_res = model.stft(mix_waves)
+			res = _ort.run(None, {'input': stft_res.cpu().numpy()})[0]
+			ten = torch.tensor(res, dtype=torch.float32)
+			tar_waves = model.istft(ten.to(device))
+			tar_waves = tar_waves.cpu()
+			tar_signal = tar_waves[:, :, trim:-trim].transpose(0, 1).reshape(2, -1).numpy()[:, :-pad]
+
+		sources.append(tar_signal)
+
+	except Exception as e:
+		print("\n\nError in demix_base() with Torch : ", e)
+
+	return np.array(sources)
+
+
+class MusicSeparationModel:
+
+	def __init__(self, config, cache_dir=None, use_cache=True, format: str=None):
+
+		self.Gdrive		= BASE_DIR
+
+		self.output_format		= config['AUDIO']['output_format'] if format is None else format.upper()
+		# Integers
+		self.normalize			= int(config['AUDIO']['normalize'])
+		self.silent				= int(config['AUDIO']['silent'])
+		self.high_pass			= int(config['PROCESS']['high_pass']) * 5				# from 0 to 100 Hz
+		self.low_pass			= 14000 + (int(config['PROCESS']['low_pass']) * 500)	# from 14 Khz to 22 Khz
+		self.chunk_size			= int(config['OPTIONS']['chunk_size'])
+		# Booleans
+		self.infra_bass			= config['AUDIO']['infra_bass']
+		self.DEBUG				= config['BONUS']['DEBUG']
+		self.GOD_MODE			= use_cache #config['BONUS']['GOD_MODE']
+		self.large_gpu			= config['BONUS']['large_gpu']
+
+		self.output = os.path.join(BASE_CACHE_DIR,"temp","karafan") if cache_dir is None else cache_dir
+		
+		self.ffmpeg = os.path.join(self.Gdrive, "ffmpeg") + (".exe" if platform.system() == 'Windows' else "")
+		
+		if torch.cuda.is_available():
+			self.GPU_device = 'cuda:0'
+			self.providers = ["CUDAExecutionProvider"]
+		# For MAC OSX
+		elif torch.backends.mps.is_available():
+			self.GPU_device = 'mps'
+			self.providers = None
+		else:
+			self.GPU_device = 'cpu'
+			self.providers = ["CPUExecutionProvider"]
+			
+		# MDX23C 8K
+		self.MDX23_overlap = 1
+		with open(os.path.join(CWD, "Data", "model_2_stem_full_band_8k.yaml")) as file:
+			data = yaml.load(file, Loader=yaml.FullLoader)
+			self.MDX23_config = ConfigDict(**data)
+
+		# Set BigShifts from Speed option
+		if config['OPTIONS']['speed']=='Fastest':
+			self.Quality_Vocal = { 'BigShifts': 1, 'BigShifts_SRS': 0 } # 1 + 0 + 0 = 1 pass
+			self.Quality_Music = { 'BigShifts': 1, 'BigShifts_SRS': 0 } # + 1  = 2 pass
+			self.Quality_Bleed = { 'BigShifts': 1, 'BigShifts_SRS': 0 }
+			self.MDX23_overlap = 1
+			self.MDX23_bleed = 1
+		elif config['OPTIONS']['speed']== 'Fast':
+			self.Quality_Vocal = { 'BigShifts': 1, 'BigShifts_SRS': 1 } # 1 + 1 + 1 = 3 pass
+			self.Quality_Music = { 'BigShifts': 1, 'BigShifts_SRS': 0 } # + 1  = 4 pass
+			self.Quality_Bleed = { 'BigShifts': 1, 'BigShifts_SRS': 1 }
+			self.MDX23_overlap = 2
+			self.MDX23_bleed = 1
+		elif config['OPTIONS']['speed']== 'Medium':
+			self.Quality_Vocal = { 'BigShifts': 1, 'BigShifts_SRS': 3 } # 1 + 3 + 1 = 5 pass
+			self.Quality_Music = { 'BigShifts': 2, 'BigShifts_SRS': 0 } # + 2  = 7 pass
+			self.Quality_Bleed = { 'BigShifts': 2, 'BigShifts_SRS': 0 }
+			self.MDX23_overlap = 4
+			self.MDX23_bleed = 2
+		elif config['OPTIONS']['speed']== 'Slow':
+			self.Quality_Vocal = { 'BigShifts': 2, 'BigShifts_SRS': 3 } # 2 + 3 + 1 = 6 pass
+			self.Quality_Music = { 'BigShifts': 3, 'BigShifts_SRS': 0 } # + 3  = 9 pass
+			self.Quality_Bleed = { 'BigShifts': 2, 'BigShifts_SRS': 1 }
+			self.MDX23_overlap = 6
+			self.MDX23_bleed = 3
+		elif config['OPTIONS']['speed']== 'Slowest':
+			self.Quality_Vocal = { 'BigShifts': 2, 'BigShifts_SRS': 4 } # 2 + 4 + 1 = 7 pass
+			self.Quality_Music = { 'BigShifts': 4, 'BigShifts_SRS': 0 } # + 4  = 11 pass
+			self.Quality_Bleed = { 'BigShifts': 2, 'BigShifts_SRS': 2 }
+			self.MDX23_overlap = 8
+			self.MDX23_bleed = 4
+		
+		# MDX-B models initialization
+
+		self.models = { 'music': [], 'vocal': [], 'bleed_music': [] , 'bleed_vocal': [] , 'remove_music': [] }
+		self.MDX = {}
+
+		# Load Models parameters
+		with open(os.path.join(CWD, "Data", "Models.csv")) as csvfile:
+			reader = csv.DictReader(csvfile, quoting=csv.QUOTE_ALL)
+			for row in reader:
+				# ignore "Other" stems for now !
+				name = row['Name']
+				if name == "" or name == "Name":  continue
+
+				# IMPORTANT : Volume Compensations are specific for each model !!!
+
+				if name == config['PROCESS']['music_1'] or name == config['PROCESS']['music_2']:
+					row['Compensation'] = float(row['Vol_Comp'])
+					self.models['music'].append(row)
+				
+				if name == config['PROCESS']['vocal_1'] or name == config['PROCESS']['vocal_2']:
+					row['Compensation'] = float(row['Vol_Comp'])
+					self.models['vocal'].append(row)
+
+				# Special case for "Bleedings Filter"
+
+				# --> it's a Music model, so look at "self.Quality_Music" for references !!
+				if name == config['PROCESS']['bleed_1'] or name == config['PROCESS']['bleed_2']:
+					row['Compensation'] = float(row['Vol_Comp'])
+					self.models['bleed_music'].append(row)
+
+				# --> it's a Vocal model, so look at "self.Quality_Vocal" for references !!
+				if name == config['PROCESS']['bleed_3'] or name == config['PROCESS']['bleed_4']:
+					row['Compensation'] = float(row['Vol_Comp'])
+					self.models['bleed_vocal'].append(row)
+
+				if name == config['PROCESS']['bleed_5'] or name == config['PROCESS']['bleed_6']:
+					row['Compensation'] = float(row['Vol_Comp'])
+					self.models['remove_music'].append(row)
+					
+		# Download Models to :
+		models_path	= os.path.join(BASE_MODELS_DIR, "karafan")
+
+		for stem in self.models:
+			for model in self.models[stem]:				
+				model['Cut_OFF']		= int(model['Cut_OFF'])
+				model['N_FFT_scale']	= int(model['N_FFT_scale'])
+				model['dim_F_set']		= int(model['dim_F_set'])
+				model['dim_T_set']		= int(model['dim_T_set'])
+
+				model['PATH'] = Download_Model(model, models_path)
+					
+		# Load Models
+		if self.large_gpu: 
+			print("Large GPU mode is enabled : Loading models now...")
+
+			for stem in self.models:
+				for model in self.models[stem]:  self.Load_MDX(model)
+	
+		self.AudioFiles = [
+			"NORMALIZED",
+			# Vocals extraction
+			"Music extract",
+			"Vocal extract",
+			"Music Bleedings",
+			# Here --> Vocals results subtracted from Original Audio = Music
+			"Vocal Bleedings",
+			"Remove Music",
+			"Vocal FINAL",
+			"Music FINAL",
+		]
+		self.AudioFiles_Mandatory = [6, 7]  # Vocal FINAL & Music FINAL
+		self.AudioFiles_Debug = [0, 1, 2]	# Music & Vocal extract
+		
+		# DEBUG
+		# Reload "Bleedings" files with GOD MODE ... or not !
+		self.AudioFiles_Debug.append(3)
+		self.AudioFiles_Debug.append(4)
+		self.AudioFiles_Debug.append(5)
+
+	
+	# ******************************************************************
+	# ****    This is the MAGIC RECIPE , the heart of KaraFan !!    ****
+	# ******************************************************************
+
+	def SEPARATE(self, audio_file, BATCH_MODE=False):
+
+		name = os.path.splitext(os.path.basename(audio_file))[0]
+		
+		#*************************************************
+		#****        DEBUG  ->  for DEVELOPERS        ****
+		#*************************************************
+		
+		# Testing SDR :
+		# Put some "SDR_song_XXX.flac" from "Gdrive > KaraFan_user > Multi-Song" in your "Music" folder
+		# That's all !!
+		# (only the song file, not "instrum.flac" or "vocals.flac" from "Stems" folder)
+
+		self.SDR_Testing = name.startswith("SDR_")
+
+		# Set to False for fast testing (1 pass instead of 2 for denoising)
+		self.Denoise = True
+
+		# You can copy this line in between steps : Save intermediate files
+		# self.Save_Audio("X - Music SUB", normalized - vocal_ensemble)
+
+		#*************************************************
+
+		self.BATCH_MODE = BATCH_MODE
+		self.song_output_path = os.path.join(self.output, name)
+		
+		
+		start_time = time.time()
+
+		# Delete previous files
+		if os.path.exists(self.song_output_path):
+			if not self.GOD_MODE:
+				print("► No GOD MODE : Re-process ALL files ...")
+				for file in os.listdir(self.song_output_path):
+					if file != "SDR_Results.txt" and file.endswith(self.output_format):
+						os.remove(os.path.join(self.song_output_path, file))
+		else:
+			# Create a folder based on input audio file's name
+			os.makedirs(self.song_output_path,exist_ok=True)
+		
+		print(f"Go with : <b>{name}</b> {audio_file=}")
+
+		# Function checks if user's file is Bad encoded, then Resample to 44.1 Khz if needed
+
+		original_audio, self.sample_rate = audio_utils.Load_Audio(audio_file, 44100, self.ffmpeg, self.song_output_path)
+
+		if original_audio is None:
+			print("► Error : Can't correct your audio file !")  # Too much Corrupted file !
+			return
+		else: print(f"Audio loaded: {original_audio.shape=} {self.sample_rate=}")
+		
+		# TODO : Get the cut-off frequency of the input audio
+		# self.original_cutoff = audio_utils.Find_Cut_OFF(original_audio, self.sample_rate)
+		
+		self.original_cutoff = self.sample_rate // 2
+		
+		print(f"{original_audio.shape[-1] // 44100} sec. - Rate : {self.sample_rate} Hz / Cut-OFF : {self.original_cutoff} Hz")
+		
+		# ****  START PROCESSING  ****
+
+		if self.normalize < 0:
+			normalized = self.Check_Already_Processed(0)
+
+			if normalized is None:
+				print(f"► Normalizing audio : {self.normalize} dB")
+				normalized = audio_utils.Normalize(original_audio, self.normalize)
+
+				self.Save_Audio(0, normalized)
+		else:
+			normalized = original_audio
+		
+		# 1 - Extract Music with MDX models (Pre-Process before Vocals extraction)
+
+		if len(self.models['music']) == 0:
+			music_ensemble = None
+		else:
+			music_extracts = []
+			for model in self.models['music']:
+				audio = self.Check_Already_Processed(1, model['Name'])
+				if audio is None:
+					audio = self.Extract_with_Model("Music", normalized, model)
+					
+					self.Save_Audio(1, audio, model['Name'])
+				
+				music_extracts.append(audio)
+			
+			if len(music_extracts) == 1:
+				music_ensemble = music_extracts[0]
+			else:
+				print("► Make Ensemble Music")
+				music_ensemble = audio_utils.Make_Ensemble('Max', music_extracts)  # Algorithm ?
+
+				del music_extracts;  gc.collect()
+		
+		# 2 - Extract Vocals with MDX models
+
+		vocal_extracts = []
+		vocal_sub = normalized if music_ensemble is None else normalized - music_ensemble
+
+		for model in self.models['vocal']:
+			audio = self.Check_Already_Processed(2, model['Name'])
+			if audio is None:
+				audio = self.Extract_with_Model("Vocal", vocal_sub, model)
+				
+				self.Save_Audio(2, audio, model['Name'])
+			
+			vocal_extracts.append(audio)
+		
+		if len(vocal_extracts) == 1:
+			vocal_ensemble = vocal_extracts[0]
+		else:
+			print("► Make Ensemble Vocals")
+			vocal_ensemble = audio_utils.Make_Ensemble('Max', vocal_extracts)  # MAX and not average, because it's Vocals !!
+
+			del vocal_extracts;  gc.collect()
+		
+		# 3 - Pass Vocals through Filters (remove music bleedings)
+
+		if len(self.models['bleed_music']) == 0:
+			vocal_final = vocal_ensemble
+		else:
+			bleed_extracts = []
+			
+			for model in self.models['bleed_music']:
+				audio = self.Check_Already_Processed(3, model['Name'])
+				if audio is None:
+					audio = self.Extract_with_Model("Bleed_Music", vocal_ensemble, model)
+
+					self.Save_Audio(3, audio, model['Name'])
+				
+				bleed_extracts.append(audio)
+				
+			if len(bleed_extracts) == 1:
+				bleed_ensemble = bleed_extracts[0]
+			else:
+				print("► Make Ensemble Music Bleedings")
+				
+				bleed_ensemble = audio_utils.Make_Ensemble('Max', bleed_extracts)
+				
+			vocal_final = vocal_ensemble - bleed_ensemble
+
+			del vocal_ensemble; del bleed_extracts; del bleed_ensemble; gc.collect()
+
+		# 4 - Apply High Pass & Low Pass filters on Vocals
+
+		print(f"► Apply High {self.high_pass} Hz & Low {self.low_pass /1000:.1f} KHz Pass on Vocals")
+
+		if self.high_pass > 0:
+			vocal_final = audio_utils.Pass_filter('highpass', self.high_pass, vocal_final, self.sample_rate, order = 16)
+
+		if self.low_pass < 22000:
+			if self.low_pass > 17000:
+				vocal_final = audio_utils.Pass_filter('lowpass', self.low_pass, vocal_final, self.sample_rate, order = 16)
+			else:
+				vocal_final = audio_utils.Pass_filter('lowpass', self.low_pass, vocal_final, self.sample_rate, order = 8)
+		
+		# 5 - Get Music by subtracting Vocals from original audio (for instrumental not captured by MDX models)
+		
+		print("► Get Music by subtracting Vocals from original audio")
+
+		music_sub = normalized - vocal_final
+
+		# 6 - Pass Music SUB through Filters
+		
+		if len(self.models['bleed_vocal']) == 0:
+			music_final = music_sub
+		else:
+			if self.DEBUG:  self.Save_Audio("4 - Music - SUB", music_sub)
+
+			bleed_extracts = [];  music_extracts = []
+			
+			# A - Remove Vocal bleedings
+			
+			for model in self.models['bleed_vocal']:
+				audio = self.Check_Already_Processed(4, model['Name'])
+				if audio is None:
+					audio = self.Extract_with_Model("Bleed_Vocal", music_sub, model)
+
+					self.Save_Audio(4, audio, model['Name'])
+				
+				bleed_extracts.append(audio)
+				
+			if len(bleed_extracts) == 1:
+				bleed_ensemble = bleed_extracts[0]
+			else:
+				bleed_ensemble = audio_utils.Make_Ensemble('Max', bleed_extracts)
+				
+			# B - Remove Music bleedings
+
+			if len(self.models['remove_music']) > 0:
+				
+				for model in self.models['remove_music']:
+					audio = self.Check_Already_Processed(5, model['Name'])
+					if audio is None:
+						audio = self.Extract_with_Model("Bleed_Music", bleed_ensemble, model)
+
+						self.Save_Audio(5, audio, model['Name'])
+					
+					music_extracts.append(audio)
+					
+				if len(music_extracts) == 1:
+					bleed_ensemble = bleed_ensemble - music_extracts[0]
+				else:
+					bleed_ensemble = bleed_ensemble - audio_utils.Make_Ensemble('Max', music_extracts)  # Algorithm ?
+
+			music_final = music_sub - bleed_ensemble
+
+			del bleed_extracts; del music_extracts; del bleed_ensemble; gc.collect()
+
+		# 7 - FINAL saving
+		
+		print("► Save Vocals FINAL !")
+
+		# Apply Infra Bass filter
+		if self.infra_bass:  vocal_final = audio_utils.Pass_filter('highpass', 18, vocal_final, self.sample_rate, order = 100)
+
+		# Apply silence filter
+		if self.silent < 0:  vocal_final = audio_utils.Silent(vocal_final, self.sample_rate, self.silent)
+
+		self.Save_Audio(6, vocal_final)
+		
+		output_vocals = remix_audio((vocal_final, self.sample_rate),to_int16=True)
+
+		print("► Save Music FINAL !")
+		
+		# Apply Infra Bass filter
+		if self.infra_bass:  music_final = audio_utils.Pass_filter('highpass', 18, music_final, self.sample_rate, order = 100)
+
+		# Apply silence filter
+		if self.silent < 0:  music_final = audio_utils.Silent(music_final, self.sample_rate, self.silent)
+
+		self.Save_Audio(7, music_final)
+
+		output_music = remix_audio((music_final, self.sample_rate),to_int16=True)
+
+		print('<b>--> Processing DONE !</b>')
+
+		# # Clear screen between each song
+		
+		elapsed_time = time.time() - start_time
+		elapsed_time = f"Elapsed Time for <b>{name}</b> : {time.strftime('%H:%M:%S', time.gmtime(elapsed_time))} sec.<br>"
+		print(elapsed_time)
+
+		if self.SDR_Testing:
+			elapsed_time = regex.sub(r"<.*?>", "", elapsed_time)   # Remove HTML tags
+			print("----------------------------------------")
+			compare.SDR(self.song_output_path, self.output_format, self.Gdrive, elapsed_time)
+
+		return output_vocals, output_music, (original_audio, self.sample_rate)
+		
+
+	def Extract_with_Model(self, type, audio, model):
+		"""
+		Explication from "Jarredou" about the 2 passes :
+
+		This helps reduce/remove the noise added by the MDX models,
+		since the phase is inverted before processing and restored afterward in one of the two passes.
+		When they are added together, only the MDX noise is out of phase and gets removed,
+		while the rest regains its original gain (0.5 + 0.5).
+		ZFTurbo also added this to Demucs in the original MVSep-MDX23 code.
+		"""
+		name = model['Name']
+
+		if type == 'Music':			quality = self.Quality_Music;  text = 'Extract Music'
+		elif type == 'Vocal':			quality = self.Quality_Vocal;  text = 'Extract Vocals'
+		elif type == 'Bleed_Music':		quality = self.Quality_Bleed;  text = 'Remove Music Bleedings'
+		elif type == 'Bleed_Vocal':		quality = self.Quality_Bleed;  text = 'Vocal Bleedings in Music'
+		elif type == 'Remove_Music':	quality = self.Quality_Bleed;  text = 'Music in Vocal Bleedings'
+		
+		text	  = f'► {text} with "{name}"'
+
+		# ONLY 1 Pass, with MDX23C models !
+		if model['Stem'] == "BOTH":
+			device = torch.device(self.GPU_device)
+			mdx23 = tfc_tdf.TFC_TDF_net(self.MDX23_config)
+			mdx23.load_state_dict( torch.load(model['PATH'], map_location = device) )
+			mdx23 = mdx23.to(device)
+			mdx23.eval()
+
+			overlap = self.MDX23_overlap if (type == "Music" or type == "Vocal") else self.MDX23_bleed
+
+			print(f"{text} (Overlap : {overlap})")
+			sources = demix_mdxv3(audio, mdx23, device, self.MDX23_config, overlap)
+			
+			source = sources['Vocals'] if "Vocal" in type else sources['Instrumental']
+			
+			# Check if source is same size than audio
+			source = librosa.util.fix_length(source, size = audio.shape[-1])
+
+			del sources; mdx23 = mdx23.cpu(); del mdx23; gc.collect()
+
+			if  torch.cuda.is_available():
+				torch.cuda.empty_cache()
+				torch.cuda.ipc_collect()
+			elif torch.backends.mps.is_available():
+				torch.mps.empty_cache()
+
+		else:
+			if not self.large_gpu:
+				# print(f'Loading model "{name}" ...')
+				self.Load_MDX(model)
+			
+			mdx_model = self.MDX[name]['model']
+			inference = self.MDX[name]['inference']
+
+			bigshifts = quality['BigShifts'] 
+			
+			if self.Denoise:
+				print(f"{text} ({quality['BigShifts']} pass)")
+				source  = 0.5 * -self.demix_full(-audio, mdx_model, inference, bigshifts, 1)[0]
+				source += 0.5 *  self.demix_full( audio, mdx_model, inference, bigshifts, 2)[0]
+			else:
+				# ONLY 1 Pass, for testing purposes
+				print(f"{text} ({quality['BigShifts']} pass) - <b>NO Denoise !</b>")
+				source = self.demix_full(audio, mdx_model, inference, bigshifts, 0)[0]
+
+			# Automatic SRS (for not FULL-BAND models !)
+			if quality['BigShifts_SRS'] > 0:
+
+				bigshifts = quality['BigShifts_SRS']
+
+				# 1 - High SRS
+
+				if model['Cut_OFF'] > 0 and model['Name'] != "Vocal Main":  # Exception !!
+
+					# This is mandatory, I don't know why, but without this,
+					# the sample rate DOWN doesn't fit the MDX model Band
+					# and produce noise in High Frequencies !! (??)
+					# @ 510 Hz -> there is less noise, but badder SDR ?!?!
+					# TODO : Test with 14600 Hz models cut-off
+					# 
+					delta = 810 if type == 'Vocal' else 1220 # Hz
+
+					audio_SRS = audio_utils.Change_sample_rate(audio, 'DOWN', self.original_cutoff, model['Cut_OFF'] + delta)
+					
+					# Limit audio to the same frequency cut-off than MDX model : To avoid SRS noise !! (That helps a little bit)
+					audio_SRS = audio_utils.Pass_filter('lowpass', model['Cut_OFF'], audio_SRS, self.sample_rate, order = 100)
+
+					# DEBUG
+					# self.Save_Audio(type + " - SRS REAL - High", audio_SRS)
+
+					if self.Denoise:
+						print(f"{text} -> SRS High ({bigshifts} pass)")
+						
+						source_SRS = 0.5 * audio_utils.Change_sample_rate(
+							-self.demix_full(-audio_SRS, mdx_model, inference, bigshifts, 1)[0], 'UP', self.original_cutoff, model['Cut_OFF'] + delta)
+						
+						source_SRS += 0.5 * audio_utils.Change_sample_rate(
+							self.demix_full( audio_SRS, mdx_model, inference, bigshifts, 2)[0], 'UP', self.original_cutoff, model['Cut_OFF'] + delta)
+					else:
+						# ONLY 1 Pass, for testing purposes
+						print(f"{text} -> SRS High ({bigshifts} pass) - <b>NO Denoise !</b>")
+						source_SRS = audio_utils.Change_sample_rate(
+							self.demix_full( audio_SRS, mdx_model, inference, bigshifts, 0)[0], 'UP', self.original_cutoff, model['Cut_OFF'] + delta)
+
+					# Check if source_SRS is same size than source
+					source_SRS = librosa.util.fix_length(source_SRS, size = source.shape[-1])
+
+					if type == 'Vocal':
+						source = audio_utils.Make_Ensemble('Max', [source, source_SRS])
+					else:
+						# OLD formula --> from Jarredou
+						
+						# vocals = Linkwitz_Riley_filter(vocals.T, 12000, 'lowpass') + Linkwitz_Riley_filter((3 * vocals_SRS.T) / 4, 12000, 'highpass')
+						# *3/4 = Dynamic SRS personal taste of "Jarredou", to avoid too much SRS noise
+						# He also told me that 12 Khz cut-off was setted for MDX23C model, but now I use the REAL cut-off of MDX models !
+
+						# Avec cutoff = 17640 hz & -60dB d'atténuation et ordre = 12 --> cut freq = 16000 hz (-1640)
+						# cut_freq = 14400 # Hz
+						# # cut_freq = 7500 # Hz
+						# # if model['Name'] == "Kim Instrum":  cut_freq = 12000 # Hz
+						source = audio_utils.Linkwitz_Riley_filter('lowpass',  16000, source,     self.sample_rate, order=12) + \
+								audio_utils.Linkwitz_Riley_filter('highpass', 16000, source_SRS, self.sample_rate, order=12)
+
+					# # new multiband ensemble
+					# vocals_low = lr_filter((weights[0] * vocals_mdxb1.T + weights[1] * vocals3.T + weights[2] * vocals_mdxb2.T) / weights.sum(), 12000, 'lowpass', order=12)
+					# vocals_mid = lr_filter(lr_filter((2 * vocals_mdxb2.T + 2 * vocals_SRS.T + vocals_demucs.T) / 5, 16500, 'lowpass', order=24), 12000, 'highpass', order=12)
+					# vocals_high = lr_filter((vocals_demucs.T + vocals_SRS.T) / 2, 16500, 'highpass', order=24)
+					# vocals = (vocals_low + vocals_mid + vocals_high) * 1.0074
+
+					del audio_SRS; del source_SRS; gc.collect()
+					
+				# 2 - Low SRS -> Bigshifts only 1 pass, else bad SDR
+				
+				if type == 'Vocal':
+					
+					cut_freq = 18550 # Hz
+
+					audio_SRS = audio_utils.Change_sample_rate(audio, 'UP', self.original_cutoff, cut_freq)
+					
+					# Limit audio to frequency cut-off (That helps a little bit)
+					audio_SRS = audio_utils.Pass_filter('lowpass', model['Cut_OFF'], audio_SRS, self.sample_rate, order = 100)
+
+					# DEBUG
+					# self.Save_Audio(type + " - SRS REAL - Low", audio_SRS)
+
+					# ONLY 1 Pass, for testing purposes
+					if self.Denoise:
+						print(f"{text} -> SRS Low (1 pass)")
+						
+						source_SRS = 0.5 * audio_utils.Change_sample_rate(
+							-self.demix_full(-audio_SRS, mdx_model, inference, 1, 1)[0], 'DOWN', self.original_cutoff, cut_freq)
+
+						source_SRS += 0.5 * audio_utils.Change_sample_rate(
+							self.demix_full( audio_SRS, mdx_model, inference, 1, 2)[0], 'DOWN', self.original_cutoff, cut_freq)
+					else:
+						print(f"{text} -> SRS Low (1 pass) - <b>NO Denoise !</b>")
+						source_SRS = audio_utils.Change_sample_rate(
+							self.demix_full( audio_SRS, mdx_model, inference, 1, 0)[0], 'DOWN', self.original_cutoff, cut_freq)
+
+					# Check if source_SRS is same size than source
+					source_SRS = librosa.util.fix_length(source_SRS, size = source.shape[-1])
+
+					source = audio_utils.Make_Ensemble('Max', [source, source_SRS])
+
+					del audio_SRS; del source_SRS; gc.collect()
+
+			if not self.large_gpu:  self.Kill_MDX(name)
+
+		source = source * model['Compensation']  # Volume Compensation
+
+		# TODO
+		# source = audio_utils.Remove_High_freq_Noise(source, model['Cut_OFF'])
+
+		return source
+	
+
+	def Load_MDX(self, model):
+		name = model['Name']
+		if name not in self.MDX:
+			self.MDX[name] = {}
+			self.MDX[name]['model'] = tfc_tdf.Conv_TDF_net_trim_model(self.GPU_device, '', 11, model)  # can use '*' to get 4 stems ??
+			self.MDX[name]['inference'] = ort.InferenceSession(
+				model['PATH'],
+				providers = self.providers,
+				provider_options = [{"device_id": 0}]
+			)
+	
+	def Kill_MDX(self, model_name):
+		if model_name in self.MDX:
+			del self.MDX[model_name]['inference']
+			del self.MDX[model_name]['model']
+			del self.MDX[model_name]
+			gc.collect()
+			if  torch.cuda.is_available():
+				torch.cuda.empty_cache()
+				torch.cuda.ipc_collect()
+			elif torch.backends.mps.is_available():
+				torch.mps.empty_cache()
+
+	def raise_aicrowd_error(self, msg):
+		# Will be used by the evaluator to provide logs, DO NOT CHANGE
+		raise NameError(msg)
+	
+		
+	def Check_Already_Processed(self, key, model_name = ""):
+		"""
+		if GOD MODE :
+			- Check if audio file is already processed, and if so, load it.
+			- Return AUDIO loaded, or NONE if not found.
+		Else :
+			- Return NONE.
+		Key :
+			index of AudioFiles list or "str" (direct filename for test mode)
+		"""
+		if not self.GOD_MODE or key not in self.AudioFiles_Debug:  return None
+
+		filename = self.AudioFiles[key]
+		if self.DEBUG:  filename = f"{key} - {filename}"
+		if model_name != "":  filename += " - ("+ model_name +")"
+
+		if self.output_format == 'PCM_16':	filename += '.wav'
+		elif self.output_format == 'FLOAT':	filename += '.wav'
+		elif self.output_format == 'WAV':	filename += '.wav'
+		elif self.output_format == "FLAC":	filename += '.flac'
+		elif self.output_format == 'MP3':	filename += '.mp3'
+
+		file = os.path.join(self.song_output_path, filename)
+		
+		if os.path.isfile(file):
+			
+			print(filename + " --> Loading ...")
+			audio, _ = audio_utils.Load_Audio(file, self.sample_rate)
+			
+			return audio
+		
+		return None
+	
+	def Save_Audio(self, key, audio, model_name = ""):
+		"""
+		Key : index of AudioFiles list or "str" (direct filename for test mode)
+		if Key is a string, it will force saving !
+		"""
+		
+		# Save only mandatory files if not in DEBUG mode
+		if type(key) is int:
+			if self.DEBUG:
+				if key not in self.AudioFiles_Debug and key not in self.AudioFiles_Mandatory:  return
+			else:
+				if key not in self.AudioFiles_Mandatory:  return
+
+		if type(key) is int:
+			filename = self.AudioFiles[key]
+			if self.DEBUG:  filename = f"{key} - {filename}"
+		else:
+			filename = key
+
+		if model_name != "":  filename += " - ("+ model_name +")"
+
+		file = os.path.join(self.song_output_path, filename)
+		
+		audio_utils.Save_Audio(file, audio, self.sample_rate, self.output_format, self.original_cutoff, self.ffmpeg)
+
+	
+	def demix_full(self, mix, use_model, infer_session, bigshifts, pass_number = 0):
+		
+		results = []
+		mix_length = int(mix.shape[1] / 44100)
+
+		# Don't depass the length of the song in seconds
+
+		if bigshifts < 1:  bigshifts = 1  # must not be <= 0 !
+		if bigshifts > mix_length:  bigshifts = mix_length - 1
+		
+		# "demix_seconds" is equal to the number of seconds to shift the mix
+		# and it's same of "bigshifts" for average results
+		demix_seconds = bigshifts
+		while bigshifts * demix_seconds > mix_length:  demix_seconds -= 1
+
+		shifts  = [x * demix_seconds for x in range(bigshifts)]
+
+		for shift in shifts:
+			
+			shift_samples = int(shift * 44100)
+			# print(f"shift_samples = {shift_samples}")
+			
+			shifted_mix = np.concatenate((mix[:, -shift_samples:], mix[:, :-shift_samples]), axis=-1)
+			# print(f"shifted_mix shape = {shifted_mix.shape}")
+			result = np.zeros((1, 2, shifted_mix.shape[-1]), dtype=np.float32)
+			divider = np.zeros((1, 2, shifted_mix.shape[-1]), dtype=np.float32)
+
+			total = 0
+			for i in range(0, shifted_mix.shape[-1], self.chunk_size):
+				total += 1
+
+				start = i
+				end = min(i + self.chunk_size, shifted_mix.shape[-1])
+				mix_part = shifted_mix[:, start:end]
+				# print(f"mix_part shape = {mix_part.shape}")
+				sources = demix_base(mix_part, self.GPU_device, use_model, infer_session)
+				result[..., start:end] += sources
+				# print(f"result shape = {result.shape}")
+				divider[..., start:end] += 1
+			
+			result /= divider
+			# print(f"result shape = {result.shape}")
+			result = np.concatenate((result[..., shift_samples:], result[..., :shift_samples]), axis=-1)
+			results.append(result)
+			
+		results = np.mean(results, axis=0)
+		return results
+	
+	#----
+
+def Download_Model(model, models_path):
+	
+	name		= model['Name']
+	repo_file	= model['Repo_FileName']
+	filename	= regex.sub(r"^(UVR-MDX-NET-|UVR_MDXNET_|\d_)*", "", repo_file)
+	file_path	= os.path.join(models_path, filename)
+
+	if not os.path.isfile(file_path):
+		print(f'Downloading model : "{name}" ...')
+
+		remote_url = 'https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/' + repo_file
+		try:
+			response = requests.get(remote_url, stream=True)
+			response.raise_for_status()  # Raise an exception in case of HTTP error code
+			
+			if response.status_code == 200:
+				
+				# total_size = int(response.headers.get('content-length', 0)) // 1048576  # MB
+				# PROGRESS.reset(total_size, unit="MB")
+				
+				with open(file_path, 'wb') as file:
+
+					for data in response.iter_content(chunk_size=1048576):
+						# PROGRESS.update()
+						file.write(data)
+			else:
+				print(f'Download of model "{name}" FAILED !!')
+		
+		except (requests.exceptions.RequestException, requests.exceptions.ChunkedEncodingError) as e:
+			print(f'Error during Downloading "{name}" !!\n\n{e}')
+			if os.path.exists(file_path):  os.remove(file_path)
+	
+	return file_path  # Path to this model
+
+
+def Process(input_file, cache_dir=None, use_cache=True, format=None, config=settings.Defaults, **kwargs):
+	model = MusicSeparationModel(config,cache_dir=cache_dir,use_cache=use_cache,format=format)
+
+	output = model.SEPARATE(input_file)
+	del model
+	gc_collect()
+	return output

lib/karafan/settings.py

@@ -0,0 +1,177 @@
+
+#   MIT License - Copyright (c) 2023 - Captain FLAM
+#
+#   https://github.com/Captain-FLAM/KaraFan
+
+import os, configparser
+
+# Default values
+Defaults = {
+	'AUDIO': {
+		'input': "Music",
+		'normalize': "-1",
+		'output': "Music",
+		'output_format': "FLAC",
+		'silent': "0",
+		'infra_bass': False,
+	},
+	'PROCESS': {
+		'music_1': "----",
+		'music_2': "----",
+		'vocal_1': "MDX23C 8K FFT",
+		'vocal_2': "----",
+		'bleed_1': "----",
+		'bleed_2': "----",
+		'high_pass': 14,
+		'low_pass': 5,
+
+		'bleed_3': "----",
+		'bleed_4': "----",
+		'bleed_5': "----",
+		'bleed_6': "----",
+	},
+	'OPTIONS': {
+		'speed': "Medium",
+		'chunk_size': 500000,
+	},
+	'BONUS': {
+		'DEBUG': True,
+		'GOD_MODE': True,
+		'KILL_on_END': False,
+		'THEME': "",
+		'large_gpu': False,
+	},
+}
+Options = {
+	'Normalize': [("NONE", "0"), ("-0.5 dB", "-0.5"), ("- 1 dB", "-1"), ("- 2 dB", "-2"), ("- 3 dB", "-3")],
+	'Format':	[("FLAC - 24 bits", "FLAC"), ("MP3 - CBR 320K", "MP3"), ("WAV - PCM 16 bits","PCM_16"), ("WAV - FLOAT 32 bits","FLOAT")],
+	'Silent':	[("NONE", "0"), ("- 45 dB", "-45"), ("- 50 dB", "-50"), ("- 55 dB", "-55"), ("- 60 dB", "-60")],
+	'Speed':	['Fastest', 'Fast', 'Medium', 'Slow', 'Slowest'],
+}
+
+Help_Dico = {
+	'input':		"- IF « Input » is a folder path, ALL audio files inside this folder will be separated by a Batch processing.<br>- Else, only the selected audio file will be processed.",
+	'normalize': 	"Normalize input audio files to avoid <b>clipping</b> and get better results.<br>Normally, <b>you do not have</b> to use this option !!<br>Only for weak or loud songs !",
+	'output':		"« Output folder » will be created based on the file's name without extension.<br>For example : if your audio input is named : « 01 - Bohemian Rhapsody<b>.MP3</b> »,<br>then output folder will be named : « 01 - Bohemian Rhapsody »",
+	'format':		"Choose your prefered audio format to save audio files.",
+	'silent':		"Make silent the parts of audio where dynamic range (RMS) goes below threshold.<br>Don't misundertand : this function is NOT a noise reduction !<br>Its behavior is to clean the final audios from « silent parts » (below -XX dB).",
+	'infra_bass':	"This will remove <b>Infra-Bass</b> in your audio files (below 15 Hz).<br>It will leave more place to others frequencies and improve the quality of your audio files.<br>It will also reduce the size of your audio files.",
+	'MDX_music':	"Make an Ensemble of extractions with Instrumental selected models.<br>This is used to remove <b>Music</b> before Vocal extractions.",
+	'MDX_vocal':	"Make an Ensemble of extractions with Vocals selected models.",
+	'MDX_bleed_1':	"Remove <b>Music Bleedings</b> in Vocal extractions.<br><br>DON'T use « <b>Instrum HQ 3</b> » as it catchs too much vocals !!",
+	'vocal_pass':	"Set the <b>Passband</b> for <b>Vocal</b> extractions to get more <b>Music</b> back !<br>Defaults : <b>70 Hz</b> (+ slope 15 Hz) & <b>16.5 KHz</b> for this frequencies (Bass - Mid / Harmonics) :<br>Male : 85 - 180 Hz / 14 K , Female : 165 - 255 Hz / 16 K , Children : 250 - 400 Hz / 20 K",
+	'MDX_bleed_2':	"Remove <b>Vocal Bleedings</b> in <b>subtracted</b> Music.<br><br>DON'T use « <b>Kim Vocal 2</b> » or « <b>Voc FT</b> » as they remove too much music !!",
+	'MDX_bleed_3':	"Remove <b>Music Bleedings</b> in Vocal Bleedings extracted <b>above</b> to get them back in Music !<br>DON'T use « Instrum HQ 3 » as it catchs too much Vocals ! 😉<br>... and : <b>ALL models</b> will carry more or less Vocal bleedings in Music Final !!",
+	'speed':		"Fastest : extract in 1 pass with <b>NO</b> SRS (<b>only</b> for Testing)<br>All « Speed » are processed with <b>DENOISE</b> (the same option as in <b>UVR 5</b> 😉)<br>Slowest : is the best quality, but it will take hours to process !! 😝",
+	'chunks':		"Chunk size for ONNX models. (default : 500,000)<br><br>Set lower to reduce GPU memory consumption OR <b>if you have GPU memory errors</b> !",
+	'debug':		"IF checked, it will save all intermediate audio files to compare in your <b>Audacity</b>.",
+	'god_mode':		"Give you the GOD's POWER : each audio file is reloaded IF it was created before,<br>NO NEED to process it again and again !!<br>You'll be warned : You have to <b>delete MANUALLY</b> each file that you want to re-process !",
+	'theme':		"Darth Vader or Stormtrooper ?<br><br>Choose your prefered theme for KaraFan !",
+	'kill_end':		"On <b>Colab</b> : KaraFan will KILL your session at end of « Processongs », to save your credits !!<br>On <b>your Laptop</b> : KaraFan will KILL your GPU, to save battery (and hot-less) !!<br>On <b>your PC</b> : KaraFan will KILL your GPU, anyway ... maybe it helps ? Try it !!",
+#	'large_gpu':	"It will load ALL models in GPU memory for faster processing of MULTIPLE audio files.<br>Requires more GB of free GPU memory.<br>Uncheck it if you have memory troubles.",
+}
+Presets = [
+	{
+		'music_1': "MDX23C 8K FFT",
+		'music_2': "----",
+		'vocal_1': "MDX23C 8K FFT",
+		'vocal_2': "----",
+		'bleed_1': "----",
+		'bleed_2': "----",
+		'bleed_3': "----",
+		'bleed_4': "----",
+		'bleed_5': "----",
+		'bleed_6': "----",
+	},
+	{
+		'music_1': "----",
+		'music_2': "----",
+		'vocal_1': "MDX23C 8K FFT",
+		'vocal_2': "----",
+		'bleed_1': "Instrum 3",
+		'bleed_2': "Instrum Main",
+		'bleed_3': "----",
+		'bleed_4': "----",
+		'bleed_5': "----",
+		'bleed_6': "----",
+	},
+	{
+		'music_1': "----",
+		'music_2': "----",
+		'vocal_1': "MDX23C 8K FFT",
+		'vocal_2': "----",
+		'bleed_1': "----",
+		'bleed_2': "----",
+		'bleed_3': "MDX23C 8K FFT",
+		'bleed_4': "----",
+		'bleed_5': "Instrum 3",
+		'bleed_6': "Instrum Main",
+	},
+	{
+		'music_1': "----",
+		'music_2': "----",
+		'vocal_1': "MDX23C 8K FFT",
+		'vocal_2': "----",
+		'bleed_1': "Instrum 3",
+		'bleed_2': "Instrum Main",
+		'bleed_3': "MDX23C 8K FFT",
+		'bleed_4': "----",
+		'bleed_5': "Instrum 3",
+		'bleed_6': "Instrum Main",
+	},
+	{
+		'music_1': "----",
+		'music_2': "----",
+		'vocal_1': "MDX23C 8K FFT",
+		'vocal_2': "Voc FT",
+		'bleed_1': "Instrum Main",
+		'bleed_2': "----",
+		'bleed_3': "MDX23C 8K FFT",
+		'bleed_4': "----",
+		'bleed_5': "----",
+		'bleed_6': "----",
+	},
+]
+
+def Load(Gdrive, isColab):
+
+	file = os.path.join(Gdrive, "KaraFan_user", "Config_Colab.ini" if isColab else "Config_PC.ini")
+	
+	if not os.path.isfile(file):
+		config = Save(Gdrive, isColab, Defaults)
+	else:
+		config = configparser.ConfigParser()
+		config.optionxform = lambda option: option  # To preserve case of Keys !!
+		config.read(file, encoding='utf-8')
+		
+		# Load default values if not present
+		for section in Defaults:
+			if section not in config:
+				config[section] = {}
+			for key in Defaults[section]:
+				if key not in config[section]:
+					config[section][key] = str(Defaults[section][key])
+	
+	if  config['AUDIO']['normalize'] not in [x[1] for x in Options['Normalize']]:
+		config['AUDIO']['normalize'] = Defaults['AUDIO']['normalize']
+	if  config['AUDIO']['output_format'] not in [x[1] for x in Options['Format']]:
+		config['AUDIO']['output_format'] = Defaults['AUDIO']['output_format']
+	if  config['AUDIO']['silent'] not in [x[1] for x in Options['Silent']]:
+		config['AUDIO']['silent'] = Defaults['AUDIO']['silent']
+	if  config['OPTIONS']['speed'] not in Options['Speed']:
+		config['OPTIONS']['speed'] = Defaults['OPTIONS']['speed']
+	
+	return config
+
+def Save(Gdrive, isColab, config_dict):
+	
+	file = os.path.join(Gdrive, "KaraFan_user", "Config_Colab.ini" if isColab else "Config_PC.ini")
+
+	config = configparser.ConfigParser()
+	config.optionxform = lambda option: option  # To preserve case of Keys !!
+	config.read_dict(config_dict)
+
+	with open(file, 'w', encoding='utf-8') as config_file:
+		config.write(config_file)
+	
+	return config

lib/karafan/setup.py

@@ -0,0 +1,102 @@
+
+#   MIT License - Copyright (c) 2023 - Captain FLAM
+#
+#   https://github.com/Captain-FLAM/KaraFan
+
+import os, gc, requests, subprocess
+
+def Check_dependencies():
+
+	# Dependencies already installed ?
+	print("Installing dependencies... This will take few minutes...", end='')
+	try:
+		subprocess.run(["pip", "install", "-r", "requirements.txt"], text=True, capture_output=True, check=True)
+		
+		print("\rInstallation done !                                     ") # Clean line
+	
+	except subprocess.CalledProcessError as e:
+		print("Error during Install dependencies :\n" + e.stderr + "\n" + e.stdout + "\n")
+		Exit_Notebook()
+
+def Install(params):
+	
+	# This is setup is only for Colab !!
+
+	if params['isColab'] == False:  return
+
+	Repository  = "https://github.com/Captain-FLAM/KaraFan"
+	Version_url = "https://raw.githubusercontent.com/Captain-FLAM/KaraFan/master/App/__init__.py"
+
+	Version = ""; Git_version = ""
+
+	Gdrive = params['Gdrive']
+	Project = params['Project']
+	DEV_MODE = params['I_AM_A_DEVELOPER']
+
+	if not os.path.exists(Gdrive):
+		print("ERROR : Google Drive path is not valid !\n")
+		Exit_Notebook()
+	
+	# Create missing folders
+	user_folder = os.path.join(Gdrive, "KaraFan_user")
+	os.makedirs(user_folder, exist_ok=True)
+	os.makedirs(os.path.join(user_folder, "Models"), exist_ok=True)
+
+	os.chdir(Project)  # For pip install
+
+	Check_dependencies()
+	
+	# Get local version
+	with open(os.path.join(Project, "App", "__init__.py"), "r") as version_file:
+		Version = version_file.readline().replace("# Version", "").strip()
+
+	# Auto-Magic update !
+	try:
+		response = requests.get(Version_url)
+		if response.status_code == requests.codes.ok:
+			Git_version = response.text.split('\n')[0].replace("# Version", "").strip()
+		else:
+			print("Unable to check version on GitHub ! Maybe you're behind a firewall ?")
+	except ValueError as e:
+		print("Error processing version data :", e)
+	except requests.exceptions.ConnectionError as e:
+		print("Connection error while trying to fetch version :", e)
+
+	if Version and Git_version:
+		if Git_version > Version:
+			print(f'A new version of "KaraFan" is available : {Git_version} !')
+
+			warning = 'You have to download the new version manually from :\n'
+			warning += Repository
+			warning +='\n... and extract it in your Project folder.\n'
+			warning +='Then, you have to "Restart" the notebook to use the new version of "KaraFan" !\n\n'
+			
+			if DEV_MODE:
+				print(warning)
+			else:
+				if os.path.exists(os.path.join(Project, ".git")):
+					try:
+						subprocess.run(["git", "-C", Project, "pull"], text=True, capture_output=True, check=True)
+
+						Check_dependencies()
+
+						print('\n\nFOR NOW : you have to go AGAIN in Colab menu, "Runtime > Restart and Run all" to use the new version of "KaraFan" !\n\n')
+
+						Exit_Notebook()
+						
+					except subprocess.CalledProcessError as e:
+						if e.returncode == 127:
+							print('WARNING : "Git" is not installed on your system !\n' + warning)
+						else:
+							print("Error during Update :\n" + e.stderr + "\n" + e.stdout)
+							Exit_Notebook()
+		else:
+			print('"KaraFan" is up to date.')
+
+def Exit_Notebook():
+	gc.collect()
+	# This trick is copyrigthed by "Captain FLAM" (2023) - MIT License
+	# That means you can use it, but you have to keep this comment in your code.
+	# After deep researches, I found this trick that nobody found before me !!!
+	from google.colab import runtime
+	runtime.unassign()

lib/karafan/sys_info.py

@@ -0,0 +1,121 @@
+
+#   MIT License - Copyright (c) 2023 - Captain FLAM
+#
+#   https://github.com/Captain-FLAM/KaraFan
+
+import os, sys, subprocess, json, regex
+from distutils import spawn
+
+def Get(font_size):
+
+	import platform, psutil
+
+	system = platform.system()
+
+	html  = '<pre style="font: bold '+ font_size +' monospace; line-height: 1.4">'
+	html += "****  System Informations  ****<br><br>"
+
+	# Get the total virtual memory size (in bytes)
+	total_virtual_memory = psutil.virtual_memory().total
+	unit_index = 0
+	units = ['B', 'KB', 'MB', 'GB', 'TB']
+
+	# Convert size into larger units until size is less than 1024
+	while total_virtual_memory >= 1024:
+		total_virtual_memory /= 1024
+		unit_index += 1
+
+	html += "Python : "+ regex.sub(r'\(.*?\)\s*', '', sys.version) +"<br>"
+	html += f"OS : {system} { platform.release() }<br>"
+	html += f"RAM : {total_virtual_memory:.2f} {units[unit_index]}<br>"
+	html += f"Current directory : { os.getcwd() }<br><br>"
+
+	html += "****    CPU Informations    ****<br><br>"
+	match system:
+		case 'Windows':  # use 'wmic'
+			try:
+				cpu_info = subprocess.check_output(['wmic', 'cpu', 'get', 'Caption,MaxClockSpeed,NumberOfCores,NumberOfLogicalProcessors', '/FORMAT:CSV'], shell=True, stderr=subprocess.STDOUT).decode('utf-8')
+				cpu_info = cpu_info.split('\n')[-2].strip()  # catch the last line
+				# Split values
+				cpu_info = cpu_info.split(',')
+				html += f"CPU : {cpu_info[1]}<br>"
+				html += f"Cores : {cpu_info[3]}<br>"
+				html += f"Threads : {cpu_info[4]}<br>"
+				html += f"MaxClock Speed : {cpu_info[2]} MHz"
+			except FileNotFoundError:
+				html += "--> Can't get CPU infos : 'wmic' tool is not available on this platform."
+
+		case 'Linux':  # use 'lscpu'
+			try:
+				cpu_info = subprocess.check_output(['lscpu', '-J'], shell=True, stderr=subprocess.STDOUT).decode('utf-8')
+				try:
+					cpu_info = json.loads(cpu_info)
+					if 'lscpu' in cpu_info:
+						sockets = cores = threads = 1
+						for item in cpu_info["lscpu"]:
+							if 'field' in item and 'data' in item:
+								data = item['data']
+								match item['field']:
+									case "Architecture:":		html += f"Arch : {data}<br>"
+									case "Model name:":			html += f"CPU : {data}<br>"
+									case "CPU max MHz:":		html += f"MaxClock Speed : {int(data)} MHz<br>"
+									case "Socket(s):":			sockets = int(data)
+									case "Core(s) per socket:":	cores   = int(data)
+									case "Thread(s) per core:":	threads = int(data)
+						
+						html += f"Cores : {cores * sockets}<br>"
+						html += f"Threads : {threads * cores * sockets}"
+
+				except json.decoder.JSONDecodeError:
+					html += "--> Can't get CPU infos : 'lscpu' doesn't work correctly on this platform ??"
+					
+			except FileNotFoundError:
+				html += "--> Can't get CPU infos : 'lscpu' tool is not available on this platform."
+
+		case 'Darwin':  # For macOS, use 'sysctl'
+			try:
+				## TODO : decode CPU infos on macOS
+				html += "CPU : " + subprocess.check_output(['sysctl', 'machdep.cpu'], shell=True, stderr=subprocess.STDOUT).decode('utf-8')
+			except FileNotFoundError:
+				html += "--> Can't get CPU infos : 'sysctl' tool is not available on this platform."
+
+		case _:
+			# For other platforms, display a generic message
+			html += "--> CPU informations are not available for this platform."
+
+	html += "<br><br>****    GPU Informations    ****<br><br>"
+	try:
+		if platform.system() == "Windows":
+			# If the platform is Windows and nvidia-smi 
+			# could not be found from the environment path, 
+			# try to find it from system drive with default installation path
+			nvidia_smi = spawn.find_executable('nvidia-smi')
+			if nvidia_smi is None:
+				nvidia_smi = "%s\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe" % os.environ['systemdrive']
+		else:
+			nvidia_smi = "nvidia-smi"
+
+		# Nvidia details information
+		gpu_info = subprocess.check_output(nvidia_smi, shell=True, stderr=subprocess.STDOUT).decode('utf-8')
+		
+		html += '<div style="line-height: 1; ">'+ gpu_info +'</div><br>'
+
+		if gpu_info.find('failed') >= 0:
+			html += "<b>GPU runtime is disabled. You can only use your CPU with available RAM.</b>"
+		elif gpu_info.find('Tesla T4') >= 0:
+			html += "You got a Tesla T4 GPU. (speeds are around  10-25 iterations/sec)"
+		elif gpu_info.find('Tesla P4') >= 0:
+			html += "You got a Tesla P4 GPU. (speeds are around  8-22 iterations/sec)"
+		elif gpu_info.find('Tesla K80') >= 0:
+			html += "You got a Tesla K80 GPU. (This is the most common and slowest gpu, speeds are around 2-10 iterations/sec)"
+		elif gpu_info.find('Tesla P100') >= 0:
+			html += "You got a Tesla P100 GPU. (This is the FASTEST gpu, speeds are around  15-42 iterations/sec)"
+	
+	except subprocess.CalledProcessError as e:
+		print(f"Erreur lors de l'exécution de nvidia-smi : {e.returncode}\n{e.output.decode('utf-8')}")
+	except FileNotFoundError:
+		html += "<b>--> Can't get GPU infos : 'nvidia-smi' tool is not available on this platform.</b>"
+
+	html += "</pre>"
+	
+	return html

lib/karafan/tfc_tdf.py

@@ -0,0 +1,235 @@
+
+#   MIT License - Copyright (c) 2023 - ZFTurbo & Jarredou
+
+import torch
+import torch.nn as nn
+from functools import partial
+
+#   https://github.com/jarredou/MVSEP-MDX23-Colab_v2/
+
+class Conv_TDF_net_trim_model(nn.Module):
+
+	def __init__(self, device, target_stem, neuron_blocks, model_params, hop=1024):
+
+		super(Conv_TDF_net_trim_model, self).__init__()
+		
+		self.dim_c = 4
+		self.dim_f = model_params['dim_F_set']
+		self.dim_t = model_params['dim_T_set']
+		self.n_fft = model_params['N_FFT_scale']
+		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.target_stem = target_stem
+
+		out_c = self.dim_c * 4 if target_stem == '*' else self.dim_c
+		self.freq_pad = torch.zeros([1, out_c, self.n_bins - self.dim_f, self.dim_t]).to(device)
+		
+		# Only used by "forward()" method for training !
+		# self.n = neuron_blocks // 2
+
+	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=None):
+		freq_pad = self.freq_pad.repeat([x.shape[0], 1, 1, 1]) if freq_pad is None else freq_pad
+		x = torch.cat([x, freq_pad], -2)
+		x = x.reshape([-1, 2, 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, 2, self.chunk_size])
+
+#   https://github.com/ZFTurbo/MVSEP-MDX23-music-separation-model
+
+class STFT:
+	def __init__(self, config):
+		self.n_fft = config.n_fft
+		self.hop_length = config.hop_length
+		self.window = torch.hann_window(window_length=self.n_fft, periodic=True)
+		self.dim_f = config.dim_f
+
+	def __call__(self, x):
+		window = self.window.to(x.device)
+		batch_dims = x.shape[:-2]
+		c, t = x.shape[-2:]
+		x = x.reshape([-1, t])
+		x = torch.stft(x, n_fft=self.n_fft, hop_length=self.hop_length, window=window, center=True, return_complex=True)
+		x = torch.view_as_real(x)
+		x = x.permute([0, 3, 1, 2])
+		x = x.reshape([*batch_dims, c, 2, -1, x.shape[-1]]).reshape([*batch_dims, c * 2, -1, x.shape[-1]])
+		return x[..., :self.dim_f, :]
+
+	def inverse(self, x):
+		window = self.window.to(x.device)
+		batch_dims = x.shape[:-3]
+		c, f, t = x.shape[-3:]
+		n = self.n_fft // 2 + 1
+		f_pad = torch.zeros([*batch_dims, c, n - f, t]).to(x.device)
+		x = torch.cat([x, f_pad], -2)
+		x = x.reshape([*batch_dims, c // 2, 2, n, t]).reshape([-1, 2, n, t])
+		x = x.permute([0, 2, 3, 1])
+		x = x[..., 0] + x[..., 1] * 1.j
+		x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop_length, window=window, center=True)
+		x = x.reshape([*batch_dims, 2, -1])
+		return x
+
+def get_norm(norm_type):
+	def norm(c, norm_type):
+		if norm_type == 'BatchNorm':
+			return nn.BatchNorm2d(c)
+		elif norm_type == 'InstanceNorm':
+			return nn.InstanceNorm2d(c, affine=True)
+		elif 'GroupNorm' in norm_type:
+			g = int(norm_type.replace('GroupNorm', ''))
+			return nn.GroupNorm(num_groups=g, num_channels=c)
+		else:
+			return nn.Identity()
+
+	return partial(norm, norm_type=norm_type)
+
+def get_act(act_type):
+	if act_type == 'gelu':
+		return nn.GELU()
+	elif act_type == 'relu':
+		return nn.ReLU()
+	elif act_type[:3] == 'elu':
+		alpha = float(act_type.replace('elu', ''))
+		return nn.ELU(alpha)
+	else:
+		raise Exception
+
+class Upscale(nn.Module):
+	def __init__(self, in_c, out_c, scale, norm, act):
+		super().__init__()
+		self.conv = nn.Sequential(norm(in_c), act, nn.ConvTranspose2d(in_channels=in_c, out_channels=out_c, kernel_size=scale, stride=scale, bias=False))
+
+	def forward(self, x):
+		return self.conv(x)
+
+class Downscale(nn.Module):
+	def __init__(self, in_c, out_c, scale, norm, act):
+		super().__init__()
+		self.conv = nn.Sequential(norm(in_c), act, nn.Conv2d(in_channels=in_c, out_channels=out_c, kernel_size=scale, stride=scale, bias=False))
+
+	def forward(self, x):
+		return self.conv(x)
+
+class TFC_TDF(nn.Module):
+	def __init__(self, in_c, c, l, f, bn, norm, act):
+		super().__init__()
+		self.blocks = nn.ModuleList()
+		
+		for i in range(l):
+			block = nn.Module()
+			block.tfc1 = nn.Sequential(norm(in_c), act, nn.Conv2d(in_c, c, 3, 1, 1, bias=False),)
+			block.tdf = nn.Sequential(norm(c), act, nn.Linear(f, f // bn, bias=False), norm(c), act, nn.Linear(f // bn, f, bias=False))
+			block.tfc2 = nn.Sequential(norm(c), act, nn.Conv2d(c, c, 3, 1, 1, bias=False))
+			block.shortcut = nn.Conv2d(in_c, c, 1, 1, 0, bias=False)
+			self.blocks.append(block)
+			in_c = c
+
+	def forward(self, x):
+		for block in self.blocks:
+			s = block.shortcut(x)
+			x = block.tfc1(x)
+			x = x + block.tdf(x)
+			x = block.tfc2(x)
+			x = x + s
+		return x
+
+class TFC_TDF_net(nn.Module):
+	
+	def __init__(self, config):
+		super().__init__()
+		
+		self.config = config
+		norm = get_norm(norm_type=config.model.norm)
+		act = get_act(act_type=config.model.act)
+		self.num_target_instruments = 1 if config.training.target_instrument else len(config.training.instruments)
+		self.num_subbands = config.model.num_subbands
+		dim_c = self.num_subbands * config.audio.num_channels * 2
+		n = config.model.num_scales
+		scale = config.model.scale
+		l = config.model.num_blocks_per_scale
+		c = config.model.num_channels
+		g = config.model.growth
+		bn = config.model.bottleneck_factor
+		f = config.audio.dim_f // self.num_subbands
+		self.first_conv = nn.Conv2d(dim_c, c, 1, 1, 0, bias=False)
+		self.encoder_blocks = nn.ModuleList()
+		
+		for i in range(n):
+			block = nn.Module()
+			block.tfc_tdf = TFC_TDF(c, c, l, f, bn, norm, act)
+			block.downscale = Downscale(c, c + g, scale, norm, act)
+			f = f // scale[1]
+			c += g
+			self.encoder_blocks.append(block)
+
+		self.bottleneck_block = TFC_TDF(c, c, l, f, bn, norm, act)
+		self.decoder_blocks = nn.ModuleList()
+		
+		for i in range(n):
+			block = nn.Module()
+			block.upscale = Upscale(c, c - g, scale, norm, act)
+			f = f * scale[1]
+			c -= g
+			block.tfc_tdf = TFC_TDF(2 * c, c, l, f, bn, norm, act)
+			self.decoder_blocks.append(block)
+
+		self.final_conv = nn.Sequential(nn.Conv2d(c + dim_c, c, 1, 1, 0, bias=False), act, nn.Conv2d(c, self.num_target_instruments * dim_c, 1, 1, 0, bias=False))
+		self.stft = STFT(config.audio)
+
+	def cac2cws(self, x):
+		k = self.num_subbands
+		b, c, f, t = x.shape
+		x = x.reshape(b, c, k, f // k, t)
+		x = x.reshape(b, c * k, f // k, t)
+		return x
+
+	def cws2cac(self, x):
+		k = self.num_subbands
+		b, c, f, t = x.shape
+		x = x.reshape(b, c // k, k, f, t)
+		x = x.reshape(b, c // k, f * k, t)
+		return x
+
+	def forward(self, x):
+		x = self.stft(x)
+		mix = x = self.cac2cws(x)
+		first_conv_out = x = self.first_conv(x)
+		x = x.transpose(-1, -2)
+		encoder_outputs = []
+		
+		for block in self.encoder_blocks:
+			x = block.tfc_tdf(x)
+			encoder_outputs.append(x)
+			x = block.downscale(x)
+			
+		x = self.bottleneck_block(x)
+		
+		for block in self.decoder_blocks:
+			x = block.upscale(x)
+			x = torch.cat([x, encoder_outputs.pop()], 1)
+			x = block.tfc_tdf(x)
+
+		x = x.transpose(-1, -2)
+		x = x * first_conv_out  
+		x = self.final_conv(torch.cat([mix, x], 1))
+		x = self.cws2cac(x)
+		
+		if self.num_target_instruments > 1:
+			b, c, f, t = x.shape
+			x = x.reshape(b, self.num_target_instruments, -1, f, t)
+			
+		x = self.stft.inverse(x)
+
+		return x

lib/mdx.py

@@ -0,0 +1,296 @@
+import json
+import os
+import warnings
+
+import numpy as np
+import torch
+from tqdm import tqdm
+from .model_utils import get_hash
+
+from .utils import gc_collect
+import onnxruntime as ort
+
+warnings.filterwarnings("ignore")
+stem_naming = {'Vocals': 'Instrumental', 'Other': 'Instruments', 'Instrumental': 'Vocals', 'Drums': 'Drumless', 'Bass': 'Bassless'}
+
+
+class MDXParams:
+    def __init__(self, device, dim_f, dim_t, n_fft, hop=1024, stem_name=None, compensation=1.000):
+        self.dim_f = dim_f
+        self.dim_t = dim_t
+        self.dim_c = 4
+        self.n_fft = n_fft
+        self.hop = hop
+        self.stem_name = stem_name
+        self.compensation = compensation
+
+        self.n_bins = self.n_fft // 2 + 1
+        self.trim = self.n_fft//2
+        self.chunk_size = hop * (self.dim_t - 1)
+        self.device = device
+        self.window = torch.hann_window(window_length=self.n_fft, periodic=True,device=self.device)
+
+        self.freq_pad = torch.zeros([1, self.dim_c, self.n_bins-self.dim_f, self.dim_t]).to(self.device)
+        self.gen_size = self.chunk_size-2*self.trim
+
+    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=None):
+        freq_pad = self.freq_pad.repeat([x.shape[0],1,1,1]) if freq_pad is None else freq_pad
+        x = torch.cat([x, freq_pad], -2)
+        x = x.reshape([-1,2,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,2,self.chunk_size])
+
+class MDXModel:
+    DEFAULT_SR = 44100
+    # Unit: seconds
+    DEFAULT_CHUNK_SIZE = 0 * DEFAULT_SR
+    DEFAULT_MARGIN_SIZE = 1 * DEFAULT_SR
+
+    DEFAULT_PROCESSOR = 0
+
+    # def __init__(self, model_path: str, params: MDXModel, processor=DEFAULT_PROCESSOR,margin=DEFAULT_MARGIN_SIZE,chunks=15):
+    def __init__(self, model_path: str, device="cpu",margin=DEFAULT_MARGIN_SIZE,chunks=15,denoise=False):
+
+        # Set the device and the provider (CPU or CUDA)
+        self.device = device #torch.device(f'cuda:{processor}') if processor >= 0 else torch.device('cpu')
+        # self.provider = ['CUDAExecutionProvider'] if processor >= 0 else ['CPUExecutionProvider']
+        self.providers = ['CUDAExecutionProvider','CPUExecutionProvider']
+
+        mp = self.get_params(model_path)
+        self.params = MDXParams(
+            self.device,
+            dim_f=mp["mdx_dim_f_set"],
+            dim_t=2 ** mp["mdx_dim_t_set"],
+            n_fft=mp["mdx_n_fft_scale_set"],
+            stem_name=mp["primary_stem"],
+            compensation=mp["compensate"]
+        )
+
+        # Load the ONNX model using ONNX Runtime
+        self.ort = ort.InferenceSession(model_path, providers=self.providers)
+        print(f"onnx load done: {model_path} ({self.params})")
+
+        self.run = lambda spec: self.ort.run(None, {'input': spec.cpu().numpy()})[0]
+        self.margin=margin
+        self.chunks=chunks
+        self.denoise=denoise
+        self.mdx_batch_size=1
+
+    def initialize_mix(self, mix, is_ckpt=False):
+        if is_ckpt:
+            pad = self.params.gen_size + self.params.trim - ((mix.shape[-1]) % self.params.gen_size)
+            mixture = np.concatenate((np.zeros((2, self.params.trim), dtype='float32'),mix, np.zeros((2, pad), dtype='float32')), 1)
+            num_chunks = mixture.shape[-1] // self.params.gen_size
+            mix_waves = [mixture[:, i * self.params.gen_size: i * self.params.gen_size + self.params.chunk_size] for i in range(num_chunks)]
+        else:
+            mix_waves = []
+            n_sample = mix.shape[1]
+            pad = self.params.gen_size - n_sample % self.params.gen_size
+            mix_p = np.concatenate((np.zeros((2,self.params.trim)), mix, np.zeros((2,pad)), np.zeros((2,self.params.trim))), 1)
+            i = 0
+            while i < n_sample + pad:
+                waves = np.array(mix_p[:, i:i+self.params.chunk_size])
+                mix_waves.append(waves)
+                i += self.params.gen_size
+                
+        mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(self.device)
+
+        return mix_waves, pad
+    
+    def demix_base(self, mix, is_ckpt=False, is_match_mix=False):
+        chunked_sources = []
+        for slice in tqdm(mix,"Processing audio:"):
+            sources = []
+            tar_waves_ = []
+            mix_p = mix[slice]
+            mix_waves, pad = self.initialize_mix(mix_p, is_ckpt=is_ckpt)
+            mix_waves = mix_waves.split(self.mdx_batch_size)
+            pad = mix_p.shape[-1] if is_ckpt else -pad
+            with torch.no_grad():
+                for mix_wave in mix_waves:
+                    tar_waves = self.run_model(mix_wave, is_ckpt=is_ckpt, is_match_mix=is_match_mix)
+                    tar_waves_.append(tar_waves)
+                tar_waves_ = np.vstack(tar_waves_)[:, :, self.params.trim:-self.params.trim] if is_ckpt else tar_waves_
+                tar_waves = np.concatenate(tar_waves_, axis=-1)[:, :pad]
+                start = 0 if slice == 0 else self.margin
+                end = None if slice == list(mix.keys())[::-1][0] or self.margin == 0 else -self.margin
+                sources.append(tar_waves[:,start:end]/self.params.compensation)
+            chunked_sources.append(sources)
+        sources = np.concatenate(chunked_sources, axis=-1)
+        
+        return sources
+
+    def run_model(self, mix, is_ckpt=False, is_match_mix=False):
+        
+        spek = self.params.stft(mix.to(self.device))*self.params.compensation
+        spek[:, :, :3, :] *= 0 
+
+        if is_match_mix:
+            spec_pred = spek.cpu().numpy()
+        else:
+            spec_pred = self.run(spek)*0.5-self.run(-spek)*0.5 if self.denoise else self.run(spek)
+
+        if is_ckpt:
+            return self.params.istft(spec_pred).cpu().detach().numpy()
+        else: 
+            return self.params.istft(torch.tensor(spec_pred).to(self.device)).cpu()[:,:,self.params.trim:-self.params.trim].transpose(0,1).reshape(2, -1).numpy()
+
+    def __del__(self):
+        del self.ort
+        gc_collect()
+
+    def get_params(self,model_path):
+        model_hash = get_hash(model_path)
+        with open(os.path.join(os.path.dirname(model_path), 'model_data.json')) as infile:
+            model_params = json.load(infile)
+        return model_params.get(model_hash)
+
+    # @staticmethod
+    # def segment(wave, combine=True, chunk_size=DEFAULT_CHUNK_SIZE, margin_size=DEFAULT_MARGIN_SIZE):
+    #     """
+    #     Segment or join segmented wave array
+
+    #     Args:
+    #         wave: (np.array) Wave array to be segmented or joined
+    #         combine: (bool) If True, combines segmented wave array. If False, segments wave array.
+    #         chunk_size: (int) Size of each segment (in samples)
+    #         margin_size: (int) Size of margin between segments (in samples)
+
+    #     Returns:
+    #         numpy array: Segmented or joined wave array
+    #     """
+
+    #     if combine:
+    #         processed_wave = None  # Initializing as None instead of [] for later numpy array concatenation
+    #         for segment_count, segment in enumerate(wave):
+    #             start = 0 if segment_count == 0 else margin_size
+    #             end = None if segment_count == len(wave) - 1 else -margin_size
+    #             if margin_size == 0:
+    #                 end = None
+    #             if processed_wave is None:  # Create array for first segment
+    #                 processed_wave = segment[:, start:end]
+    #             else:  # Concatenate to existing array for subsequent segments
+    #                 processed_wave = np.concatenate((processed_wave, segment[:, start:end]), axis=-1)
+
+    #     else:
+    #         processed_wave = []
+    #         sample_count = wave.shape[-1]
+
+    #         if chunk_size <= 0 or chunk_size > sample_count:
+    #             chunk_size = sample_count
+
+    #         if margin_size > chunk_size:
+    #             margin_size = chunk_size
+
+    #         for segment_count, skip in enumerate(range(0, sample_count, chunk_size)):
+
+    #             margin = 0 if segment_count == 0 else margin_size
+    #             end = min(skip + chunk_size + margin_size, sample_count)
+    #             start = skip - margin
+
+    #             cut = wave[:, start:end].copy()
+    #             processed_wave.append(cut)
+
+    #             if end == sample_count:
+    #                 break
+
+    #     return processed_wave
+
+    # def pad_wave(self, wave):
+    #     """
+    #     Pad the wave array to match the required chunk size
+
+    #     Args:
+    #         wave: (np.array) Wave array to be padded
+
+    #     Returns:
+    #         tuple: (padded_wave, pad, trim)
+    #             - padded_wave: Padded wave array
+    #             - pad: Number of samples that were padded
+    #             - trim: Number of samples that were trimmed
+    #     """
+    #     n_sample = wave.shape[1]
+    #     trim = self.params.n_fft // 2
+    #     gen_size = self.params.chunk_size - 2 * trim
+    #     pad = gen_size - n_sample % gen_size
+
+    #     # Padded wave
+    #     wave_p = np.concatenate((np.zeros((2, trim)), wave, np.zeros((2, pad)), np.zeros((2, trim))), 1)
+
+    #     mix_waves = []
+    #     for i in range(0, n_sample + pad, gen_size):
+    #         waves = np.array(wave_p[:, i:i + self.params.chunk_size])
+    #         mix_waves.append(waves)
+
+    #     mix_waves = torch.tensor(mix_waves, dtype=torch.float32).cpu()
+
+    #     return mix_waves, pad, trim
+
+    # # def _process_wave(self, mix_waves, trim, pad, q: queue.Queue, _id: int):
+    # def _process_wave(self, mix_waves, trim, pad):
+    #     """
+    #     Process each wave segment in a multi-threaded environment
+
+    #     Args:
+    #         mix_waves: (torch.Tensor) Wave segments to be processed
+    #         trim: (int) Number of samples trimmed during padding
+    #         pad: (int) Number of samples padded during padding
+    #         q: (queue.Queue) Queue to hold the processed wave segments
+    #         _id: (int) Identifier of the processed wave segment
+
+    #     Returns:
+    #         numpy array: Processed wave segment
+    #     """
+    #     mix_waves = mix_waves.split(1)
+    #     with torch.no_grad():
+    #         pw = []
+    #         for mix_wave in mix_waves:
+    #             # self.prog.update()
+    #             spec = self.stft(mix_wave.to(self.device))
+    #             processed_spec = torch.tensor(self.process(spec))
+    #             processed_wav = self.istft(processed_spec)
+    #             processed_wav = processed_wav[:, :, trim:-trim].transpose(0, 1).reshape(2, -1).cpu().numpy()
+    #             pw.append(processed_wav)
+
+    #     processed_signal = np.concatenate(pw, axis=-1)[:, :-pad]
+        
+    #     return processed_signal
+
+    # def process_wave(self, wave: np.array, mt_threads=1):
+    #     """
+    #     Process the wave array in a multi-threaded environment
+
+    #     Args:
+    #         wave: (np.array) Wave array to be processed
+    #         mt_threads: (int) Number of threads to be used for processing
+
+    #     Returns:
+    #         numpy array: Processed wave array
+    #     """
+    #     # self.prog = tqdm(total=0)
+    #     # chunk_size = wave.shape[-1] // (mt_threads if not self.chunks else (self.chunks*mt_threads))
+    #     chunk_size = self.params.chunk_size
+    #     margin = self.margin
+    #     waves = self.segment(wave, False, chunk_size, margin)
+        
+    #     processed_batches = []
+    #     for batch in tqdm(waves,desc="processing waves"):
+    #         mix_waves, pad, trim = self.pad_wave(batch)
+    #         processed_waves = self._process_wave(mix_waves, trim, pad)
+    #         processed_batches.append(processed_waves)
+
+    #     assert len(processed_batches) == len(waves), 'Incomplete processed batches, please reduce batch size!'
+    #     segmented_mix = self.segment(processed_batches, True, chunk_size, margin)
+
+    #     return segmented_mix

lib/model_utils.py

@@ -0,0 +1,53 @@
+
+import hashlib
+import os
+import torch.nn.functional as F
+import librosa
+import torch
+
+def get_hash(model_path):
+    try:
+        with open(model_path, 'rb') as f:
+            f.seek(- 10000 * 1024, 2)
+            model_hash = hashlib.md5(f.read()).hexdigest()
+    except:
+        model_hash = hashlib.md5(open(model_path, 'rb').read()).hexdigest()
+
+    return model_hash
+
+def load_hubert(model_path,config):
+    try:
+        from fairseq import checkpoint_utils
+        models, _, _ = checkpoint_utils.load_model_ensemble_and_task([model_path],suffix="",)
+        hubert_model = models[0]
+        hubert_model = hubert_model.to(config.device)
+        if config.is_half:
+            hubert_model = hubert_model.half()
+        else:
+            hubert_model = hubert_model.float()
+        hubert_model.eval()
+        return hubert_model
+    except Exception as e:
+        print(e)
+        return None
+    
+def change_rms(data1, sr1, data2, sr2, rate):  # 1是输入音频，2是输出音频,rate是2的占比
+    # print(data1.max(),data2.max())
+    rms1 = librosa.feature.rms(
+        y=data1, frame_length=sr1 // 2 * 2, hop_length=sr1 // 2
+    )  # 每半秒一个点
+    rms2 = librosa.feature.rms(y=data2, frame_length=sr2 // 2 * 2, hop_length=sr2 // 2)
+    rms1 = torch.from_numpy(rms1)
+    rms1 = F.interpolate(
+        rms1.unsqueeze(0), size=data2.shape[0], mode="linear"
+    ).squeeze()
+    rms2 = torch.from_numpy(rms2)
+    rms2 = F.interpolate(
+        rms2.unsqueeze(0), size=data2.shape[0], mode="linear"
+    ).squeeze()
+    rms2 = torch.max(rms2, torch.zeros_like(rms2) + 1e-6)
+    data2 *= (
+        torch.pow(rms1, torch.tensor(1 - rate))
+        * torch.pow(rms2, torch.tensor(rate - 1))
+    ).numpy()
+    return data2

lib/rmvpe.py

@@ -0,0 +1,684 @@
+import torch, numpy as np, pdb
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch, pdb
+import numpy as np
+import torch.nn.functional as F
+from scipy.signal import get_window
+from librosa.util import pad_center, tiny, normalize
+
+
+###stft codes from https://github.com/pseeth/torch-stft/blob/master/torch_stft/util.py
+def window_sumsquare(
+    window,
+    n_frames,
+    hop_length=200,
+    win_length=800,
+    n_fft=800,
+    dtype=np.float32,
+    norm=None,
+):
+    """
+    # from librosa 0.6
+    Compute the sum-square envelope of a window function at a given hop length.
+    This is used to estimate modulation effects induced by windowing
+    observations in short-time fourier transforms.
+    Parameters
+    ----------
+    window : string, tuple, number, callable, or list-like
+        Window specification, as in `get_window`
+    n_frames : int > 0
+        The number of analysis frames
+    hop_length : int > 0
+        The number of samples to advance between frames
+    win_length : [optional]
+        The length of the window function.  By default, this matches `n_fft`.
+    n_fft : int > 0
+        The length of each analysis frame.
+    dtype : np.dtype
+        The data type of the output
+    Returns
+    -------
+    wss : np.ndarray, shape=`(n_fft + hop_length * (n_frames - 1))`
+        The sum-squared envelope of the window function
+    """
+    if win_length is None:
+        win_length = n_fft
+
+    n = n_fft + hop_length * (n_frames - 1)
+    x = np.zeros(n, dtype=dtype)
+
+    # Compute the squared window at the desired length
+    win_sq = get_window(window, win_length, fftbins=True)
+    win_sq = normalize(win_sq, norm=norm) ** 2
+    win_sq = pad_center(win_sq, n_fft)
+
+    # Fill the envelope
+    for i in range(n_frames):
+        sample = i * hop_length
+        x[sample : min(n, sample + n_fft)] += win_sq[: max(0, min(n_fft, n - sample))]
+    return x
+
+
+class STFT(torch.nn.Module):
+    def __init__(
+        self, filter_length=1024, hop_length=512, win_length=None, window="hann"
+    ):
+        """
+        This module implements an STFT using 1D convolution and 1D transpose convolutions.
+        This is a bit tricky so there are some cases that probably won't work as working
+        out the same sizes before and after in all overlap add setups is tough. Right now,
+        this code should work with hop lengths that are half the filter length (50% overlap
+        between frames).
+        Keyword Arguments:
+            filter_length {int} -- Length of filters used (default: {1024})
+            hop_length {int} -- Hop length of STFT (restrict to 50% overlap between frames) (default: {512})
+            win_length {[type]} -- Length of the window function applied to each frame (if not specified, it
+                equals the filter length). (default: {None})
+            window {str} -- Type of window to use (options are bartlett, hann, hamming, blackman, blackmanharris)
+                (default: {'hann'})
+        """
+        super(STFT, self).__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length if win_length else filter_length
+        self.window = window
+        self.forward_transform = None
+        self.pad_amount = int(self.filter_length / 2)
+        scale = self.filter_length / self.hop_length
+        fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+        cutoff = int((self.filter_length / 2 + 1))
+        fourier_basis = np.vstack(
+            [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
+        )
+        forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
+        inverse_basis = torch.FloatTensor(
+            np.linalg.pinv(scale * fourier_basis).T[:, None, :]
+        )
+
+        assert filter_length >= self.win_length
+        # get window and zero center pad it to filter_length
+        fft_window = get_window(window, self.win_length, fftbins=True)
+        fft_window = pad_center(fft_window, size=filter_length)
+        fft_window = torch.from_numpy(fft_window).float()
+
+        # window the bases
+        forward_basis *= fft_window
+        inverse_basis *= fft_window
+
+        self.register_buffer("forward_basis", forward_basis.float())
+        self.register_buffer("inverse_basis", inverse_basis.float())
+
+    def transform(self, input_data):
+        """Take input data (audio) to STFT domain.
+        Arguments:
+            input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+        Returns:
+            magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+            phase {tensor} -- Phase of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+        """
+        num_batches = input_data.shape[0]
+        num_samples = input_data.shape[-1]
+
+        self.num_samples = num_samples
+
+        # similar to librosa, reflect-pad the input
+        input_data = input_data.view(num_batches, 1, num_samples)
+        # print(1234,input_data.shape)
+        input_data = F.pad(
+            input_data.unsqueeze(1),
+            (self.pad_amount, self.pad_amount, 0, 0, 0, 0),
+            mode="reflect",
+        ).squeeze(1)
+        # print(2333,input_data.shape,self.forward_basis.shape,self.hop_length)
+        # pdb.set_trace()
+        forward_transform = F.conv1d(
+            input_data, self.forward_basis, stride=self.hop_length, padding=0
+        )
+
+        cutoff = int((self.filter_length / 2) + 1)
+        real_part = forward_transform[:, :cutoff, :]
+        imag_part = forward_transform[:, cutoff:, :]
+
+        magnitude = torch.sqrt(real_part**2 + imag_part**2)
+        # phase = torch.atan2(imag_part.data, real_part.data)
+
+        return magnitude  # , phase
+
+    def inverse(self, magnitude, phase):
+        """Call the inverse STFT (iSTFT), given magnitude and phase tensors produced
+        by the ```transform``` function.
+        Arguments:
+            magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+            phase {tensor} -- Phase of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+        Returns:
+            inverse_transform {tensor} -- Reconstructed audio given magnitude and phase. Of
+                shape (num_batch, num_samples)
+        """
+        recombine_magnitude_phase = torch.cat(
+            [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
+        )
+
+        inverse_transform = F.conv_transpose1d(
+            recombine_magnitude_phase,
+            self.inverse_basis,
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        if self.window is not None:
+            window_sum = window_sumsquare(
+                self.window,
+                magnitude.size(-1),
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                n_fft=self.filter_length,
+                dtype=np.float32,
+            )
+            # remove modulation effects
+            approx_nonzero_indices = torch.from_numpy(
+                np.where(window_sum > tiny(window_sum))[0]
+            )
+            window_sum = torch.from_numpy(window_sum).to(inverse_transform.device)
+            inverse_transform[:, :, approx_nonzero_indices] /= window_sum[
+                approx_nonzero_indices
+            ]
+
+            # scale by hop ratio
+            inverse_transform *= float(self.filter_length) / self.hop_length
+
+        inverse_transform = inverse_transform[..., self.pad_amount :]
+        inverse_transform = inverse_transform[..., : self.num_samples]
+        inverse_transform = inverse_transform.squeeze(1)
+
+        return inverse_transform
+
+    def forward(self, input_data):
+        """Take input data (audio) to STFT domain and then back to audio.
+        Arguments:
+            input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+        Returns:
+            reconstruction {tensor} -- Reconstructed audio given magnitude and phase. Of
+                shape (num_batch, num_samples)
+        """
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction
+
+
+from time import time as ttime
+
+
+class BiGRU(nn.Module):
+    def __init__(self, input_features, hidden_features, num_layers):
+        super(BiGRU, self).__init__()
+        self.gru = nn.GRU(
+            input_features,
+            hidden_features,
+            num_layers=num_layers,
+            batch_first=True,
+            bidirectional=True,
+        )
+
+    def forward(self, x):
+        return self.gru(x)[0]
+
+
+class ConvBlockRes(nn.Module):
+    def __init__(self, in_channels, out_channels, momentum=0.01):
+        super(ConvBlockRes, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=(1, 1),
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+            nn.Conv2d(
+                in_channels=out_channels,
+                out_channels=out_channels,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=(1, 1),
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+        )
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
+            self.is_shortcut = True
+        else:
+            self.is_shortcut = False
+
+    def forward(self, x):
+        if self.is_shortcut:
+            return self.conv(x) + self.shortcut(x)
+        else:
+            return self.conv(x) + x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        in_size,
+        n_encoders,
+        kernel_size,
+        n_blocks,
+        out_channels=16,
+        momentum=0.01,
+    ):
+        super(Encoder, self).__init__()
+        self.n_encoders = n_encoders
+        self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
+        self.layers = nn.ModuleList()
+        self.latent_channels = []
+        for i in range(self.n_encoders):
+            self.layers.append(
+                ResEncoderBlock(
+                    in_channels, out_channels, kernel_size, n_blocks, momentum=momentum
+                )
+            )
+            self.latent_channels.append([out_channels, in_size])
+            in_channels = out_channels
+            out_channels *= 2
+            in_size //= 2
+        self.out_size = in_size
+        self.out_channel = out_channels
+
+    def forward(self, x):
+        concat_tensors = []
+        x = self.bn(x)
+        for i in range(self.n_encoders):
+            _, x = self.layers[i](x)
+            concat_tensors.append(_)
+        return x, concat_tensors
+
+
+class ResEncoderBlock(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01
+    ):
+        super(ResEncoderBlock, self).__init__()
+        self.n_blocks = n_blocks
+        self.conv = nn.ModuleList()
+        self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
+        for i in range(n_blocks - 1):
+            self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
+        self.kernel_size = kernel_size
+        if self.kernel_size is not None:
+            self.pool = nn.AvgPool2d(kernel_size=kernel_size)
+
+    def forward(self, x):
+        for i in range(self.n_blocks):
+            x = self.conv[i](x)
+        if self.kernel_size is not None:
+            return x, self.pool(x)
+        else:
+            return x
+
+
+class Intermediate(nn.Module):  #
+    def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
+        super(Intermediate, self).__init__()
+        self.n_inters = n_inters
+        self.layers = nn.ModuleList()
+        self.layers.append(
+            ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)
+        )
+        for i in range(self.n_inters - 1):
+            self.layers.append(
+                ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)
+            )
+
+    def forward(self, x):
+        for i in range(self.n_inters):
+            x = self.layers[i](x)
+        return x
+
+
+class ResDecoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
+        super(ResDecoderBlock, self).__init__()
+        out_padding = (0, 1) if stride == (1, 2) else (1, 1)
+        self.n_blocks = n_blocks
+        self.conv1 = nn.Sequential(
+            nn.ConvTranspose2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(3, 3),
+                stride=stride,
+                padding=(1, 1),
+                output_padding=out_padding,
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+        )
+        self.conv2 = nn.ModuleList()
+        self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
+        for i in range(n_blocks - 1):
+            self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))
+
+    def forward(self, x, concat_tensor):
+        x = self.conv1(x)
+        x = torch.cat((x, concat_tensor), dim=1)
+        for i in range(self.n_blocks):
+            x = self.conv2[i](x)
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList()
+        self.n_decoders = n_decoders
+        for i in range(self.n_decoders):
+            out_channels = in_channels // 2
+            self.layers.append(
+                ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum)
+            )
+            in_channels = out_channels
+
+    def forward(self, x, concat_tensors):
+        for i in range(self.n_decoders):
+            x = self.layers[i](x, concat_tensors[-1 - i])
+        return x
+
+
+class DeepUnet(nn.Module):
+    def __init__(
+        self,
+        kernel_size,
+        n_blocks,
+        en_de_layers=5,
+        inter_layers=4,
+        in_channels=1,
+        en_out_channels=16,
+    ):
+        super(DeepUnet, self).__init__()
+        self.encoder = Encoder(
+            in_channels, 128, en_de_layers, kernel_size, n_blocks, en_out_channels
+        )
+        self.intermediate = Intermediate(
+            self.encoder.out_channel // 2,
+            self.encoder.out_channel,
+            inter_layers,
+            n_blocks,
+        )
+        self.decoder = Decoder(
+            self.encoder.out_channel, en_de_layers, kernel_size, n_blocks
+        )
+
+    def forward(self, x):
+        x, concat_tensors = self.encoder(x)
+        x = self.intermediate(x)
+        x = self.decoder(x, concat_tensors)
+        return x
+
+
+class E2E(nn.Module):
+    def __init__(
+        self,
+        n_blocks,
+        n_gru,
+        kernel_size,
+        en_de_layers=5,
+        inter_layers=4,
+        in_channels=1,
+        en_out_channels=16,
+    ):
+        super(E2E, self).__init__()
+        self.unet = DeepUnet(
+            kernel_size,
+            n_blocks,
+            en_de_layers,
+            inter_layers,
+            in_channels,
+            en_out_channels,
+        )
+        self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
+        if n_gru:
+            self.fc = nn.Sequential(
+                BiGRU(3 * 128, 256, n_gru),
+                nn.Linear(512, 360),
+                nn.Dropout(0.25),
+                nn.Sigmoid(),
+            )
+        else:
+            self.fc = nn.Sequential(
+                nn.Linear(3 * nn.N_MELS, nn.N_CLASS), nn.Dropout(0.25), nn.Sigmoid()
+            )
+
+    def forward(self, mel):
+        # print(mel.shape)
+        mel = mel.transpose(-1, -2).unsqueeze(1)
+        x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
+        x = self.fc(x)
+        # print(x.shape)
+        return x
+
+
+from librosa.filters import mel
+
+
+class MelSpectrogram(torch.nn.Module):
+    def __init__(
+        self,
+        is_half,
+        n_mel_channels,
+        sampling_rate,
+        win_length,
+        hop_length,
+        n_fft=None,
+        mel_fmin=0,
+        mel_fmax=None,
+        clamp=1e-5,
+    ):
+        super().__init__()
+        n_fft = win_length if n_fft is None else n_fft
+        self.hann_window = {}
+        mel_basis = mel(
+            sr=sampling_rate,
+            n_fft=n_fft,
+            n_mels=n_mel_channels,
+            fmin=mel_fmin,
+            fmax=mel_fmax,
+            htk=True,
+        )
+        mel_basis = torch.from_numpy(mel_basis).float()
+        self.register_buffer("mel_basis", mel_basis)
+        self.n_fft = win_length if n_fft is None else n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.sampling_rate = sampling_rate
+        self.n_mel_channels = n_mel_channels
+        self.clamp = clamp
+        self.is_half = is_half
+
+    def forward(self, audio, keyshift=0, speed=1, center=True):
+        factor = 2 ** (keyshift / 12)
+        n_fft_new = int(np.round(self.n_fft * factor))
+        win_length_new = int(np.round(self.win_length * factor))
+        hop_length_new = int(np.round(self.hop_length * speed))
+        keyshift_key = str(keyshift) + "_" + str(audio.device)
+        if keyshift_key not in self.hann_window:
+            self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(
+                # "cpu"if(audio.device.type=="privateuseone") else audio.device
+                audio.device
+            )
+        # fft = torch.stft(#doesn't support pytorch_dml
+        #     # audio.cpu() if(audio.device.type=="privateuseone")else audio,
+        #     audio,
+        #     n_fft=n_fft_new,
+        #     hop_length=hop_length_new,
+        #     win_length=win_length_new,
+        #     window=self.hann_window[keyshift_key],
+        #     center=center,
+        #     return_complex=True,
+        # )
+        # magnitude = torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
+        # print(1111111111)
+        # print(222222222222222,audio.device,self.is_half)
+        if hasattr(self, "stft") == False:
+            # print(n_fft_new,hop_length_new,win_length_new,audio.shape)
+            self.stft = STFT(
+                filter_length=n_fft_new,
+                hop_length=hop_length_new,
+                win_length=win_length_new,
+                window="hann",
+            ).to(audio.device)
+        magnitude = self.stft.transform(audio)  # phase
+        # if (audio.device.type == "privateuseone"):
+        #     magnitude=magnitude.to(audio.device)
+        if keyshift != 0:
+            size = self.n_fft // 2 + 1
+            resize = magnitude.size(1)
+            if resize < size:
+                magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
+            magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
+        mel_output = torch.matmul(self.mel_basis, magnitude)
+        if self.is_half == True:
+            mel_output = mel_output.half()
+        log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
+        # print(log_mel_spec.device.type)
+        return log_mel_spec
+
+
+class RMVPE:
+    def __init__(self, model_path, is_half, onnx=False, device=None):
+        self.resample_kernel = {}
+        self.resample_kernel = {}
+        self.is_half = is_half
+        self.onnx = onnx
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
+        self.mel_extractor = MelSpectrogram(
+            is_half, 128, 16000, 1024, 160, None, 30, 8000
+        ).to(device)
+        if onnx:
+            import onnxruntime as ort
+
+            ort_session = ort.InferenceSession(
+                "rmvpe.onnx", providers=["DmlExecutionProvider"]
+            )
+            self.model = ort_session
+        else:
+            model = E2E(4, 1, (2, 2))
+            ckpt = torch.load(model_path, map_location="cpu")
+            model.load_state_dict(ckpt)
+            model.eval()
+            if is_half == True:
+                model = model.half()
+            self.model = model
+            self.model = self.model.to(device)
+        cents_mapping = 20 * np.arange(360) + 1997.3794084376191
+        self.cents_mapping = np.pad(cents_mapping, (4, 4))  # 368
+
+    def mel2hidden(self, mel):
+        with torch.no_grad():
+            n_frames = mel.shape[-1]
+            padding = min(32 * ((n_frames - 1) // 32 + 1) - n_frames, n_frames)
+            mel = F.pad(mel, (0, padding), mode="reflect")
+            #if "privateuseone" in str(self.device):
+            if self.onnx:
+                onnx_input_name = self.model.get_inputs()[0].name
+                onnx_outputs_names = self.model.get_outputs()[0].name
+                hidden = self.model.run(
+                    [onnx_outputs_names],
+                    input_feed={onnx_input_name: mel.cpu().numpy()},
+                )[0]
+            else:
+                hidden = self.model(mel)
+            return hidden[:, :n_frames]
+
+    def decode(self, hidden, thred=0.03):
+        cents_pred = self.to_local_average_cents(hidden, thred=thred)
+        f0 = 10 * (2 ** (cents_pred / 1200))
+        f0[f0 == 10] = 0
+        # f0 = np.array([10 * (2 ** (cent_pred / 1200)) if cent_pred else 0 for cent_pred in cents_pred])
+        return f0
+
+    def infer_from_audio(self, audio, thred=0.03):
+        if self.onnx == False:
+            audio = torch.from_numpy(audio).float().to(self.device).unsqueeze(0)
+            mel = self.mel_extractor(audio, center=True)
+            hidden = self.mel2hidden(mel)
+            hidden = hidden.squeeze(0).cpu().numpy()
+            if self.is_half == True:
+                hidden = hidden.astype("float32")
+            f0 = self.decode(hidden, thred=thred)
+            return f0
+        else:
+            # torch.cuda.synchronize()
+            t0 = ttime()
+            mel = self.mel_extractor(
+                torch.from_numpy(audio).float().to(self.device).unsqueeze(0), center=True
+            )
+            # print(123123123,mel.device.type)
+            # torch.cuda.synchronize()
+            t1 = ttime()
+            hidden = self.mel2hidden(mel)
+            # torch.cuda.synchronize()
+            t2 = ttime()
+            # print(234234,hidden.device.type)
+            if not self.onnx:
+                hidden = hidden.squeeze(0).cpu().numpy()
+            else:
+                hidden = hidden[0]
+            if self.is_half == True:
+                hidden = hidden.astype("float32")
+            f0 = self.decode(hidden, thred=thred)
+            # torch.cuda.synchronize()
+            t3 = ttime()
+            # print("hmvpe:%s\t%s\t%s\t%s"%(t1-t0,t2-t1,t3-t2,t3-t0))
+            return f0
+    
+    def infer_from_audio_with_pitch(self, audio, thred=0.03, f0_min=50, f0_max=1100):
+        audio = torch.from_numpy(audio).float().to(self.device).unsqueeze(0)
+        mel = self.mel_extractor(audio, center=True)
+        hidden = self.mel2hidden(mel)
+        hidden = hidden.squeeze(0).cpu().numpy()
+        if self.is_half == True:
+            hidden = hidden.astype("float32")
+        f0 = self.decode(hidden, thred=thred)
+        f0[(f0 < f0_min) | (f0 > f0_max)] = 0  
+        return f0
+
+    def to_local_average_cents(self, salience, thred=0.05):
+        # t0 = ttime()
+        center = np.argmax(salience, axis=1)  # 帧长#index
+        salience = np.pad(salience, ((0, 0), (4, 4)))  # 帧长,368
+        # t1 = ttime()
+        center += 4
+        todo_salience = []
+        todo_cents_mapping = []
+        starts = center - 4
+        ends = center + 5
+        for idx in range(salience.shape[0]):
+            todo_salience.append(salience[:, starts[idx] : ends[idx]][idx])
+            todo_cents_mapping.append(self.cents_mapping[starts[idx] : ends[idx]])
+        # t2 = ttime()
+        todo_salience = np.array(todo_salience)  # 帧长，9
+        todo_cents_mapping = np.array(todo_cents_mapping)  # 帧长，9
+        product_sum = np.sum(todo_salience * todo_cents_mapping, 1)
+        weight_sum = np.sum(todo_salience, 1)  # 帧长
+        devided = product_sum / weight_sum  # 帧长
+        # t3 = ttime()
+        maxx = np.max(salience, axis=1)  # 帧长
+        devided[maxx <= thred] = 0
+        # t4 = ttime()
+        # print("decode:%s\t%s\t%s\t%s" % (t1 - t0, t2 - t1, t3 - t2, t4 - t3))
+        return devided

lib/separate.py

@@ -0,0 +1,942 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from .uvr5_pack.demucs.apply import apply_model, demucs_segments
+from .uvr5_pack.demucs.hdemucs import HDemucs
+from .uvr5_pack.demucs.model_v2 import auto_load_demucs_model_v2
+from .uvr5_pack.demucs.pretrained import get_model as _gm
+from .uvr5_pack.demucs.utils import apply_model_v1
+from .uvr5_pack.demucs.utils import apply_model_v2
+from .uvr5_pack import spec_utils
+from .uvr5_pack.vr_network import nets
+from .uvr5_pack.vr_network import nets_new
+#from lib_v5.vr_network.model_param_init import ModelParameters
+from pathlib import Path
+from .uvr5_pack.constants import *
+from .uvr5_pack.error_handling import *
+import audioread
+import gzip
+import librosa
+import math
+import numpy as np
+import onnxruntime as ort
+import os
+import torch
+import warnings
+import pydub
+import soundfile as sf
+import traceback
+from .uvr5_pack import mdxnet as MdxnetSet
+
+if TYPE_CHECKING:
+    from .uvr5_pack.UVR import ModelData
+
+warnings.filterwarnings("ignore")
+cpu = torch.device('cpu')
+
+class SeperateAttributes:
+    def __init__(self, model_data: ModelData, process_data: dict, main_model_primary_stem_4_stem=None, main_process_method=None):
+        
+        self.list_all_models: list
+        self.process_data = process_data
+        self.progress_value = 0
+        self.set_progress_bar = process_data['set_progress_bar']
+        self.write_to_console = process_data['write_to_console']
+        self.audio_file = process_data['audio_file']
+        self.audio_file_base = process_data['audio_file_base']
+        self.export_path = process_data['export_path']
+        self.cached_source_callback = process_data['cached_source_callback']
+        self.cached_model_source_holder = process_data['cached_model_source_holder']
+        self.is_4_stem_ensemble = process_data['is_4_stem_ensemble']
+        self.list_all_models = process_data['list_all_models']
+        self.process_iteration = process_data['process_iteration']
+        self.mixer_path = model_data.mixer_path
+        self.model_samplerate = model_data.model_samplerate
+        self.model_capacity = model_data.model_capacity
+        self.is_vr_51_model = model_data.is_vr_51_model
+        self.is_pre_proc_model = model_data.is_pre_proc_model
+        self.is_secondary_model_activated = model_data.is_secondary_model_activated if not self.is_pre_proc_model else False
+        self.is_secondary_model = model_data.is_secondary_model if not self.is_pre_proc_model else True
+        self.process_method = model_data.process_method
+        self.model_path = model_data.model_path
+        self.model_name = model_data.model_name
+        self.model_basename = model_data.model_basename
+        self.wav_type_set = model_data.wav_type_set
+        self.mp3_bit_set = model_data.mp3_bit_set
+        self.save_format = model_data.save_format
+        self.is_gpu_conversion = model_data.is_gpu_conversion
+        self.is_normalization = model_data.is_normalization
+        self.is_primary_stem_only = model_data.is_primary_stem_only if not self.is_secondary_model else model_data.is_primary_model_primary_stem_only
+        self.is_secondary_stem_only = model_data.is_secondary_stem_only if not self.is_secondary_model else model_data.is_primary_model_secondary_stem_only      
+        self.is_ensemble_mode = model_data.is_ensemble_mode
+        self.secondary_model = model_data.secondary_model #
+        self.primary_model_primary_stem = model_data.primary_model_primary_stem
+        self.primary_stem = model_data.primary_stem #
+        self.secondary_stem = model_data.secondary_stem #
+        self.is_invert_spec = model_data.is_invert_spec #
+        self.is_mixer_mode = model_data.is_mixer_mode #
+        self.secondary_model_scale = model_data.secondary_model_scale #
+        self.is_demucs_pre_proc_model_inst_mix = model_data.is_demucs_pre_proc_model_inst_mix #
+        self.primary_source_map = {}
+        self.secondary_source_map = {}
+        self.primary_source = None
+        self.secondary_source = None
+        self.secondary_source_primary = None
+        self.secondary_source_secondary = None
+
+        if not model_data.process_method == DEMUCS_ARCH_TYPE:
+            if process_data['is_ensemble_master'] and not self.is_4_stem_ensemble:
+                if not model_data.ensemble_primary_stem == self.primary_stem:
+                    self.is_primary_stem_only, self.is_secondary_stem_only = self.is_secondary_stem_only, self.is_primary_stem_only
+            
+            if self.is_secondary_model and not process_data['is_ensemble_master']:
+                if not self.primary_model_primary_stem == self.primary_stem and not main_model_primary_stem_4_stem:
+                    self.is_primary_stem_only, self.is_secondary_stem_only = self.is_secondary_stem_only, self.is_primary_stem_only
+                    
+            if main_model_primary_stem_4_stem:
+                self.is_primary_stem_only = True if main_model_primary_stem_4_stem == self.primary_stem else False
+                self.is_secondary_stem_only = True if not main_model_primary_stem_4_stem == self.primary_stem else False
+
+            if self.is_pre_proc_model:
+                self.is_primary_stem_only = True if self.primary_stem == INST_STEM else False
+                self.is_secondary_stem_only = True if self.secondary_stem == INST_STEM else False
+
+        if model_data.process_method == MDX_ARCH_TYPE:
+            self.is_mdx_ckpt = model_data.is_mdx_ckpt
+            self.primary_model_name, self.primary_sources = self.cached_source_callback(MDX_ARCH_TYPE, model_name=self.model_basename)
+            self.is_denoise = model_data.is_denoise
+            self.mdx_batch_size = model_data.mdx_batch_size
+            self.compensate = model_data.compensate
+            self.dim_f, self.dim_t = model_data.mdx_dim_f_set, 2**model_data.mdx_dim_t_set
+            self.n_fft = model_data.mdx_n_fft_scale_set
+            self.chunks = model_data.chunks
+            self.margin = model_data.margin
+            self.adjust = 1
+            self.dim_c = 4
+            self.hop = 1024
+            
+            if self.is_gpu_conversion >= 0 and torch.cuda.is_available():
+                self.device, self.run_type = torch.device('cuda:0'), ['CUDAExecutionProvider']
+            else:
+                self.device, self.run_type = torch.device('cpu'), ['CPUExecutionProvider']
+
+        if model_data.process_method == DEMUCS_ARCH_TYPE:
+            self.demucs_stems = model_data.demucs_stems if not main_process_method in [MDX_ARCH_TYPE, VR_ARCH_TYPE] else None
+            self.secondary_model_4_stem = model_data.secondary_model_4_stem
+            self.secondary_model_4_stem_scale = model_data.secondary_model_4_stem_scale
+            self.primary_stem = model_data.ensemble_primary_stem if process_data['is_ensemble_master'] else model_data.primary_stem
+            self.secondary_stem = model_data.ensemble_secondary_stem if process_data['is_ensemble_master'] else model_data.secondary_stem
+            self.is_chunk_demucs = model_data.is_chunk_demucs
+            self.segment = model_data.segment
+            self.demucs_version = model_data.demucs_version
+            self.demucs_source_list = model_data.demucs_source_list
+            self.demucs_source_map = model_data.demucs_source_map
+            self.is_demucs_combine_stems = model_data.is_demucs_combine_stems
+            self.demucs_stem_count = model_data.demucs_stem_count
+            self.pre_proc_model = model_data.pre_proc_model
+            
+            if self.is_secondary_model and not process_data['is_ensemble_master']:
+                if not self.demucs_stem_count == 2 and model_data.primary_model_primary_stem == INST_STEM:
+                    self.primary_stem = VOCAL_STEM
+                    self.secondary_stem = INST_STEM
+                else:
+                    self.primary_stem = model_data.primary_model_primary_stem
+                    self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
+            
+            if self.is_chunk_demucs:
+                self.chunks_demucs = model_data.chunks_demucs
+                self.margin_demucs = model_data.margin_demucs
+            else:
+                self.chunks_demucs = 0
+                self.margin_demucs = 44100
+                
+            self.shifts = model_data.shifts
+            self.is_split_mode = model_data.is_split_mode if not self.demucs_version == DEMUCS_V4 else True
+            self.overlap = model_data.overlap
+            self.primary_model_name, self.primary_sources = self.cached_source_callback(DEMUCS_ARCH_TYPE, model_name=self.model_basename)
+
+        if model_data.process_method == VR_ARCH_TYPE:
+            self.primary_model_name, self.primary_sources = self.cached_source_callback(VR_ARCH_TYPE, model_name=self.model_basename)
+            self.mp = model_data.vr_model_param
+            self.high_end_process = model_data.is_high_end_process
+            self.is_tta = model_data.is_tta
+            self.is_post_process = model_data.is_post_process
+            self.is_gpu_conversion = model_data.is_gpu_conversion
+            self.batch_size = model_data.batch_size
+            self.window_size = model_data.window_size
+            self.input_high_end_h = None
+            self.post_process_threshold = model_data.post_process_threshold
+            self.aggressiveness = {'value': model_data.aggression_setting, 
+                                   'split_bin': self.mp.param['band'][1]['crop_stop'], 
+                                   'aggr_correction': self.mp.param.get('aggr_correction')}
+            
+    def start_inference_console_write(self):
+        
+        if self.is_secondary_model and not self.is_pre_proc_model:
+            self.write_to_console(INFERENCE_STEP_2_SEC(self.process_method, self.model_basename))
+        
+        if self.is_pre_proc_model:
+            self.write_to_console(INFERENCE_STEP_2_PRE(self.process_method, self.model_basename))
+        
+    def running_inference_console_write(self, is_no_write=False):
+        
+        self.write_to_console(DONE, base_text='') if not is_no_write else None
+        self.set_progress_bar(0.05) if not is_no_write else None
+        
+        if self.is_secondary_model and not self.is_pre_proc_model:
+            self.write_to_console(INFERENCE_STEP_1_SEC)
+        elif self.is_pre_proc_model:
+            self.write_to_console(INFERENCE_STEP_1_PRE)
+        else:
+            self.write_to_console(INFERENCE_STEP_1)
+        
+    def running_inference_progress_bar(self, length, is_match_mix=False):
+        if not is_match_mix:
+            self.progress_value += 1
+
+            if (0.8/length*self.progress_value) >= 0.8:
+                length = self.progress_value + 1
+  
+            self.set_progress_bar(0.1, (0.8/length*self.progress_value))
+        
+    def load_cached_sources(self, is_4_stem_demucs=False):
+        
+        if self.is_secondary_model and not self.is_pre_proc_model:
+            self.write_to_console(INFERENCE_STEP_2_SEC_CACHED_MODOEL(self.process_method, self.model_basename))
+        elif self.is_pre_proc_model:
+            self.write_to_console(INFERENCE_STEP_2_PRE_CACHED_MODOEL(self.process_method, self.model_basename))
+        else:
+            self.write_to_console(INFERENCE_STEP_2_PRIMARY_CACHED)
+
+        if not is_4_stem_demucs:
+            primary_stem, secondary_stem = gather_sources(self.primary_stem, self.secondary_stem, self.primary_sources)
+            
+            return primary_stem, secondary_stem
+            
+    def cache_source(self, secondary_sources):
+        
+        model_occurrences = self.list_all_models.count(self.model_basename)
+        
+        if not model_occurrences <= 1:
+            if self.process_method == MDX_ARCH_TYPE:
+                self.cached_model_source_holder(MDX_ARCH_TYPE, secondary_sources, self.model_basename)
+                
+            if self.process_method == VR_ARCH_TYPE:
+                self.cached_model_source_holder(VR_ARCH_TYPE, secondary_sources, self.model_basename)
+
+            if self.process_method == DEMUCS_ARCH_TYPE:
+                self.cached_model_source_holder(DEMUCS_ARCH_TYPE, secondary_sources, self.model_basename)
+                
+    def write_audio(self, stem_path, stem_source, samplerate, secondary_model_source=None, model_scale=None):
+                
+        if not self.is_secondary_model:
+            if self.is_secondary_model_activated:
+                if isinstance(secondary_model_source, np.ndarray):
+                    secondary_model_scale = model_scale if model_scale else self.secondary_model_scale
+                    stem_source = spec_utils.average_dual_sources(stem_source, secondary_model_source, secondary_model_scale)
+            
+            sf.write(stem_path, stem_source, samplerate, subtype=self.wav_type_set)
+            save_format(stem_path, self.save_format, self.mp3_bit_set) if not self.is_ensemble_mode else None
+            
+            self.write_to_console(DONE, base_text='')
+            self.set_progress_bar(0.95)
+
+    def run_mixer(self, mix, sources):
+        try:
+            if self.is_mixer_mode and len(sources) == 4:
+                mixer = MdxnetSet.Mixer(self.device, self.mixer_path).eval()
+                with torch.no_grad():
+                    mix = torch.tensor(mix, dtype=torch.float32)
+                    sources_ = torch.tensor(sources).detach()
+                    x = torch.cat([sources_, mix.unsqueeze(0)], 0)
+                    sources_ = mixer(x)
+                final_source = np.array(sources_)
+            else:
+                final_source = sources
+        except Exception as e:
+            error_name = f'{type(e).__name__}'
+            traceback_text = ''.join(traceback.format_tb(e.__traceback__))
+            message = f'{error_name}: "{e}"\n{traceback_text}"'
+            print('Mixer Failed: ', message)
+            final_source = sources
+            
+        return final_source
+
+class  SeperateMDX(SeperateAttributes):        
+
+    def seperate(self):
+        samplerate = 44100
+          
+        if self.primary_model_name == self.model_basename and self.primary_sources:
+            self.primary_source, self.secondary_source = self.load_cached_sources()
+        else:
+            self.start_inference_console_write()
+
+            if self.is_mdx_ckpt:
+                model_params = torch.load(self.model_path, map_location=lambda storage, loc: storage)['hyper_parameters']
+                self.dim_c, self.hop = model_params['dim_c'], model_params['hop_length']
+                separator = MdxnetSet.ConvTDFNet(**model_params)
+                self.model_run = separator.load_from_checkpoint(self.model_path).to(self.device).eval()
+            else:
+                ort_ = ort.InferenceSession(self.model_path, providers=self.run_type)
+                self.model_run = lambda spek:ort_.run(None, {'input': spek.cpu().numpy()})[0]
+
+            self.initialize_model_settings()
+            self.running_inference_console_write()
+            mdx_net_cut = True if self.primary_stem in MDX_NET_FREQ_CUT else False
+            mix, raw_mix, samplerate = prepare_mix(self.audio_file, self.chunks, self.margin, mdx_net_cut=mdx_net_cut)
+            source = self.demix_base(mix, is_ckpt=self.is_mdx_ckpt)[0]
+            self.write_to_console(DONE, base_text='')            
+
+        if self.is_secondary_model_activated:
+            if self.secondary_model:
+                self.secondary_source_primary, self.secondary_source_secondary = process_secondary_model(self.secondary_model, self.process_data, main_process_method=self.process_method)
+        
+        if not self.is_secondary_stem_only:
+            self.write_to_console(f'{SAVING_STEM[0]}{self.primary_stem}{SAVING_STEM[1]}') if not self.is_secondary_model else None
+            primary_stem_path = os.path.join(self.export_path, f'{self.audio_file_base}_({self.primary_stem}).wav')
+            if not isinstance(self.primary_source, np.ndarray):
+                self.primary_source = spec_utils.normalize(source, self.is_normalization).T
+            self.primary_source_map = {self.primary_stem: self.primary_source}
+            self.write_audio(primary_stem_path, self.primary_source, samplerate, self.secondary_source_primary)
+
+        if not self.is_primary_stem_only:
+            self.write_to_console(f'{SAVING_STEM[0]}{self.secondary_stem}{SAVING_STEM[1]}') if not self.is_secondary_model else None
+            secondary_stem_path = os.path.join(self.export_path, f'{self.audio_file_base}_({self.secondary_stem}).wav')
+            if not isinstance(self.secondary_source, np.ndarray):
+                raw_mix = self.demix_base(raw_mix, is_match_mix=True)[0] if mdx_net_cut else raw_mix
+                self.secondary_source, raw_mix = spec_utils.normalize_two_stem(source*self.compensate, raw_mix, self.is_normalization)
+            
+                if self.is_invert_spec:
+                    self.secondary_source = spec_utils.invert_stem(raw_mix, self.secondary_source)
+                else:
+                    self.secondary_source = (-self.secondary_source.T+raw_mix.T)
+
+            self.secondary_source_map = {self.secondary_stem: self.secondary_source}
+            self.write_audio(secondary_stem_path, self.secondary_source, samplerate, self.secondary_source_secondary)
+
+        torch.cuda.empty_cache()
+        secondary_sources = {**self.primary_source_map, **self.secondary_source_map}
+
+        self.cache_source(secondary_sources)
+
+        if self.is_secondary_model:
+            return secondary_sources
+
+    def initialize_model_settings(self):
+        self.n_bins = self.n_fft//2+1
+        self.trim = self.n_fft//2
+        self.chunk_size = self.hop * (self.dim_t-1)
+        self.window = torch.hann_window(window_length=self.n_fft, periodic=False).to(self.device)
+        self.freq_pad = torch.zeros([1, self.dim_c, self.n_bins-self.dim_f, self.dim_t]).to(self.device)
+        self.gen_size = self.chunk_size-2*self.trim
+
+    def initialize_mix(self, mix, is_ckpt=False):
+        if is_ckpt:
+            pad = self.gen_size + self.trim - ((mix.shape[-1]) % self.gen_size)
+            mixture = np.concatenate((np.zeros((2, self.trim), dtype='float32'),mix, np.zeros((2, pad), dtype='float32')), 1)
+            num_chunks = mixture.shape[-1] // self.gen_size
+            mix_waves = [mixture[:, i * self.gen_size: i * self.gen_size + self.chunk_size] for i in range(num_chunks)]
+        else:
+            mix_waves = []
+            n_sample = mix.shape[1]
+            pad = self.gen_size - n_sample%self.gen_size
+            mix_p = np.concatenate((np.zeros((2,self.trim)), mix, np.zeros((2,pad)), np.zeros((2,self.trim))), 1)
+            i = 0
+            while i < n_sample + pad:
+                waves = np.array(mix_p[:, i:i+self.chunk_size])
+                mix_waves.append(waves)
+                i += self.gen_size
+                
+        mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(self.device)
+
+        return mix_waves, pad
+    
+    def demix_base(self, mix, is_ckpt=False, is_match_mix=False):
+        chunked_sources = []
+        for slice in mix:
+            sources = []
+            tar_waves_ = []
+            mix_p = mix[slice]
+            mix_waves, pad = self.initialize_mix(mix_p, is_ckpt=is_ckpt)
+            mix_waves = mix_waves.split(self.mdx_batch_size)
+            pad = mix_p.shape[-1] if is_ckpt else -pad
+            with torch.no_grad():
+                for mix_wave in mix_waves:
+                    self.running_inference_progress_bar(len(mix)*len(mix_waves), is_match_mix=is_match_mix)
+                    tar_waves = self.run_model(mix_wave, is_ckpt=is_ckpt, is_match_mix=is_match_mix)
+                    tar_waves_.append(tar_waves)
+                tar_waves_ = np.vstack(tar_waves_)[:, :, self.trim:-self.trim] if is_ckpt else tar_waves_
+                tar_waves = np.concatenate(tar_waves_, axis=-1)[:, :pad]
+                start = 0 if slice == 0 else self.margin
+                end = None if slice == list(mix.keys())[::-1][0] or self.margin == 0 else -self.margin
+                sources.append(tar_waves[:,start:end]*(1/self.adjust))
+            chunked_sources.append(sources)
+        sources = np.concatenate(chunked_sources, axis=-1)
+        
+        return sources
+
+    def run_model(self, mix, is_ckpt=False, is_match_mix=False):
+        
+        spek = self.stft(mix.to(self.device))*self.adjust
+        spek[:, :, :3, :] *= 0 
+
+        if is_match_mix:
+            spec_pred = spek.cpu().numpy()
+        else:
+            spec_pred = -self.model_run(-spek)*0.5+self.model_run(spek)*0.5 if self.is_denoise else self.model_run(spek)
+
+        if is_ckpt:
+            return self.istft(spec_pred).cpu().detach().numpy()
+        else: 
+            return self.istft(torch.tensor(spec_pred).to(self.device)).to(cpu)[:,:,self.trim:-self.trim].transpose(0,1).reshape(2, -1).numpy()
+    
+    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=None):
+        freq_pad = self.freq_pad.repeat([x.shape[0],1,1,1]) if freq_pad is None else freq_pad
+        x = torch.cat([x, freq_pad], -2)
+        x = x.reshape([-1,2,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,2,self.chunk_size])
+
+class SeperateDemucs(SeperateAttributes):        
+
+    def seperate(self):
+
+        samplerate = 44100
+        source = None
+        model_scale = None
+        stem_source = None
+        stem_source_secondary = None
+        inst_mix = None
+        inst_raw_mix = None
+        raw_mix = None
+        inst_source = None
+        is_no_write = False
+        is_no_piano_guitar = False
+
+        if self.primary_model_name == self.model_basename and type(self.primary_sources) is dict and not self.pre_proc_model:
+            self.primary_source, self.secondary_source = self.load_cached_sources()
+        elif self.primary_model_name == self.model_basename and isinstance(self.primary_sources, np.ndarray) and not self.pre_proc_model:
+            source = self.primary_sources
+            self.load_cached_sources(is_4_stem_demucs=True)
+        else:
+            self.start_inference_console_write()
+
+            if self.is_gpu_conversion >= 0:
+                self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 
+            else:
+                self.device = torch.device('cpu')
+            
+            if self.demucs_version == DEMUCS_V1:
+                if str(self.model_path).endswith(".gz"):
+                    self.model_path = gzip.open(self.model_path, "rb")
+                klass, args, kwargs, state = torch.load(self.model_path)
+                self.demucs = klass(*args, **kwargs)
+                self.demucs.to(self.device) 
+                self.demucs.load_state_dict(state)
+            elif self.demucs_version == DEMUCS_V2:
+                self.demucs = auto_load_demucs_model_v2(self.demucs_source_list, self.model_path)
+                self.demucs.to(self.device) 
+                self.demucs.load_state_dict(torch.load(self.model_path))
+                self.demucs.eval()
+            else:  
+                self.demucs = HDemucs(sources=self.demucs_source_list)
+                self.demucs = _gm(name=os.path.splitext(os.path.basename(self.model_path))[0], 
+                                  repo=Path(os.path.dirname(self.model_path)))
+                self.demucs = demucs_segments(self.segment, self.demucs)
+                self.demucs.to(self.device)
+                self.demucs.eval()
+
+            if self.pre_proc_model:
+                if self.primary_stem not in [VOCAL_STEM, INST_STEM]:
+                    is_no_write = True
+                    self.write_to_console(DONE, base_text='')
+                    mix_no_voc = process_secondary_model(self.pre_proc_model, self.process_data, is_pre_proc_model=True)
+                    inst_mix, inst_raw_mix, inst_samplerate = prepare_mix(mix_no_voc[INST_STEM], self.chunks_demucs, self.margin_demucs)
+                    self.process_iteration()
+                    self.running_inference_console_write(is_no_write=is_no_write)
+                    inst_source = self.demix_demucs(inst_mix)
+                    inst_source = self.run_mixer(inst_raw_mix, inst_source)
+                    self.process_iteration()
+
+            self.running_inference_console_write(is_no_write=is_no_write) if not self.pre_proc_model else None
+            mix, raw_mix, samplerate = prepare_mix(self.audio_file, self.chunks_demucs, self.margin_demucs)
+            
+            if self.primary_model_name == self.model_basename and isinstance(self.primary_sources, np.ndarray) and self.pre_proc_model:
+                source = self.primary_sources
+            else:
+                source = self.demix_demucs(mix)
+                source = self.run_mixer(raw_mix, source)
+            
+            self.write_to_console(DONE, base_text='')
+            
+            del self.demucs
+            torch.cuda.empty_cache()
+            
+        if isinstance(inst_source, np.ndarray):
+            source_reshape = spec_utils.reshape_sources(inst_source[self.demucs_source_map[VOCAL_STEM]], source[self.demucs_source_map[VOCAL_STEM]])
+            inst_source[self.demucs_source_map[VOCAL_STEM]] = source_reshape
+            source = inst_source
+
+        if isinstance(source, np.ndarray):
+            if len(source) == 2:
+                self.demucs_source_map = DEMUCS_2_SOURCE_MAPPER
+            else:
+                self.demucs_source_map = DEMUCS_6_SOURCE_MAPPER if len(source) == 6 else DEMUCS_4_SOURCE_MAPPER
+
+                if len(source) == 6 and self.process_data['is_ensemble_master'] or len(source) == 6 and self.is_secondary_model:
+                    is_no_piano_guitar = True
+                    six_stem_other_source = list(source)
+                    six_stem_other_source = [i for n, i in enumerate(source) if n in [self.demucs_source_map[OTHER_STEM], self.demucs_source_map[GUITAR_STEM], self.demucs_source_map[PIANO_STEM]]]
+                    other_source = np.zeros_like(six_stem_other_source[0])
+                    for i in six_stem_other_source:
+                        other_source += i
+                    source_reshape = spec_utils.reshape_sources(source[self.demucs_source_map[OTHER_STEM]], other_source)
+                    source[self.demucs_source_map[OTHER_STEM]] = source_reshape
+
+        if (self.demucs_stems == ALL_STEMS and not self.process_data['is_ensemble_master']) or self.is_4_stem_ensemble:
+            self.cache_source(source)
+            
+            for stem_name, stem_value in self.demucs_source_map.items():
+                if self.is_secondary_model_activated and not self.is_secondary_model and not stem_value >= 4:
+                    if self.secondary_model_4_stem[stem_value]:
+                        model_scale = self.secondary_model_4_stem_scale[stem_value]
+                        stem_source_secondary = process_secondary_model(self.secondary_model_4_stem[stem_value], self.process_data, main_model_primary_stem_4_stem=stem_name, is_4_stem_demucs=True)
+                        if isinstance(stem_source_secondary, np.ndarray):
+                            stem_source_secondary = stem_source_secondary[1 if self.secondary_model_4_stem[stem_value].demucs_stem_count == 2 else stem_value]
+                            stem_source_secondary = spec_utils.normalize(stem_source_secondary, self.is_normalization).T
+                        elif type(stem_source_secondary) is dict:
+                            stem_source_secondary = stem_source_secondary[stem_name]
+                            
+                stem_source_secondary = None if stem_value >= 4 else stem_source_secondary
+                self.write_to_console(f'{SAVING_STEM[0]}{stem_name}{SAVING_STEM[1]}') if not self.is_secondary_model else None
+                stem_path = os.path.join(self.export_path, f'{self.audio_file_base}_({stem_name}).wav')
+                stem_source = spec_utils.normalize(source[stem_value], self.is_normalization).T
+                self.write_audio(stem_path, stem_source, samplerate, secondary_model_source=stem_source_secondary, model_scale=model_scale)
+
+            if self.is_secondary_model:    
+                return source
+        else:
+            if self.is_secondary_model_activated:
+                if self.secondary_model:
+                    self.secondary_source_primary, self.secondary_source_secondary = process_secondary_model(self.secondary_model, self.process_data, main_process_method=self.process_method)
+
+            if not self.is_secondary_stem_only:
+                self.write_to_console(f'{SAVING_STEM[0]}{self.primary_stem}{SAVING_STEM[1]}') if not self.is_secondary_model else None
+                primary_stem_path = os.path.join(self.export_path, f'{self.audio_file_base}_({self.primary_stem}).wav')
+                if not isinstance(self.primary_source, np.ndarray):
+                    self.primary_source = spec_utils.normalize(source[self.demucs_source_map[self.primary_stem]], self.is_normalization).T
+                self.primary_source_map = {self.primary_stem: self.primary_source}
+                self.write_audio(primary_stem_path, self.primary_source, samplerate, self.secondary_source_primary)
+
+            if not self.is_primary_stem_only:
+                def secondary_save(sec_stem_name, source, raw_mixture=None, is_inst_mixture=False):
+                    secondary_source = self.secondary_source if not is_inst_mixture else None
+                    self.write_to_console(f'{SAVING_STEM[0]}{sec_stem_name}{SAVING_STEM[1]}') if not self.is_secondary_model else None
+                    secondary_stem_path = os.path.join(self.export_path, f'{self.audio_file_base}_({sec_stem_name}).wav')
+                    secondary_source_secondary = None
+                    
+                    if not isinstance(secondary_source, np.ndarray):
+                        if self.is_demucs_combine_stems:
+                            source = list(source)
+                            if is_inst_mixture:
+                                source = [i for n, i in enumerate(source) if not n in [self.demucs_source_map[self.primary_stem], self.demucs_source_map[VOCAL_STEM]]]
+                            else:
+                                source.pop(self.demucs_source_map[self.primary_stem])
+                                
+                            source = source[:len(source) - 2] if is_no_piano_guitar else source
+                            secondary_source = np.zeros_like(source[0])
+                            for i in source:
+                                secondary_source += i
+                            secondary_source = spec_utils.normalize(secondary_source, self.is_normalization).T
+                        else:
+                            if not isinstance(raw_mixture, np.ndarray):
+                                raw_mixture = prepare_mix(self.audio_file, self.chunks_demucs, self.margin_demucs, is_missing_mix=True)
+       
+                            secondary_source, raw_mixture = spec_utils.normalize_two_stem(source[self.demucs_source_map[self.primary_stem]], raw_mixture, self.is_normalization)
+                            
+                            if self.is_invert_spec:
+                                secondary_source = spec_utils.invert_stem(raw_mixture, secondary_source)
+                            else:
+                                raw_mixture = spec_utils.reshape_sources(secondary_source, raw_mixture)
+                                secondary_source = (-secondary_source.T+raw_mixture.T)
+                            
+                    if not is_inst_mixture:
+                        self.secondary_source = secondary_source
+                        secondary_source_secondary = self.secondary_source_secondary
+                        self.secondary_source_map = {self.secondary_stem: self.secondary_source}
+
+                    self.write_audio(secondary_stem_path, secondary_source, samplerate, secondary_source_secondary)
+
+                secondary_save(self.secondary_stem, source, raw_mixture=raw_mix)
+                
+                if self.is_demucs_pre_proc_model_inst_mix and self.pre_proc_model and not self.is_4_stem_ensemble:
+                    secondary_save(f"{self.secondary_stem} {INST_STEM}", source, raw_mixture=inst_raw_mix, is_inst_mixture=True)
+
+            secondary_sources = {**self.primary_source_map, **self.secondary_source_map}
+
+            self.cache_source(secondary_sources)
+            
+            if self.is_secondary_model:    
+                return secondary_sources
+    
+    def demix_demucs(self, mix):
+        processed = {}
+
+        set_progress_bar = None if self.is_chunk_demucs else self.set_progress_bar
+
+        for nmix in mix:
+            self.progress_value += 1
+            self.set_progress_bar(0.1, (0.8/len(mix)*self.progress_value)) if self.is_chunk_demucs else None
+            cmix = mix[nmix]
+            cmix = torch.tensor(cmix, dtype=torch.float32)
+            ref = cmix.mean(0)        
+            cmix = (cmix - ref.mean()) / ref.std()
+            mix_infer = cmix 
+            
+            with torch.no_grad():
+                if self.demucs_version == DEMUCS_V1:
+                    sources = apply_model_v1(self.demucs, 
+                                                mix_infer.to(self.device), 
+                                                self.shifts, 
+                                                self.is_split_mode,
+                                                set_progress_bar=set_progress_bar)
+                elif self.demucs_version == DEMUCS_V2:
+                    sources = apply_model_v2(self.demucs, 
+                                                mix_infer.to(self.device), 
+                                                self.shifts,
+                                                self.is_split_mode,
+                                                self.overlap,
+                                                set_progress_bar=set_progress_bar)
+                else:
+                    sources = apply_model(self.demucs, 
+                                            mix_infer[None], 
+                                            self.shifts,
+                                            self.is_split_mode,
+                                            self.overlap,
+                                            static_shifts=1 if self.shifts == 0 else self.shifts,
+                                            set_progress_bar=set_progress_bar,
+                                            device=self.device)[0]
+            
+            sources = (sources * ref.std() + ref.mean()).cpu().numpy()
+            sources[[0,1]] = sources[[1,0]]
+            start = 0 if nmix == 0 else self.margin_demucs
+            end = None if nmix == list(mix.keys())[::-1][0] else -self.margin_demucs
+            if self.margin_demucs == 0:
+                end = None
+            processed[nmix] = sources[:,:,start:end].copy()
+            sources = list(processed.values())
+        sources = np.concatenate(sources, axis=-1)
+                        
+        return sources
+
+class SeperateVR(SeperateAttributes):        
+
+    def seperate(self):
+        if self.primary_model_name == self.model_basename and self.primary_sources:
+            self.primary_source, self.secondary_source = self.load_cached_sources()
+        else:
+            self.start_inference_console_write()
+            if self.is_gpu_conversion >= 0:
+                if OPERATING_SYSTEM == 'Darwin':
+                    device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
+                else:
+                    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 
+            else:
+                device = torch.device('cpu')
+
+            nn_arch_sizes = [
+                31191, # default
+                33966, 56817, 123821, 123812, 129605, 218409, 537238, 537227]
+            vr_5_1_models = [56817, 218409]
+            model_size = math.ceil(os.stat(self.model_path).st_size / 1024)
+            nn_arch_size = min(nn_arch_sizes, key=lambda x:abs(x-model_size))
+
+            if nn_arch_size in vr_5_1_models or self.is_vr_51_model:
+                self.model_run = nets_new.CascadedNet(self.mp.param['bins'] * 2, nn_arch_size, nout=self.model_capacity[0], nout_lstm=self.model_capacity[1])
+            else:
+                self.model_run = nets.determine_model_capacity(self.mp.param['bins'] * 2, nn_arch_size)
+                            
+            self.model_run.load_state_dict(torch.load(self.model_path, map_location=cpu)) 
+            self.model_run.to(device) 
+
+            self.running_inference_console_write()
+                        
+            y_spec, v_spec = self.inference_vr(self.loading_mix(), device, self.aggressiveness)
+            self.write_to_console(DONE, base_text='')
+            
+        if self.is_secondary_model_activated:
+            if self.secondary_model:
+                self.secondary_source_primary, self.secondary_source_secondary = process_secondary_model(self.secondary_model, self.process_data, main_process_method=self.process_method)
+
+        if not self.is_secondary_stem_only:
+            self.write_to_console(f'{SAVING_STEM[0]}{self.primary_stem}{SAVING_STEM[1]}') if not self.is_secondary_model else None
+            primary_stem_path = os.path.join(self.export_path, f'{self.audio_file_base}_({self.primary_stem}).wav')
+            if not isinstance(self.primary_source, np.ndarray):
+                self.primary_source = spec_utils.normalize(self.spec_to_wav(y_spec), self.is_normalization).T
+                if not self.model_samplerate == 44100:
+                    self.primary_source = librosa.resample(self.primary_source.T, orig_sr=self.model_samplerate, target_sr=44100).T
+                
+            self.primary_source_map = {self.primary_stem: self.primary_source}
+            
+            self.write_audio(primary_stem_path, self.primary_source, 44100, self.secondary_source_primary)
+
+        if not self.is_primary_stem_only:
+            self.write_to_console(f'{SAVING_STEM[0]}{self.secondary_stem}{SAVING_STEM[1]}') if not self.is_secondary_model else None
+            secondary_stem_path = os.path.join(self.export_path, f'{self.audio_file_base}_({self.secondary_stem}).wav')
+            if not isinstance(self.secondary_source, np.ndarray):
+                self.secondary_source = self.spec_to_wav(v_spec)
+                self.secondary_source = spec_utils.normalize(self.spec_to_wav(v_spec), self.is_normalization).T
+                if not self.model_samplerate == 44100:
+                    self.secondary_source = librosa.resample(self.secondary_source.T, orig_sr=self.model_samplerate, target_sr=44100).T
+            
+            self.secondary_source_map = {self.secondary_stem: self.secondary_source}
+            
+            self.write_audio(secondary_stem_path, self.secondary_source, 44100, self.secondary_source_secondary)
+            
+        torch.cuda.empty_cache()
+        secondary_sources = {**self.primary_source_map, **self.secondary_source_map}
+        self.cache_source(secondary_sources)
+
+        if self.is_secondary_model:
+            return secondary_sources
+            
+    def loading_mix(self):
+
+        X_wave, X_spec_s = {}, {}
+        
+        bands_n = len(self.mp.param['band'])
+        
+        for d in range(bands_n, 0, -1):        
+            bp = self.mp.param['band'][d]
+        
+            if OPERATING_SYSTEM == 'Darwin':
+                wav_resolution = 'polyphase' if SYSTEM_PROC == ARM or ARM in SYSTEM_ARCH else bp['res_type']
+            else:
+                wav_resolution = bp['res_type']
+        
+            if d == bands_n: # high-end band
+                X_wave[d], _ = librosa.load(self.audio_file, sr=bp['sr'], mono=False, dtype=np.float32, res_type=wav_resolution)
+                    
+                if not np.any(X_wave[d]) and self.audio_file.endswith('.mp3'):
+                    X_wave[d] = rerun_mp3(self.audio_file, bp['sr'])
+
+                if X_wave[d].ndim == 1:
+                    X_wave[d] = np.asarray([X_wave[d], X_wave[d]])
+            else: # lower bands
+                X_wave[d] = librosa.resample(X_wave[d+1], orig_sr=self.mp.param['band'][d+1]['sr'], target_sr=bp['sr'], res_type=wav_resolution)
+                
+            X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(X_wave[d], bp['hl'], bp['n_fft'], self.mp.param['mid_side'], 
+                                                            self.mp.param['mid_side_b2'], self.mp.param['reverse'])
+            
+            if d == bands_n and self.high_end_process != 'none':
+                self.input_high_end_h = (bp['n_fft']//2 - bp['crop_stop']) + (self.mp.param['pre_filter_stop'] - self.mp.param['pre_filter_start'])
+                self.input_high_end = X_spec_s[d][:, bp['n_fft']//2-self.input_high_end_h:bp['n_fft']//2, :]
+
+        X_spec = spec_utils.combine_spectrograms(X_spec_s, self.mp)
+        
+        del X_wave, X_spec_s
+
+        return X_spec
+
+    def inference_vr(self, X_spec, device, aggressiveness):
+        def _execute(X_mag_pad, roi_size):
+            X_dataset = []
+            patches = (X_mag_pad.shape[2] - 2 * self.model_run.offset) // roi_size
+            total_iterations = patches//self.batch_size if not self.is_tta else (patches//self.batch_size)*2
+            for i in range(patches):
+                start = i * roi_size
+                X_mag_window = X_mag_pad[:, :, start:start + self.window_size]
+                X_dataset.append(X_mag_window)
+
+            X_dataset = np.asarray(X_dataset)
+            self.model_run.eval()
+            with torch.no_grad():
+                mask = []
+                for i in range(0, patches, self.batch_size):
+                    self.progress_value += 1
+                    if self.progress_value >= total_iterations:
+                        self.progress_value = total_iterations
+                    self.set_progress_bar(0.1, 0.8/total_iterations*self.progress_value)
+                    X_batch = X_dataset[i: i + self.batch_size]
+                    X_batch = torch.from_numpy(X_batch).to(device)
+                    pred = self.model_run.predict_mask(X_batch)
+                    if not pred.size()[3] > 0:
+                        raise Exception(ERROR_MAPPER[WINDOW_SIZE_ERROR])
+                    pred = pred.detach().cpu().numpy()
+                    pred = np.concatenate(pred, axis=2)
+                    mask.append(pred)
+                if len(mask) == 0:
+                    raise Exception(ERROR_MAPPER[WINDOW_SIZE_ERROR])
+                
+                mask = np.concatenate(mask, axis=2)
+            return mask
+
+        def postprocess(mask, X_mag, X_phase):
+            
+            is_non_accom_stem = False
+            for stem in NON_ACCOM_STEMS:
+                if stem == self.primary_stem:
+                    is_non_accom_stem = True
+                    
+            mask = spec_utils.adjust_aggr(mask, is_non_accom_stem, aggressiveness)
+
+            if self.is_post_process:
+                mask = spec_utils.merge_artifacts(mask, thres=self.post_process_threshold)
+
+            y_spec = mask * X_mag * np.exp(1.j * X_phase)
+            v_spec = (1 - mask) * X_mag * np.exp(1.j * X_phase)
+        
+            return y_spec, v_spec
+        X_mag, X_phase = spec_utils.preprocess(X_spec)
+        n_frame = X_mag.shape[2]
+        pad_l, pad_r, roi_size = spec_utils.make_padding(n_frame, self.window_size, self.model_run.offset)
+        X_mag_pad = np.pad(X_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode='constant')
+        X_mag_pad /= X_mag_pad.max()
+        mask = _execute(X_mag_pad, roi_size)
+        
+        if self.is_tta:
+            pad_l += roi_size // 2
+            pad_r += roi_size // 2
+            X_mag_pad = np.pad(X_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode='constant')
+            X_mag_pad /= X_mag_pad.max()
+            mask_tta = _execute(X_mag_pad, roi_size)
+            mask_tta = mask_tta[:, :, roi_size // 2:]
+            mask = (mask[:, :, :n_frame] + mask_tta[:, :, :n_frame]) * 0.5
+        else:
+            mask = mask[:, :, :n_frame]
+
+        y_spec, v_spec = postprocess(mask, X_mag, X_phase)
+        
+        return y_spec, v_spec
+
+    def spec_to_wav(self, spec):
+        
+        if self.high_end_process.startswith('mirroring'):        
+            input_high_end_ = spec_utils.mirroring(self.high_end_process, spec, self.input_high_end, self.mp)
+            wav = spec_utils.cmb_spectrogram_to_wave(spec, self.mp, self.input_high_end_h, input_high_end_)       
+        else:
+            wav = spec_utils.cmb_spectrogram_to_wave(spec, self.mp)
+            
+        return wav
+   
+def process_secondary_model(secondary_model: ModelData, process_data, main_model_primary_stem_4_stem=None, is_4_stem_demucs=False, main_process_method=None, is_pre_proc_model=False):
+        
+    if not is_pre_proc_model:
+        process_iteration = process_data['process_iteration']
+        process_iteration()
+    
+    if secondary_model.process_method == VR_ARCH_TYPE:
+        seperator = SeperateVR(secondary_model, process_data, main_model_primary_stem_4_stem=main_model_primary_stem_4_stem, main_process_method=main_process_method)
+    if secondary_model.process_method == MDX_ARCH_TYPE:
+        seperator = SeperateMDX(secondary_model, process_data, main_model_primary_stem_4_stem=main_model_primary_stem_4_stem, main_process_method=main_process_method)
+    if secondary_model.process_method == DEMUCS_ARCH_TYPE:
+        seperator = SeperateDemucs(secondary_model, process_data, main_model_primary_stem_4_stem=main_model_primary_stem_4_stem, main_process_method=main_process_method)
+        
+    secondary_sources = seperator.seperate()
+
+    if type(secondary_sources) is dict and not is_4_stem_demucs and not is_pre_proc_model:
+        return gather_sources(secondary_model.primary_model_primary_stem, STEM_PAIR_MAPPER[secondary_model.primary_model_primary_stem], secondary_sources)
+    else:
+        return secondary_sources
+    
+def gather_sources(primary_stem_name, secondary_stem_name, secondary_sources: dict):
+    
+    source_primary = False
+    source_secondary = False
+
+    for key, value in secondary_sources.items():
+        if key in primary_stem_name:
+            source_primary = value
+        if key in secondary_stem_name:
+            source_secondary = value
+
+    return source_primary, source_secondary
+        
+def prepare_mix(mix, chunk_set, margin_set, mdx_net_cut=False, is_missing_mix=False):
+
+    audio_path = mix
+    samplerate = 44100
+
+    if not isinstance(mix, np.ndarray):
+        mix, samplerate = librosa.load(mix, mono=False, sr=44100)
+    else:
+        mix = mix.T
+
+    if not np.any(mix) and audio_path.endswith('.mp3'):
+        mix = rerun_mp3(audio_path)
+
+    if mix.ndim == 1:
+        mix = np.asfortranarray([mix,mix])
+
+    def get_segmented_mix(chunk_set=chunk_set):
+        segmented_mix = {}
+        
+        samples = mix.shape[-1]
+        margin = margin_set
+        chunk_size = chunk_set*44100
+        assert not margin == 0, 'margin cannot be zero!'
+        
+        if margin > chunk_size:
+            margin = chunk_size
+        if chunk_set == 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
+            segmented_mix[skip] = mix[:,start:end].copy()
+            if end == samples:
+                break
+            
+        return segmented_mix
+
+    if is_missing_mix:
+        return mix
+    else:
+        segmented_mix = get_segmented_mix()
+        raw_mix = get_segmented_mix(chunk_set=0) if mdx_net_cut else mix
+        return segmented_mix, raw_mix, samplerate
+
+def rerun_mp3(audio_file, sample_rate=44100):
+
+    with audioread.audio_open(audio_file) as f:
+        track_length = int(f.duration)
+
+    return librosa.load(audio_file, duration=track_length, mono=False, sr=sample_rate)[0]
+
+def save_format(audio_path, save_format, mp3_bit_set):
+    
+    if not save_format == WAV:
+        
+        if OPERATING_SYSTEM == 'Darwin':
+            FFMPEG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'ffmpeg')
+            pydub.AudioSegment.converter = FFMPEG_PATH
+        
+        musfile = pydub.AudioSegment.from_wav(audio_path)
+        
+        if save_format == FLAC:
+            audio_path_flac = audio_path.replace(".wav", ".flac")
+            musfile.export(audio_path_flac, format="flac")  
+        
+        if save_format == MP3:
+            audio_path_mp3 = audio_path.replace(".wav", ".mp3")
+            musfile.export(audio_path_mp3, format="mp3", bitrate=mp3_bit_set)
+        
+        try:
+            os.remove(audio_path)
+        except Exception as e:
+            print(e)

lib/separators.py

@@ -0,0 +1,367 @@
+import hashlib
+import json
+from multiprocessing.pool import ThreadPool
+import os
+from .utils import ObjectNamespace
+import numpy as np
+import torch
+from tqdm import tqdm
+from .mdx import MDXModel
+from .uvr5_pack.constants import MDX_NET_FREQ_CUT
+from .uvr5_pack.vr_network.model_param_init import ModelParameters
+from .uvr5_pack.vr_network.nets_new import CascadedNet
+from .uvr5_pack.vr_network.nets import CascadedASPPNet
+from .audio import remix_audio
+import librosa
+from .uvr5_pack import spec_utils
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+
+class UVR5Base:
+    
+    def __init__(self, agg, model_path, device, is_half,**kwargs):
+        self.model_path = model_path
+        self.device = device
+        self.data = {
+            # Processing Options
+            "postprocess": False,
+            "tta": False,
+            # Constants
+            "window_size": 512,
+            "agg": agg,
+            "high_end_process": "mirroring",
+        }
+        mp = ModelParameters(os.path.join(dir_path,"uvr5_pack","vr_network","modelparams","4band_v2.json"))
+        model = CascadedASPPNet(mp.param["bins"] * 2)
+        cpk = torch.load(model_path, map_location=self.device)
+        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
+    
+    @staticmethod
+    def get_model_params(model_path):
+        try:
+            with open(model_path, 'rb') as f:
+                f.seek(- 10000 * 1024, 2)
+                model_hash = hashlib.md5(f.read()).hexdigest()
+        except:
+            model_hash = hashlib.md5(open(model_path, 'rb').read()).hexdigest()
+
+        print(model_hash)
+        model_settings_json = os.path.splitext(model_path)[0]+".json"
+        model_data_json = os.path.join(os.path.dirname(model_path),"model_data.json")
+
+        if os.path.isfile(model_settings_json):
+            return json.load(open(model_settings_json))
+        elif os.path.isfile(model_data_json):
+            with open(model_data_json,"r") as d:
+                hash_mapper = json.loads(d.read())
+
+            for hash, settings in hash_mapper.items():
+                if model_hash in hash:
+                    return settings
+        return None
+
+    def inference(self, X_spec, aggressiveness):
+        """
+        data ： dic configs
+        """
+        data = self.data
+        device = self.device
+        model = self.model
+
+        def _execute(
+            X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half=True
+        ):
+            model.eval()
+            with torch.no_grad():
+                preds = []
+
+                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 = model.predict(X_mag_window, aggressiveness)
+
+                    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 = spec_utils.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)
+
+    def process_vocals(self,v_spec_m,input_high_end,input_high_end_h,return_dict={}):
+        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)
+        print(f"vocals done: {wav_vocals.shape}")
+        return_dict["vocals"] = remix_audio((wav_vocals,return_dict["sr"]),to_int16=True,axis=0)
+        return return_dict["vocals"]
+    
+    def process_instrumental(self,y_spec_m,input_high_end,input_high_end_h,return_dict={}):
+        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)
+        print(f"instruments done: {wav_instrument.shape}")
+        return_dict["instrumentals"] = remix_audio((wav_instrument,return_dict["sr"]),to_int16=True,axis=0)
+        return return_dict["instrumentals"] 
+    
+    def process_audio(self,y_spec_m,v_spec_m,input_high_end,input_high_end_h):
+        return_dict = {
+            "sr": self.mp.param["sr"]
+        }
+        
+        with ThreadPool(2) as pool:
+            pool.apply(self.process_vocals, args=(v_spec_m,input_high_end,input_high_end_h,return_dict))
+            pool.apply(self.process_instrumental, args=(y_spec_m,input_high_end,input_high_end_h,return_dict))
+ 
+        return return_dict
+
+    def run_inference(self, music_file):
+        X_wave,  X_spec_s = {}, {}
+        bands_n = len(self.mp.param["band"])
+        # print(bands_n)
+        for d in range(bands_n, 0, -1):
+            bp = self.mp.param["band"][d]
+            if d == bands_n:  # high-end band
+                input_audio = librosa.load(music_file,sr=bp["sr"],res_type=bp["res_type"])
+                X_wave[d] = input_audio[0]
+                if X_wave[d].ndim == 1:
+                    X_wave[d] = np.asfortranarray([X_wave[d], X_wave[d]])
+            else:  # lower bands
+                X_wave[d] = librosa.resample(
+                    X_wave[d + 1],
+                    orig_sr=self.mp.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=bp["res_type"],
+                )
+            # Stft of wave source
+            X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(
+                X_wave[d],
+                bp["hl"],
+                bp["n_fft"],
+                self.mp.param["mid_side"],
+                self.mp.param["mid_side_b2"],
+                self.mp.param["reverse"],
+            )
+            # pdb.set_trace()
+            if d == bands_n and self.data["high_end_process"] != "none":
+                input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (
+                    self.mp.param["pre_filter_stop"] - self.mp.param["pre_filter_start"]
+                )
+                input_high_end = X_spec_s[d][
+                    :, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, :
+                ]
+
+        X_spec_m = spec_utils.combine_spectrograms(X_spec_s, self.mp)
+        aggresive_set = float(self.data["agg"] / 100)
+        aggressiveness = {
+            "value": aggresive_set,
+            "split_bin": self.mp.param["band"][1]["crop_stop"],
+        }
+
+        # pred, X_mag, X_phase = run_inference(X_spec_m, self.device, self.models, aggressiveness, self.data)
+        with torch.no_grad():
+            pred, X_mag, X_phase = self.inference(X_spec_m, aggressiveness)
+
+    #     return pred, X_mag, X_phase, X_spec_m, input_high_end,input_high_end_h
+
+    # def run_process_audio(self, pred, X_mag, X_phase, X_spec_m, input_high_end,input_high_end_h):
+        # Postprocess
+        if self.data["postprocess"]:
+            pred_inv = np.clip(X_mag - pred, 0, np.inf)
+            pred = spec_utils.mask_silence(pred, pred_inv)
+        y_spec_m = pred * X_phase
+        v_spec_m = X_spec_m - y_spec_m
+
+        
+        return_dict = self.process_audio(y_spec_m,v_spec_m,input_high_end,input_high_end_h)
+        return_dict["input_audio"] = input_audio
+        
+        return return_dict
+
+class UVR5New(UVR5Base):
+    def __init__(self, agg, model_path, device, is_half, dereverb, **kwargs):
+        self.model_path = model_path
+        self.device = device
+        self.data = {
+            # Processing Options
+            "postprocess": False,
+            "tta": False,
+            # Constants
+            "window_size": 512,
+            "agg": agg,
+            "high_end_process": "mirroring",
+        }
+        mp = ModelParameters(os.path.join(dir_path,"uvr5_pack","vr_network","modelparams","4band_v3.json"))
+        nout = 64 if dereverb else 48
+        model = CascadedNet(mp.param["bins"] * 2, nout)
+        cpk = torch.load(model_path, map_location=self.device)
+        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
+    
+class MDXNet:
+    def __init__(self, model_path, chunks=15,denoise=False,num_threads=1,device="cpu",**kwargs):
+
+        self.chunks = chunks
+        self.sr = 44100
+        
+        self.args = ObjectNamespace(**kwargs)
+        self.denoise = denoise
+        self.num_threads = num_threads
+
+        self.device = device
+
+        self.is_mdx_ckpt = "ckpt" in model_path
+
+        self.model = MDXModel(model_path, device=self.device,chunks=self.chunks,margin=self.sr)
+
+    def __del__(self):
+        del self.model
+
+    def process_audio(self,primary,secondary,target_sr=None):
+        target_sr =  self.sr if target_sr is None else target_sr
+        # foreground is processed data
+        vocals,instrumental = (secondary,primary) if "instrument" in self.model.params.stem_name.lower() else (primary,secondary)
+        with ThreadPool(2) as pool:
+            results = pool.starmap(remix_audio, [
+                ((instrumental,self.sr),target_sr,False,True,self.sr!=target_sr,True),
+                ((vocals,self.sr),target_sr,False,True,self.sr!=target_sr,True)
+            ])
+
+        return_dict = {
+            "sr": target_sr,
+            "instrumentals": results[0],
+            "vocals": results[1]
+        }
+        return return_dict
+    
+    def run_inference(self, audio_path):
+        
+        mdx_net_cut = True if self.model.params.stem_name in MDX_NET_FREQ_CUT else False
+        mix, raw_mix, samplerate = prepare_mix(audio_path, self.model.chunks, self.model.margin, mdx_net_cut=mdx_net_cut)
+        wave_processed = self.model.demix_base(mix, is_ckpt=self.is_mdx_ckpt)[0]
+        
+    
+        raw_mix = self.model.demix_base(raw_mix, is_match_mix=True)[0] if mdx_net_cut else raw_mix
+
+        return_dict = self.process_audio(primary=wave_processed,secondary=(raw_mix-wave_processed),target_sr=samplerate)
+        return_dict["input_audio"] = (raw_mix, samplerate)
+
+        return return_dict
+
+    
+def prepare_mix(mix, chunk_set, margin_set, mdx_net_cut=False):
+
+    samplerate = 44100
+
+    if not isinstance(mix, np.ndarray):
+        mix, samplerate = librosa.load(mix, mono=False, sr=44100)
+    else:
+        mix = mix.T
+
+    if mix.ndim == 1:
+        mix = np.asfortranarray([mix,mix])
+
+    def get_segmented_mix(chunk_set=chunk_set):
+        segmented_mix = {}
+        
+        samples = mix.shape[-1]
+        margin = margin_set
+        chunk_size = chunk_set*44100
+        assert not margin == 0, 'margin cannot be zero!'
+        
+        if margin > chunk_size:
+            margin = chunk_size
+        if chunk_set == 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
+            segmented_mix[skip] = mix[:,start:end].copy()
+            if end == samples:
+                break
+            
+        return segmented_mix
+
+    
+    segmented_mix = get_segmented_mix()
+    raw_mix = get_segmented_mix(chunk_set=0) if mdx_net_cut else mix
+    return segmented_mix, raw_mix, samplerate

lib/slicer2.py

@@ -0,0 +1,264 @@
+import numpy as np
+from scipy import signal
+
+# This function is obtained from librosa.
+def get_rms(
+    y,
+    frame_length=2048,
+    hop_length=512,
+    pad_mode="constant",
+):
+    padding = (int(frame_length // 2), int(frame_length // 2))
+    y = np.pad(y, padding, mode=pad_mode)
+
+    axis = -1
+    # put our new within-frame axis at the end for now
+    out_strides = y.strides + tuple([y.strides[axis]])
+    # Reduce the shape on the framing axis
+    x_shape_trimmed = list(y.shape)
+    x_shape_trimmed[axis] -= frame_length - 1
+    out_shape = tuple(x_shape_trimmed) + tuple([frame_length])
+    xw = np.lib.stride_tricks.as_strided(y, shape=out_shape, strides=out_strides)
+    if axis < 0:
+        target_axis = axis - 1
+    else:
+        target_axis = axis + 1
+    xw = np.moveaxis(xw, -1, target_axis)
+    # Downsample along the target axis
+    slices = [slice(None)] * xw.ndim
+    slices[axis] = slice(0, None, hop_length)
+    x = xw[tuple(slices)]
+
+    # Calculate power
+    power = np.mean(np.abs(x) ** 2, axis=-2, keepdims=True)
+
+    return np.sqrt(power)
+
+
+class Slicer:
+    def __init__(
+        self,
+        sr: int,
+        threshold: float = -40.0,
+        min_length: int = 5000,
+        min_interval: int = 300,
+        hop_size: int = 20,
+        max_sil_kept: int = 5000,
+    ):
+        if not min_length >= min_interval >= hop_size:
+            raise ValueError(
+                "The following condition must be satisfied: min_length >= min_interval >= hop_size"
+            )
+        if not max_sil_kept >= hop_size:
+            raise ValueError(
+                "The following condition must be satisfied: max_sil_kept >= hop_size"
+            )
+        min_interval = sr * min_interval / 1000
+        self.threshold = 10 ** (threshold / 20.0)
+        self.hop_size = round(sr * hop_size / 1000)
+        self.win_size = min(round(min_interval), 4 * self.hop_size)
+        self.min_length = round(sr * min_length / 1000 / self.hop_size)
+        self.min_interval = round(min_interval / self.hop_size)
+        self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)
+        self.bh, self.ah = signal.butter(N=5, Wn=48, btype="high", fs=sr)
+
+    def _apply_slice(self, waveform, begin, end):
+        if len(waveform.shape) > 1:
+            return waveform[
+                :, begin * self.hop_size : min(waveform.shape[1], end * self.hop_size)
+            ]
+        else:
+            return waveform[
+                begin * self.hop_size : min(waveform.shape[0], end * self.hop_size)
+            ]
+
+    # @timeit
+    def slice(self, waveform):
+        # zero phased digital filter cause pre-ringing noise...
+        waveform = signal.lfilter(self.bh, self.ah, waveform)
+
+        if len(waveform.shape) > 1:
+            samples = waveform.mean(axis=0)
+        else:
+            samples = waveform
+        if samples.shape[0] <= self.min_length:
+            return [waveform]
+        rms_list = get_rms(
+            y=samples, frame_length=self.win_size, hop_length=self.hop_size
+        ).squeeze(0)
+        sil_tags = []
+        silence_start = None
+        clip_start = 0
+        for i, rms in enumerate(rms_list):
+            # Keep looping while frame is silent.
+            if rms < self.threshold:
+                # Record start of silent frames.
+                if silence_start is None:
+                    silence_start = i
+                continue
+            # Keep looping while frame is not silent and silence start has not been recorded.
+            if silence_start is None:
+                continue
+            # Clear recorded silence start if interval is not enough or clip is too short
+            is_leading_silence = silence_start == 0 and i > self.max_sil_kept
+            need_slice_middle = (
+                i - silence_start >= self.min_interval
+                and i - clip_start >= self.min_length
+            )
+            if not is_leading_silence and not need_slice_middle:
+                silence_start = None
+                continue
+            # Need slicing. Record the range of silent frames to be removed.
+            if i - silence_start <= self.max_sil_kept:
+                pos = rms_list[silence_start : i + 1].argmin() + silence_start
+                if silence_start == 0:
+                    sil_tags.append((0, pos))
+                else:
+                    sil_tags.append((pos, pos))
+                clip_start = pos
+            elif i - silence_start <= self.max_sil_kept * 2:
+                pos = rms_list[
+                    i - self.max_sil_kept : silence_start + self.max_sil_kept + 1
+                ].argmin()
+                pos += i - self.max_sil_kept
+                pos_l = (
+                    rms_list[
+                        silence_start : silence_start + self.max_sil_kept + 1
+                    ].argmin()
+                    + silence_start
+                )
+                pos_r = (
+                    rms_list[i - self.max_sil_kept : i + 1].argmin()
+                    + i
+                    - self.max_sil_kept
+                )
+                if silence_start == 0:
+                    sil_tags.append((0, pos_r))
+                    clip_start = pos_r
+                else:
+                    sil_tags.append((min(pos_l, pos), max(pos_r, pos)))
+                    clip_start = max(pos_r, pos)
+            else:
+                pos_l = (
+                    rms_list[
+                        silence_start : silence_start + self.max_sil_kept + 1
+                    ].argmin()
+                    + silence_start
+                )
+                pos_r = (
+                    rms_list[i - self.max_sil_kept : i + 1].argmin()
+                    + i
+                    - self.max_sil_kept
+                )
+                if silence_start == 0:
+                    sil_tags.append((0, pos_r))
+                else:
+                    sil_tags.append((pos_l, pos_r))
+                clip_start = pos_r
+            silence_start = None
+        # Deal with trailing silence.
+        total_frames = rms_list.shape[0]
+        if (
+            silence_start is not None
+            and total_frames - silence_start >= self.min_interval
+        ):
+            silence_end = min(total_frames, silence_start + self.max_sil_kept)
+            pos = rms_list[silence_start : silence_end + 1].argmin() + silence_start
+            sil_tags.append((pos, total_frames + 1))
+        # Apply and return slices.
+        if len(sil_tags) == 0:
+            return [waveform]
+        else:
+            chunks = []
+            if sil_tags[0][0] > 0:
+                chunks.append(self._apply_slice(waveform, 0, sil_tags[0][0]))
+            for i in range(len(sil_tags) - 1):
+                chunks.append(
+                    self._apply_slice(waveform, sil_tags[i][1], sil_tags[i + 1][0])
+                )
+            if sil_tags[-1][1] < total_frames:
+                chunks.append(
+                    self._apply_slice(waveform, sil_tags[-1][1], total_frames)
+                )
+            return chunks
+
+
+def main():
+    import os.path
+    from argparse import ArgumentParser
+
+    import librosa
+    import soundfile
+
+    parser = ArgumentParser()
+    parser.add_argument("audio", type=str, help="The audio to be sliced")
+    parser.add_argument(
+        "--out", type=str, help="Output directory of the sliced audio clips"
+    )
+    parser.add_argument(
+        "--db_thresh",
+        type=float,
+        required=False,
+        default=-40,
+        help="The dB threshold for silence detection",
+    )
+    parser.add_argument(
+        "--min_length",
+        type=int,
+        required=False,
+        default=5000,
+        help="The minimum milliseconds required for each sliced audio clip",
+    )
+    parser.add_argument(
+        "--min_interval",
+        type=int,
+        required=False,
+        default=300,
+        help="The minimum milliseconds for a silence part to be sliced",
+    )
+    parser.add_argument(
+        "--hop_size",
+        type=int,
+        required=False,
+        default=10,
+        help="Frame length in milliseconds",
+    )
+    parser.add_argument(
+        "--max_sil_kept",
+        type=int,
+        required=False,
+        default=500,
+        help="The maximum silence length kept around the sliced clip, presented in milliseconds",
+    )
+    args = parser.parse_args()
+    out = args.out
+    if out is None:
+        out = os.path.dirname(os.path.abspath(args.audio))
+    audio, sr = librosa.load(args.audio, sr=None, mono=False)
+    slicer = Slicer(
+        sr=sr,
+        threshold=args.db_thresh,
+        min_length=args.min_length,
+        min_interval=args.min_interval,
+        hop_size=args.hop_size,
+        max_sil_kept=args.max_sil_kept,
+    )
+    chunks = slicer.slice(audio)
+    if not os.path.exists(out):
+        os.makedirs(out)
+    for i, chunk in enumerate(chunks):
+        if len(chunk.shape) > 1:
+            chunk = chunk.T
+        soundfile.write(
+            os.path.join(
+                out,
+                f"%s_%d.wav"
+                % (os.path.basename(args.audio).rsplit(".", maxsplit=1)[0], i),
+            ),
+            chunk,
+            sr,
+        )
+
+
+if __name__ == "__main__":
+    main()

lib/train/cmd.txt

@@ -0,0 +1 @@
+python train_nsf_sim_cache_sid.py -c configs/mi_mix40k_nsf_co256_cs1sid_ms2048.json -m ft-mi

lib/train/data_utils.py

@@ -0,0 +1,512 @@
+import os, traceback
+import numpy as np
+import torch
+import torch.utils.data
+
+from .mel_processing import spectrogram_torch
+from .utils import load_wav_to_torch, load_filepaths_and_text
+
+
+class TextAudioLoaderMultiNSFsid(torch.utils.data.Dataset):
+    """
+    1) loads audio, text pairs
+    2) normalizes text and converts them to sequences of integers
+    3) computes spectrograms from audio files.
+    """
+
+    def __init__(self, audiopaths_and_text, hparams):
+        self.audiopaths_and_text = load_filepaths_and_text(audiopaths_and_text)
+        self.max_wav_value = hparams.max_wav_value
+        self.sampling_rate = hparams.sampling_rate
+        self.filter_length = hparams.filter_length
+        self.hop_length = hparams.hop_length
+        self.win_length = hparams.win_length
+        self.sampling_rate = hparams.sampling_rate
+        self.min_text_len = getattr(hparams, "min_text_len", 1)
+        self.max_text_len = getattr(hparams, "max_text_len", 5000)
+        self._filter()
+
+    def _filter(self):
+        """
+        Filter text & store spec lengths
+        """
+        # Store spectrogram lengths for Bucketing
+        # wav_length ~= file_size / (wav_channels * Bytes per dim) = file_size / (1 * 2)
+        # spec_length = wav_length // hop_length
+        audiopaths_and_text_new = []
+        lengths = []
+        for audiopath, text, pitch, pitchf, dv in self.audiopaths_and_text:
+            if self.min_text_len <= len(text) and len(text) <= self.max_text_len:
+                audiopaths_and_text_new.append([audiopath, text, pitch, pitchf, dv])
+                lengths.append(os.path.getsize(audiopath) // (3 * self.hop_length))
+        self.audiopaths_and_text = audiopaths_and_text_new
+        self.lengths = lengths
+
+    def get_sid(self, sid):
+        sid = torch.LongTensor([int(sid)])
+        return sid
+
+    def get_audio_text_pair(self, audiopath_and_text):
+        # separate filename and text
+        file = audiopath_and_text[0]
+        phone = audiopath_and_text[1]
+        pitch = audiopath_and_text[2]
+        pitchf = audiopath_and_text[3]
+        dv = audiopath_and_text[4]
+
+        phone, pitch, pitchf = self.get_labels(phone, pitch, pitchf)
+        spec, wav = self.get_audio(file)
+        dv = self.get_sid(dv)
+
+        len_phone = phone.size()[0]
+        len_spec = spec.size()[-1]
+        # print(123,phone.shape,pitch.shape,spec.shape)
+        if len_phone != len_spec:
+            len_min = min(len_phone, len_spec)
+            # amor
+            len_wav = len_min * self.hop_length
+
+            spec = spec[:, :len_min]
+            wav = wav[:, :len_wav]
+
+            phone = phone[:len_min, :]
+            pitch = pitch[:len_min]
+            pitchf = pitchf[:len_min]
+
+        return (spec, wav, phone, pitch, pitchf, dv)
+
+    def get_labels(self, phone, pitch, pitchf):
+        phone = np.load(phone)
+        phone = np.repeat(phone, 2, axis=0)
+        pitch = np.load(pitch)
+        pitchf = np.load(pitchf)
+        n_num = min(phone.shape[0], 900)  # DistributedBucketSampler
+        # print(234,phone.shape,pitch.shape)
+        phone = phone[:n_num, :]
+        pitch = pitch[:n_num]
+        pitchf = pitchf[:n_num]
+        phone = torch.FloatTensor(phone)
+        pitch = torch.LongTensor(pitch)
+        pitchf = torch.FloatTensor(pitchf)
+        return phone, pitch, pitchf
+
+    def get_audio(self, filename):
+        audio, sampling_rate = load_wav_to_torch(filename)
+        if sampling_rate != self.sampling_rate:
+            raise ValueError(
+                "{} SR doesn't match target {} SR".format(
+                    sampling_rate, self.sampling_rate
+                )
+            )
+        audio_norm = audio
+        #        audio_norm = audio / self.max_wav_value
+        #        audio_norm = audio / np.abs(audio).max()
+
+        audio_norm = audio_norm.unsqueeze(0)
+        spec_filename = filename.replace(".wav", ".spec.pt")
+        if os.path.exists(spec_filename):
+            try:
+                spec = torch.load(spec_filename)
+            except:
+                print(spec_filename, traceback.format_exc())
+                spec = spectrogram_torch(
+                    audio_norm,
+                    self.filter_length,
+                    self.sampling_rate,
+                    self.hop_length,
+                    self.win_length,
+                    center=False,
+                )
+                spec = torch.squeeze(spec, 0)
+                torch.save(spec, spec_filename, _use_new_zipfile_serialization=False)
+        else:
+            spec = spectrogram_torch(
+                audio_norm,
+                self.filter_length,
+                self.sampling_rate,
+                self.hop_length,
+                self.win_length,
+                center=False,
+            )
+            spec = torch.squeeze(spec, 0)
+            torch.save(spec, spec_filename, _use_new_zipfile_serialization=False)
+        return spec, audio_norm
+
+    def __getitem__(self, index):
+        return self.get_audio_text_pair(self.audiopaths_and_text[index])
+
+    def __len__(self):
+        return len(self.audiopaths_and_text)
+
+
+class TextAudioCollateMultiNSFsid:
+    """Zero-pads model inputs and targets"""
+
+    def __init__(self, return_ids=False):
+        self.return_ids = return_ids
+
+    def __call__(self, batch):
+        """Collate's training batch from normalized text and aduio
+        PARAMS
+        ------
+        batch: [text_normalized, spec_normalized, wav_normalized]
+        """
+        # Right zero-pad all one-hot text sequences to max input length
+        _, ids_sorted_decreasing = torch.sort(
+            torch.LongTensor([x[0].size(1) for x in batch]), dim=0, descending=True
+        )
+
+        max_spec_len = max([x[0].size(1) for x in batch])
+        max_wave_len = max([x[1].size(1) for x in batch])
+        spec_lengths = torch.LongTensor(len(batch))
+        wave_lengths = torch.LongTensor(len(batch))
+        spec_padded = torch.FloatTensor(len(batch), batch[0][0].size(0), max_spec_len)
+        wave_padded = torch.FloatTensor(len(batch), 1, max_wave_len)
+        spec_padded.zero_()
+        wave_padded.zero_()
+
+        max_phone_len = max([x[2].size(0) for x in batch])
+        phone_lengths = torch.LongTensor(len(batch))
+        phone_padded = torch.FloatTensor(
+            len(batch), max_phone_len, batch[0][2].shape[1]
+        )  # (spec, wav, phone, pitch)
+        pitch_padded = torch.LongTensor(len(batch), max_phone_len)
+        pitchf_padded = torch.FloatTensor(len(batch), max_phone_len)
+        phone_padded.zero_()
+        pitch_padded.zero_()
+        pitchf_padded.zero_()
+        # dv = torch.FloatTensor(len(batch), 256)#gin=256
+        sid = torch.LongTensor(len(batch))
+
+        for i in range(len(ids_sorted_decreasing)):
+            row = batch[ids_sorted_decreasing[i]]
+
+            spec = row[0]
+            spec_padded[i, :, : spec.size(1)] = spec
+            spec_lengths[i] = spec.size(1)
+
+            wave = row[1]
+            wave_padded[i, :, : wave.size(1)] = wave
+            wave_lengths[i] = wave.size(1)
+
+            phone = row[2]
+            phone_padded[i, : phone.size(0), :] = phone
+            phone_lengths[i] = phone.size(0)
+
+            pitch = row[3]
+            pitch_padded[i, : pitch.size(0)] = pitch
+            pitchf = row[4]
+            pitchf_padded[i, : pitchf.size(0)] = pitchf
+
+            # dv[i] = row[5]
+            sid[i] = row[5]
+
+        return (
+            phone_padded,
+            phone_lengths,
+            pitch_padded,
+            pitchf_padded,
+            spec_padded,
+            spec_lengths,
+            wave_padded,
+            wave_lengths,
+            # dv
+            sid,
+        )
+
+
+class TextAudioLoader(torch.utils.data.Dataset):
+    """
+    1) loads audio, text pairs
+    2) normalizes text and converts them to sequences of integers
+    3) computes spectrograms from audio files.
+    """
+
+    def __init__(self, audiopaths_and_text, hparams):
+        self.audiopaths_and_text = load_filepaths_and_text(audiopaths_and_text)
+        self.max_wav_value = hparams.max_wav_value
+        self.sampling_rate = hparams.sampling_rate
+        self.filter_length = hparams.filter_length
+        self.hop_length = hparams.hop_length
+        self.win_length = hparams.win_length
+        self.sampling_rate = hparams.sampling_rate
+        self.min_text_len = getattr(hparams, "min_text_len", 1)
+        self.max_text_len = getattr(hparams, "max_text_len", 5000)
+        self._filter()
+
+    def _filter(self):
+        """
+        Filter text & store spec lengths
+        """
+        # Store spectrogram lengths for Bucketing
+        # wav_length ~= file_size / (wav_channels * Bytes per dim) = file_size / (1 * 2)
+        # spec_length = wav_length // hop_length
+        audiopaths_and_text_new = []
+        lengths = []
+        for audiopath, text, dv in self.audiopaths_and_text:
+            if self.min_text_len <= len(text) and len(text) <= self.max_text_len:
+                audiopaths_and_text_new.append([audiopath, text, dv])
+                lengths.append(os.path.getsize(audiopath) // (3 * self.hop_length))
+        self.audiopaths_and_text = audiopaths_and_text_new
+        self.lengths = lengths
+
+    def get_sid(self, sid):
+        sid = torch.LongTensor([int(sid)])
+        return sid
+
+    def get_audio_text_pair(self, audiopath_and_text):
+        # separate filename and text
+        file = audiopath_and_text[0]
+        phone = audiopath_and_text[1]
+        dv = audiopath_and_text[2]
+
+        phone = self.get_labels(phone)
+        spec, wav = self.get_audio(file)
+        dv = self.get_sid(dv)
+
+        len_phone = phone.size()[0]
+        len_spec = spec.size()[-1]
+        if len_phone != len_spec:
+            len_min = min(len_phone, len_spec)
+            len_wav = len_min * self.hop_length
+            spec = spec[:, :len_min]
+            wav = wav[:, :len_wav]
+            phone = phone[:len_min, :]
+        return (spec, wav, phone, dv)
+
+    def get_labels(self, phone):
+        phone = np.load(phone)
+        phone = np.repeat(phone, 2, axis=0)
+        n_num = min(phone.shape[0], 900)  # DistributedBucketSampler
+        phone = phone[:n_num, :]
+        phone = torch.FloatTensor(phone)
+        return phone
+
+    def get_audio(self, filename):
+        audio, sampling_rate = load_wav_to_torch(filename)
+        if sampling_rate != self.sampling_rate:
+            raise ValueError(
+                "{} SR doesn't match target {} SR".format(
+                    sampling_rate, self.sampling_rate
+                )
+            )
+        audio_norm = audio
+        #        audio_norm = audio / self.max_wav_value
+        #        audio_norm = audio / np.abs(audio).max()
+
+        audio_norm = audio_norm.unsqueeze(0)
+        spec_filename = filename.replace(".wav", ".spec.pt")
+        if os.path.exists(spec_filename):
+            try:
+                spec = torch.load(spec_filename)
+            except:
+                print(spec_filename, traceback.format_exc())
+                spec = spectrogram_torch(
+                    audio_norm,
+                    self.filter_length,
+                    self.sampling_rate,
+                    self.hop_length,
+                    self.win_length,
+                    center=False,
+                )
+                spec = torch.squeeze(spec, 0)
+                torch.save(spec, spec_filename, _use_new_zipfile_serialization=False)
+        else:
+            spec = spectrogram_torch(
+                audio_norm,
+                self.filter_length,
+                self.sampling_rate,
+                self.hop_length,
+                self.win_length,
+                center=False,
+            )
+            spec = torch.squeeze(spec, 0)
+            torch.save(spec, spec_filename, _use_new_zipfile_serialization=False)
+        return spec, audio_norm
+
+    def __getitem__(self, index):
+        return self.get_audio_text_pair(self.audiopaths_and_text[index])
+
+    def __len__(self):
+        return len(self.audiopaths_and_text)
+
+
+class TextAudioCollate:
+    """Zero-pads model inputs and targets"""
+
+    def __init__(self, return_ids=False):
+        self.return_ids = return_ids
+
+    def __call__(self, batch):
+        """Collate's training batch from normalized text and aduio
+        PARAMS
+        ------
+        batch: [text_normalized, spec_normalized, wav_normalized]
+        """
+        # Right zero-pad all one-hot text sequences to max input length
+        _, ids_sorted_decreasing = torch.sort(
+            torch.LongTensor([x[0].size(1) for x in batch]), dim=0, descending=True
+        )
+
+        max_spec_len = max([x[0].size(1) for x in batch])
+        max_wave_len = max([x[1].size(1) for x in batch])
+        spec_lengths = torch.LongTensor(len(batch))
+        wave_lengths = torch.LongTensor(len(batch))
+        spec_padded = torch.FloatTensor(len(batch), batch[0][0].size(0), max_spec_len)
+        wave_padded = torch.FloatTensor(len(batch), 1, max_wave_len)
+        spec_padded.zero_()
+        wave_padded.zero_()
+
+        max_phone_len = max([x[2].size(0) for x in batch])
+        phone_lengths = torch.LongTensor(len(batch))
+        phone_padded = torch.FloatTensor(
+            len(batch), max_phone_len, batch[0][2].shape[1]
+        )
+        phone_padded.zero_()
+        sid = torch.LongTensor(len(batch))
+
+        for i in range(len(ids_sorted_decreasing)):
+            row = batch[ids_sorted_decreasing[i]]
+
+            spec = row[0]
+            spec_padded[i, :, : spec.size(1)] = spec
+            spec_lengths[i] = spec.size(1)
+
+            wave = row[1]
+            wave_padded[i, :, : wave.size(1)] = wave
+            wave_lengths[i] = wave.size(1)
+
+            phone = row[2]
+            phone_padded[i, : phone.size(0), :] = phone
+            phone_lengths[i] = phone.size(0)
+
+            sid[i] = row[3]
+
+        return (
+            phone_padded,
+            phone_lengths,
+            spec_padded,
+            spec_lengths,
+            wave_padded,
+            wave_lengths,
+            sid,
+        )
+
+
+class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
+    """
+    Maintain similar input lengths in a batch.
+    Length groups are specified by boundaries.
+    Ex) boundaries = [b1, b2, b3] -> any batch is included either {x | b1 < length(x) <=b2} or {x | b2 < length(x) <= b3}.
+
+    It removes samples which are not included in the boundaries.
+    Ex) boundaries = [b1, b2, b3] -> any x s.t. length(x) <= b1 or length(x) > b3 are discarded.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        boundaries,
+        num_replicas=None,
+        rank=None,
+        shuffle=True,
+    ):
+        super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)
+        self.lengths = dataset.lengths
+        self.batch_size = batch_size
+        self.boundaries = boundaries
+
+        self.buckets, self.num_samples_per_bucket = self._create_buckets()
+        self.total_size = sum(self.num_samples_per_bucket)
+        self.num_samples = self.total_size // self.num_replicas
+
+    def _create_buckets(self):
+        buckets = [[] for _ in range(len(self.boundaries) - 1)]
+        for i in range(len(self.lengths)):
+            length = self.lengths[i]
+            idx_bucket = self._bisect(length)
+            if idx_bucket != -1:
+                buckets[idx_bucket].append(i)
+
+        for i in range(len(buckets) - 1, -1, -1):  #
+            if len(buckets[i]) == 0:
+                buckets.pop(i)
+                self.boundaries.pop(i + 1)
+
+        num_samples_per_bucket = []
+        for i in range(len(buckets)):
+            len_bucket = len(buckets[i])
+            total_batch_size = self.num_replicas * self.batch_size
+            rem = (
+                total_batch_size - (len_bucket % total_batch_size)
+            ) % total_batch_size
+            num_samples_per_bucket.append(len_bucket + rem)
+        return buckets, num_samples_per_bucket
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+
+        indices = []
+        if self.shuffle:
+            for bucket in self.buckets:
+                indices.append(torch.randperm(len(bucket), generator=g).tolist())
+        else:
+            for bucket in self.buckets:
+                indices.append(list(range(len(bucket))))
+
+        batches = []
+        for i in range(len(self.buckets)):
+            bucket = self.buckets[i]
+            len_bucket = len(bucket)
+            ids_bucket = indices[i]
+            num_samples_bucket = self.num_samples_per_bucket[i]
+
+            # add extra samples to make it evenly divisible
+            rem = num_samples_bucket - len_bucket
+            ids_bucket = (
+                ids_bucket
+                + ids_bucket * (rem // len_bucket)
+                + ids_bucket[: (rem % len_bucket)]
+            )
+
+            # subsample
+            ids_bucket = ids_bucket[self.rank :: self.num_replicas]
+
+            # batching
+            for j in range(len(ids_bucket) // self.batch_size):
+                batch = [
+                    bucket[idx]
+                    for idx in ids_bucket[
+                        j * self.batch_size : (j + 1) * self.batch_size
+                    ]
+                ]
+                batches.append(batch)
+
+        if self.shuffle:
+            batch_ids = torch.randperm(len(batches), generator=g).tolist()
+            batches = [batches[i] for i in batch_ids]
+        self.batches = batches
+
+        assert len(self.batches) * self.batch_size == self.num_samples
+        return iter(self.batches)
+
+    def _bisect(self, x, lo=0, hi=None):
+        if hi is None:
+            hi = len(self.boundaries) - 1
+
+        if hi > lo:
+            mid = (hi + lo) // 2
+            if self.boundaries[mid] < x and x <= self.boundaries[mid + 1]:
+                return mid
+            elif x <= self.boundaries[mid]:
+                return self._bisect(x, lo, mid)
+            else:
+                return self._bisect(x, mid + 1, hi)
+        else:
+            return -1
+
+    def __len__(self):
+        return self.num_samples // self.batch_size

lib/train/losses.py

@@ -0,0 +1,58 @@
+import torch
+
+
+def feature_loss(fmap_r, fmap_g):
+    loss = 0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            rl = rl.float().detach()
+            gl = gl.float()
+            loss += torch.mean(torch.abs(rl - gl))
+
+    return loss * 2
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    r_losses = []
+    g_losses = []
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        dr = dr.float()
+        dg = dg.float()
+        r_loss = torch.mean((1 - dr) ** 2)
+        g_loss = torch.mean(dg**2)
+        loss += r_loss + g_loss
+        r_losses.append(r_loss.item())
+        g_losses.append(g_loss.item())
+
+    return loss, r_losses, g_losses
+
+
+def generator_loss(disc_outputs):
+    loss = 0
+    gen_losses = []
+    for dg in disc_outputs:
+        dg = dg.float()
+        l = torch.mean((1 - dg) ** 2)
+        gen_losses.append(l)
+        loss += l
+
+    return loss, gen_losses
+
+
+def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
+    """
+    z_p, logs_q: [b, h, t_t]
+    m_p, logs_p: [b, h, t_t]
+    """
+    z_p = z_p.float()
+    logs_q = logs_q.float()
+    m_p = m_p.float()
+    logs_p = logs_p.float()
+    z_mask = z_mask.float()
+
+    kl = logs_p - logs_q - 0.5
+    kl += 0.5 * ((z_p - m_p) ** 2) * torch.exp(-2.0 * logs_p)
+    kl = torch.sum(kl * z_mask)
+    l = kl / torch.sum(z_mask)
+    return l

lib/train/mel_processing.py

@@ -0,0 +1,130 @@
+import torch
+import torch.utils.data
+from librosa.filters import mel as librosa_mel_fn
+
+
+MAX_WAV_VALUE = 32768.0
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    """
+    PARAMS
+    ------
+    C: compression factor
+    """
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    """
+    PARAMS
+    ------
+    C: compression factor used to compress
+    """
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    return dynamic_range_compression_torch(magnitudes)
+
+
+def spectral_de_normalize_torch(magnitudes):
+    return dynamic_range_decompression_torch(magnitudes)
+
+
+# Reusable banks
+mel_basis = {}
+hann_window = {}
+
+
+def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
+    """Convert waveform into Linear-frequency Linear-amplitude spectrogram.
+
+    Args:
+        y             :: (B, T) - Audio waveforms
+        n_fft
+        sampling_rate
+        hop_size
+        win_size
+        center
+    Returns:
+        :: (B, Freq, Frame) - Linear-frequency Linear-amplitude spectrogram
+    """
+    # Validation
+    if torch.min(y) < -1.07:
+        print("min value is ", torch.min(y))
+    if torch.max(y) > 1.07:
+        print("max value is ", torch.max(y))
+
+    # Window - Cache if needed
+    global hann_window
+    dtype_device = str(y.dtype) + "_" + str(y.device)
+    wnsize_dtype_device = str(win_size) + "_" + dtype_device
+    if wnsize_dtype_device not in hann_window:
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
+            dtype=y.dtype, device=y.device
+        )
+
+    # Padding
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
+    y = y.squeeze(1)
+
+    # Complex Spectrogram :: (B, T) -> (B, Freq, Frame, RealComplex=2)
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window[wnsize_dtype_device],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=False,
+    )
+
+    # Linear-frequency Linear-amplitude spectrogram :: (B, Freq, Frame, RealComplex=2) -> (B, Freq, Frame)
+    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
+    return spec
+
+
+def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
+    # MelBasis - Cache if needed
+    global mel_basis
+    dtype_device = str(spec.dtype) + "_" + str(spec.device)
+    fmax_dtype_device = str(fmax) + "_" + dtype_device
+    if fmax_dtype_device not in mel_basis:
+        mel = librosa_mel_fn(
+            sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
+        )
+        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
+            dtype=spec.dtype, device=spec.device
+        )
+
+    # Mel-frequency Log-amplitude spectrogram :: (B, Freq=num_mels, Frame)
+    melspec = torch.matmul(mel_basis[fmax_dtype_device], spec)
+    melspec = spectral_normalize_torch(melspec)
+    return melspec
+
+
+def mel_spectrogram_torch(
+    y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
+):
+    """Convert waveform into Mel-frequency Log-amplitude spectrogram.
+
+    Args:
+        y       :: (B, T)           - Waveforms
+    Returns:
+        melspec :: (B, Freq, Frame) - Mel-frequency Log-amplitude spectrogram
+    """
+    # Linear-frequency Linear-amplitude spectrogram :: (B, T) -> (B, Freq, Frame)
+    spec = spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center)
+
+    # Mel-frequency Log-amplitude spectrogram :: (B, Freq, Frame) -> (B, Freq=num_mels, Frame)
+    melspec = spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax)
+
+    return melspec

lib/train/process_ckpt.py

@@ -0,0 +1,217 @@
+import torch, traceback, os
+from collections import OrderedDict
+from .. import get_cwd
+from ...i18n import i18n
+CWD = get_cwd()
+
+def show_info(path):
+    try:
+        a = torch.load(path, map_location="cpu")
+        return "模型信息:%s\n采样率:%s\n模型是否输入音高引导:%s\n版本:%s" % (
+            a.get("info", "None"),
+            a.get("sr", "None"),
+            a.get("f0", "None"),
+            a.get("version", "None"),
+        )
+    except:
+        return traceback.format_exc()
+
+
+def extract_small_model(path, name, sr, if_f0, info, version):
+    try:
+        ckpt = torch.load(path, map_location="cpu")
+        if "model" in ckpt:
+            ckpt = ckpt["model"]
+        opt = OrderedDict()
+        opt["weight"] = {}
+        for key in ckpt.keys():
+            if "enc_q" in key:
+                continue
+            opt["weight"][key] = ckpt[key].half()
+        if sr == "40k":
+            opt["config"] = [
+                1025,
+                32,
+                192,
+                192,
+                768,
+                2,
+                6,
+                3,
+                0,
+                "1",
+                [3, 7, 11],
+                [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                [10, 10, 2, 2],
+                512,
+                [16, 16, 4, 4],
+                109,
+                256,
+                40000,
+            ]
+        elif sr == "48k":
+            if version == "v1":
+                opt["config"] = [
+                    1025,
+                    32,
+                    192,
+                    192,
+                    768,
+                    2,
+                    6,
+                    3,
+                    0,
+                    "1",
+                    [3, 7, 11],
+                    [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                    [10, 6, 2, 2, 2],
+                    512,
+                    [16, 16, 4, 4, 4],
+                    109,
+                    256,
+                    48000,
+                ]
+            else:
+                opt["config"] = [
+                    1025,
+                    32,
+                    192,
+                    192,
+                    768,
+                    2,
+                    6,
+                    3,
+                    0,
+                    "1",
+                    [3, 7, 11],
+                    [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                    [12, 10, 2, 2],
+                    512,
+                    [24, 20, 4, 4],
+                    109,
+                    256,
+                    48000,
+                ]
+        elif sr == "32k":
+            if version == "v1":
+                opt["config"] = [
+                    513,
+                    32,
+                    192,
+                    192,
+                    768,
+                    2,
+                    6,
+                    3,
+                    0,
+                    "1",
+                    [3, 7, 11],
+                    [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                    [10, 4, 2, 2, 2],
+                    512,
+                    [16, 16, 4, 4, 4],
+                    109,
+                    256,
+                    32000,
+                ]
+            else:
+                opt["config"] = [
+                    513,
+                    32,
+                    192,
+                    192,
+                    768,
+                    2,
+                    6,
+                    3,
+                    0,
+                    "1",
+                    [3, 7, 11],
+                    [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                    [10, 8, 2, 2],
+                    512,
+                    [20, 16, 4, 4],
+                    109,
+                    256,
+                    32000,
+                ]
+        if info == "":
+            info = "Extracted model."
+        opt["info"] = info
+        opt["version"] = version
+        opt["sr"] = sr
+        opt["f0"] = int(if_f0)
+        torch.save(opt, "weights/%s.pth" % name)
+        return "Success."
+    except:
+        return traceback.format_exc()
+
+
+def change_info(path, info, name):
+    try:
+        ckpt = torch.load(path, map_location="cpu")
+        ckpt["info"] = info
+        if name == "":
+            name = os.path.basename(path)
+        torch.save(ckpt, "weights/%s" % name)
+        return "Success."
+    except:
+        return traceback.format_exc()
+
+
+def merge(path1, path2, alpha1, sr, f0, info, name, version):
+    try:
+
+        def extract(ckpt):
+            a = ckpt["model"]
+            opt = OrderedDict()
+            opt["weight"] = {}
+            for key in a.keys():
+                if "enc_q" in key:
+                    continue
+                opt["weight"][key] = a[key]
+            return opt
+
+        ckpt1 = torch.load(path1, map_location="cpu")
+        ckpt2 = torch.load(path2, map_location="cpu")
+        cfg = ckpt1["config"]
+        if "model" in ckpt1:
+            ckpt1 = extract(ckpt1)
+        else:
+            ckpt1 = ckpt1["weight"]
+        if "model" in ckpt2:
+            ckpt2 = extract(ckpt2)
+        else:
+            ckpt2 = ckpt2["weight"]
+        if sorted(list(ckpt1.keys())) != sorted(list(ckpt2.keys())):
+            return "Fail to merge the models. The model architectures are not the same."
+        opt = OrderedDict()
+        opt["weight"] = {}
+        for key in ckpt1.keys():
+            # try:
+            if key == "emb_g.weight" and ckpt1[key].shape != ckpt2[key].shape:
+                min_shape0 = min(ckpt1[key].shape[0], ckpt2[key].shape[0])
+                opt["weight"][key] = (
+                    alpha1 * (ckpt1[key][:min_shape0].float())
+                    + (1 - alpha1) * (ckpt2[key][:min_shape0].float())
+                ).half()
+            else:
+                opt["weight"][key] = (
+                    alpha1 * (ckpt1[key].float()) + (1 - alpha1) * (ckpt2[key].float())
+                ).half()
+        # except:
+        #     pdb.set_trace()
+        opt["config"] = cfg
+        """
+        if(sr=="40k"):opt["config"] = [1025, 32, 192, 192, 768, 2, 6, 3, 0, "1", [3, 7, 11], [[1, 3, 5], [1, 3, 5], [1, 3, 5]], [10, 10, 2, 2], 512, [16, 16, 4, 4,4], 109, 256, 40000]
+        elif(sr=="48k"):opt["config"] = [1025, 32, 192, 192, 768, 2, 6, 3, 0, "1", [3, 7, 11], [[1, 3, 5], [1, 3, 5], [1, 3, 5]], [10,6,2,2,2], 512, [16, 16, 4, 4], 109, 256, 48000]
+        elif(sr=="32k"):opt["config"] = [513, 32, 192, 192, 768, 2, 6, 3, 0, "1", [3, 7, 11], [[1, 3, 5], [1, 3, 5], [1, 3, 5]], [10, 4, 2, 2, 2], 512, [16, 16, 4, 4,4], 109, 256, 32000]
+        """
+        opt["sr"] = sr
+        opt["f0"] = 1 if f0 == i18n("是") else 0
+        opt["version"] = version
+        opt["info"] = info
+        torch.save(opt, "weights/%s.pth" % name)
+        return "Success."
+    except:
+        return traceback.format_exc()

lib/train/utils.py

@@ -0,0 +1,491 @@
+import os, traceback
+import glob
+import sys
+import argparse
+import logging
+import json
+import subprocess
+import numpy as np
+from scipy.io.wavfile import read
+import torch
+
+MATPLOTLIB_FLAG = False
+
+logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+logger = logging
+
+
+def load_checkpoint_d(checkpoint_path, combd, sbd, optimizer=None, load_opt=1):
+    assert os.path.isfile(checkpoint_path)
+    checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
+
+    ##################
+    def go(model, bkey):
+        saved_state_dict = checkpoint_dict[bkey]
+        if hasattr(model, "module"):
+            state_dict = model.module.state_dict()
+        else:
+            state_dict = model.state_dict()
+        new_state_dict = {}
+        for k, v in state_dict.items():  # 模型需要的shape
+            try:
+                new_state_dict[k] = saved_state_dict[k]
+                if saved_state_dict[k].shape != state_dict[k].shape:
+                    print(
+                        "shape-%s-mismatch|need-%s|get-%s"
+                        % (k, state_dict[k].shape, saved_state_dict[k].shape)
+                    )  #
+                    raise KeyError
+            except:
+                # logger.info(traceback.format_exc())
+                logger.info("%s is not in the checkpoint" % k)  # pretrain缺失的
+                new_state_dict[k] = v  # 模型自带的随机值
+        if hasattr(model, "module"):
+            model.module.load_state_dict(new_state_dict, strict=False)
+        else:
+            model.load_state_dict(new_state_dict, strict=False)
+        return model
+
+    go(combd, "combd")
+    model = go(sbd, "sbd")
+    #############
+    logger.info("Loaded model weights")
+
+    iteration = checkpoint_dict["iteration"]
+    learning_rate = checkpoint_dict["learning_rate"]
+    if (
+        optimizer is not None and load_opt == 1
+    ):  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
+        #   try:
+        optimizer.load_state_dict(checkpoint_dict["optimizer"])
+    #   except:
+    #     traceback.print_exc()
+    logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration))
+    return model, optimizer, learning_rate, iteration
+
+
+# def load_checkpoint(checkpoint_path, model, optimizer=None):
+#   assert os.path.isfile(checkpoint_path)
+#   checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
+#   iteration = checkpoint_dict['iteration']
+#   learning_rate = checkpoint_dict['learning_rate']
+#   if optimizer is not None:
+#     optimizer.load_state_dict(checkpoint_dict['optimizer'])
+#   # print(1111)
+#   saved_state_dict = checkpoint_dict['model']
+#   # print(1111)
+#
+#   if hasattr(model, 'module'):
+#     state_dict = model.module.state_dict()
+#   else:
+#     state_dict = model.state_dict()
+#   new_state_dict= {}
+#   for k, v in state_dict.items():
+#     try:
+#       new_state_dict[k] = saved_state_dict[k]
+#     except:
+#       logger.info("%s is not in the checkpoint" % k)
+#       new_state_dict[k] = v
+#   if hasattr(model, 'module'):
+#     model.module.load_state_dict(new_state_dict)
+#   else:
+#     model.load_state_dict(new_state_dict)
+#   logger.info("Loaded checkpoint '{}' (epoch {})" .format(
+#     checkpoint_path, iteration))
+#   return model, optimizer, learning_rate, iteration
+def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
+    assert os.path.isfile(checkpoint_path)
+    checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
+
+    saved_state_dict = checkpoint_dict["model"]
+    if hasattr(model, "module"):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+    new_state_dict = {}
+    for k, v in state_dict.items():  # 模型需要的shape
+        try:
+            new_state_dict[k] = saved_state_dict[k]
+            if saved_state_dict[k].shape != state_dict[k].shape:
+                print(
+                    "shape-%s-mismatch|need-%s|get-%s"
+                    % (k, state_dict[k].shape, saved_state_dict[k].shape)
+                )  #
+                raise KeyError
+        except:
+            # logger.info(traceback.format_exc())
+            logger.info("%s is not in the checkpoint" % k)  # pretrain缺失的
+            new_state_dict[k] = v  # 模型自带的随机值
+    if hasattr(model, "module"):
+        model.module.load_state_dict(new_state_dict, strict=False)
+    else:
+        model.load_state_dict(new_state_dict, strict=False)
+    logger.info("Loaded model weights")
+
+    iteration = checkpoint_dict["iteration"]
+    learning_rate = checkpoint_dict["learning_rate"]
+    if (
+        optimizer is not None and load_opt == 1
+    ):  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
+        #   try:
+        optimizer.load_state_dict(checkpoint_dict["optimizer"])
+    #   except:
+    #     traceback.print_exc()
+    logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration))
+    return model, optimizer, learning_rate, iteration
+
+
+def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
+    logger.info(
+        "Saving model and optimizer state at epoch {} to {}".format(
+            iteration, checkpoint_path
+        )
+    )
+    if hasattr(model, "module"):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+    torch.save(
+        {
+            "model": state_dict,
+            "iteration": iteration,
+            "optimizer": optimizer.state_dict(),
+            "learning_rate": learning_rate,
+        },
+        checkpoint_path,
+    )
+
+
+def save_checkpoint_d(combd, sbd, optimizer, learning_rate, iteration, checkpoint_path):
+    logger.info(
+        "Saving model and optimizer state at epoch {} to {}".format(
+            iteration, checkpoint_path
+        )
+    )
+    if hasattr(combd, "module"):
+        state_dict_combd = combd.module.state_dict()
+    else:
+        state_dict_combd = combd.state_dict()
+    if hasattr(sbd, "module"):
+        state_dict_sbd = sbd.module.state_dict()
+    else:
+        state_dict_sbd = sbd.state_dict()
+    torch.save(
+        {
+            "combd": state_dict_combd,
+            "sbd": state_dict_sbd,
+            "iteration": iteration,
+            "optimizer": optimizer.state_dict(),
+            "learning_rate": learning_rate,
+        },
+        checkpoint_path,
+    )
+
+
+def summarize(
+    writer,
+    global_step,
+    scalars={},
+    histograms={},
+    images={},
+    audios={},
+    audio_sampling_rate=22050,
+):
+    for k, v in scalars.items():
+        writer.add_scalar(k, v, global_step)
+    for k, v in histograms.items():
+        writer.add_histogram(k, v, global_step)
+    for k, v in images.items():
+        writer.add_image(k, v, global_step, dataformats="HWC")
+    for k, v in audios.items():
+        writer.add_audio(k, v, global_step, audio_sampling_rate)
+
+
+def latest_checkpoint_path(dir_path, regex="G_*.pth"):
+    f_list = glob.glob(os.path.join(dir_path, regex))
+    f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
+    x = f_list[-1]
+    print(x)
+    return x
+
+
+def plot_spectrogram_to_numpy(spectrogram):
+    global MATPLOTLIB_FLAG
+    if not MATPLOTLIB_FLAG:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        MATPLOTLIB_FLAG = True
+        mpl_logger = logging.getLogger("matplotlib")
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+    import numpy as np
+
+    fig, ax = plt.subplots(figsize=(10, 2))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+    plt.colorbar(im, ax=ax)
+    plt.xlabel("Frames")
+    plt.ylabel("Channels")
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def plot_alignment_to_numpy(alignment, info=None):
+    global MATPLOTLIB_FLAG
+    if not MATPLOTLIB_FLAG:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        MATPLOTLIB_FLAG = True
+        mpl_logger = logging.getLogger("matplotlib")
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+    import numpy as np
+
+    fig, ax = plt.subplots(figsize=(6, 4))
+    im = ax.imshow(
+        alignment.transpose(), aspect="auto", origin="lower", interpolation="none"
+    )
+    fig.colorbar(im, ax=ax)
+    xlabel = "Decoder timestep"
+    if info is not None:
+        xlabel += "\n\n" + info
+    plt.xlabel(xlabel)
+    plt.ylabel("Encoder timestep")
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def load_wav_to_torch(full_path):
+    sampling_rate, data = read(full_path)
+    return torch.FloatTensor(data.astype(np.float32)), sampling_rate
+
+
+def load_filepaths_and_text(filename, split="|"):
+    with open(filename, encoding="utf-8") as f:
+        filepaths_and_text = [line.strip().split(split) for line in f]
+    return filepaths_and_text
+
+
+def get_hparams(init=True):
+    """
+    todo:
+      结尾七人组：
+        保存频率、总epoch                     done
+        bs                                    done
+        pretrainG、pretrainD                  done
+        卡号：os.en["CUDA_VISIBLE_DEVICES"]   done
+        if_latest                             done
+      模型：if_f0                             done
+      采样率：自动选择config                  done
+      是否缓存数据集进GPU:if_cache_data_in_gpu done
+
+      -m:
+        自动决定training_files路径,改掉train_nsf_load_pretrain.py里的hps.data.training_files    done
+      -c不要了
+    """
+    parser = argparse.ArgumentParser()
+    # parser.add_argument('-c', '--config', type=str, default="configs/40k.json",help='JSON file for configuration')
+    parser.add_argument(
+        "-se",
+        "--save_every_epoch",
+        type=int,
+        required=True,
+        help="checkpoint save frequency (epoch)",
+    )
+    parser.add_argument(
+        "-te", "--total_epoch", type=int, required=True, help="total_epoch"
+    )
+    parser.add_argument(
+        "-pg", "--pretrainG", type=str, default="", help="Pretrained Discriminator path"
+    )
+    parser.add_argument(
+        "-pd", "--pretrainD", type=str, default="", help="Pretrained Generator path"
+    )
+    parser.add_argument("-g", "--gpus", type=str, default="0", help="split by -")
+    parser.add_argument(
+        "-bs", "--batch_size", type=int, required=True, help="batch size"
+    )
+    parser.add_argument(
+        "-e", "--experiment_dir", type=str, required=True, help="experiment dir"
+    )
+    parser.add_argument(
+        "-n", "--name", type=str, help="name of the model, uses experiment_dir as default"
+    )
+    parser.add_argument(
+        "-sr", "--sample_rate", type=str, required=True, help="sample rate, 32k/40k/48k"
+    )
+    parser.add_argument(
+        "-sw",
+        "--save_every_weights",
+        type=str,
+        default="0",
+        help="save the extracted model in weights directory when saving checkpoints",
+    )
+    parser.add_argument(
+        "-v", "--version", type=str, required=True, help="model version"
+    )
+    parser.add_argument(
+        "-f0",
+        "--if_f0",
+        type=int,
+        required=True,
+        help="use f0 as one of the inputs of the model, 1 or 0",
+    )
+    parser.add_argument(
+        "-l",
+        "--if_latest",
+        type=int,
+        required=True,
+        help="if only save the latest G/D pth file, 1 or 0",
+    )
+    parser.add_argument(
+        "-c",
+        "--if_cache_data_in_gpu",
+        type=int,
+        required=True,
+        help="if caching the dataset in GPU memory, 1 or 0",
+    )
+
+
+    args = parser.parse_args()
+    name = args.name or args.experiment_dir
+    experiment_dir = os.path.join("./logs", args.experiment_dir)
+
+    if not os.path.exists(experiment_dir):
+        os.makedirs(experiment_dir)
+
+    if args.version == "v1" or args.sample_rate == "40k":
+        config_path = "configs/%s.json" % args.sample_rate
+    else:
+        config_path = "configs/%s_v2.json" % args.sample_rate
+    config_save_path = os.path.join(experiment_dir, "config.json")
+    if init:
+        with open(config_path, "r") as f:
+            data = f.read()
+        with open(config_save_path, "w") as f:
+            f.write(data)
+    else:
+        with open(config_save_path, "r") as f:
+            data = f.read()
+    config = json.loads(data)
+
+    hparams = HParams(**config)
+    hparams.model_dir = hparams.experiment_dir = experiment_dir
+    hparams.save_every_epoch = args.save_every_epoch
+    hparams.name = name
+    hparams.total_epoch = args.total_epoch
+    hparams.pretrainG = args.pretrainG
+    hparams.pretrainD = args.pretrainD
+    hparams.version = args.version
+    hparams.gpus = args.gpus
+    hparams.train.batch_size = args.batch_size
+    hparams.sample_rate = args.sample_rate
+    hparams.if_f0 = args.if_f0
+    hparams.if_latest = args.if_latest
+    hparams.save_every_weights = args.save_every_weights
+    hparams.if_cache_data_in_gpu = args.if_cache_data_in_gpu
+    hparams.data.training_files = os.path.join(experiment_dir,"filelist.txt")
+    return hparams
+
+
+def get_hparams_from_dir(model_dir):
+    config_save_path = os.path.join(model_dir, "config.json")
+    with open(config_save_path, "r") as f:
+        data = f.read()
+    config = json.loads(data)
+
+    hparams = HParams(**config)
+    hparams.model_dir = model_dir
+    return hparams
+
+
+def get_hparams_from_file(config_path):
+    with open(config_path, "r") as f:
+        data = f.read()
+    config = json.loads(data)
+
+    hparams = HParams(**config)
+    return hparams
+
+
+def check_git_hash(model_dir):
+    source_dir = os.path.dirname(os.path.realpath(__file__))
+    if not os.path.exists(os.path.join(source_dir, ".git")):
+        logger.warn(
+            "{} is not a git repository, therefore hash value comparison will be ignored.".format(
+                source_dir
+            )
+        )
+        return
+
+    cur_hash = subprocess.getoutput("git rev-parse HEAD")
+
+    path = os.path.join(model_dir, "githash")
+    if os.path.exists(path):
+        saved_hash = open(path).read()
+        if saved_hash != cur_hash:
+            logger.warn(
+                "git hash values are different. {}(saved) != {}(current)".format(
+                    saved_hash[:8], cur_hash[:8]
+                )
+            )
+    else:
+        open(path, "w").write(cur_hash)
+
+
+def get_logger(model_dir, filename="train.log"):
+    global logger
+    logger = logging.getLogger(os.path.basename(model_dir))
+    logger.setLevel(logging.DEBUG)
+
+    formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+    h = logging.FileHandler(os.path.join(model_dir, filename))
+    h.setLevel(logging.DEBUG)
+    h.setFormatter(formatter)
+    logger.addHandler(h)
+    return logger
+
+
+class HParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if type(v) == dict:
+                v = HParams(**v)
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()

lib/utils.py

@@ -0,0 +1,108 @@
+from datetime import datetime
+import gc
+import glob
+import hashlib
+# from importlib.util import LazyLoader, find_spec, module_from_spec
+import multiprocessing
+import os
+import platform
+from time import sleep
+import numpy as np
+import psutil
+import requests
+import torch
+from . import ObjectNamespace, get_cwd
+
+torch.manual_seed(1337)
+CWD = get_cwd()
+
+def get_hash(*args):
+    return hashlib.md5("".join(str(data) for data in args).encode()).hexdigest()
+
+def get_file_hash(fname,size=10*1024**2):
+    if hasattr(fname,"read"):
+        with open(fname,"rb") as f: return hashlib.md5(f.read(-size)).hexdigest()
+
+def get_subprocesses(pid = os.getpid()):
+    # Get a list of all subprocesses started by the current process
+    subprocesses = psutil.Process(pid).children(recursive=True)
+    python_processes = [p for p in subprocesses if p.status()=="running"]
+    for p in python_processes:
+        cpu_percent = p.cpu_percent()
+        memory_percent = p.memory_percent()
+        process = ObjectNamespace(**{
+            'pid': p.pid,
+            "name": p.name(),
+            'cpu_percent': f"{cpu_percent:.2f}%",
+            'memory_percent': f"{memory_percent:.2f}%",
+            'status': p.status(),
+            'time_started': datetime.fromtimestamp(p.create_time()).isoformat(),
+            'kill': p.kill
+            })
+        yield process
+
+def get_filenames(root=CWD,folder="**",exts=["*"],name_filters=[""],filter_func=any,format_func=lambda x:x):
+    fnames = []
+    for ext in exts:
+        fnames.extend(glob.glob(f"{root}/{folder}/*.{ext}",recursive=True))
+    return sorted([format_func(ele) for ele in fnames if filter_func([nf.lower() in os.path.basename(ele).lower() for nf in name_filters])])
+
+def get_index(arr,value):
+    if arr is not None:
+        if value in arr: return arr.index(value)
+        elif value is not None:
+            for i,item in enumerate(arr):
+                k1, k2 = str(item), str(value)
+                if (k1 in k2) or (k2 in k1): return i
+    return 0
+
+def gc_collect():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+    elif torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+    gc.set_threshold(100,10,1)
+    gc.collect()
+    
+def get_optimal_torch_device(index = 0) -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device(
+            f"cuda:{index % torch.cuda.device_count()}"
+        )  # Very fast
+    elif torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+def get_optimal_threads(offset=0):
+    cores = multiprocessing.cpu_count() - offset
+    return int(max(np.floor(cores * (1-psutil.cpu_percent())),1))
+
+def pid_is_active(pid: int):        
+    """ Check For the existence of a unix pid. https://stackoverflow.com/a/568285"""
+    try:
+        if platform.system() == "Windows":
+            return psutil.pid_exists(pid)
+        elif platform.system() == "Linux":
+            os.kill(pid, 0)
+    except Exception as e:
+        print(e)
+        return False
+    else:
+        return True
+    
+def poll_url(url,timeout=10):
+    for i in range(timeout): # wait for server to start up
+        try:
+            with requests.get(url) as req:
+                if req.status_code==200: return True
+        except Exception:
+            sleep(1.)
+            print(f"waited {i+1} seconds...")
+    return False
+
+def get_merge_func(merge_type: str):
+    if merge_type=="min": return np.nanmin
+    elif merge_type=="max": return np.nanmax
+    elif merge_type=="median": return np.nanmedian
+    else: return np.nanmean

lib/uvr5_pack/UVR.py

@@ -0,0 +1,272 @@
+class ModelData():
+    def __init__(self, model_name: str, 
+                 selected_process_method=ENSEMBLE_MODE, 
+                 is_secondary_model=False, 
+                 primary_model_primary_stem=None, 
+                 is_primary_model_primary_stem_only=False, 
+                 is_primary_model_secondary_stem_only=False, 
+                 is_pre_proc_model=False,
+                 is_dry_check=False):
+
+        self.is_gpu_conversion = 0 if root.is_gpu_conversion_var.get() else -1
+        self.is_normalization = root.is_normalization_var.get()
+        self.is_primary_stem_only = root.is_primary_stem_only_var.get()
+        self.is_secondary_stem_only = root.is_secondary_stem_only_var.get()
+        self.is_denoise = root.is_denoise_var.get()
+        self.mdx_batch_size = 1 if root.mdx_batch_size_var.get() == DEF_OPT else int(root.mdx_batch_size_var.get())
+        self.is_mdx_ckpt = False
+        self.wav_type_set = root.wav_type_set
+        self.mp3_bit_set = root.mp3_bit_set_var.get()
+        self.save_format = root.save_format_var.get()
+        self.is_invert_spec = root.is_invert_spec_var.get()
+        self.is_mixer_mode = root.is_mixer_mode_var.get()
+        self.demucs_stems = root.demucs_stems_var.get()
+        self.demucs_source_list = []
+        self.demucs_stem_count = 0
+        self.mixer_path = MDX_MIXER_PATH
+        self.model_name = model_name
+        self.process_method = selected_process_method
+        self.model_status = False if self.model_name == CHOOSE_MODEL or self.model_name == NO_MODEL else True
+        self.primary_stem = None
+        self.secondary_stem = None
+        self.is_ensemble_mode = False
+        self.ensemble_primary_stem = None
+        self.ensemble_secondary_stem = None
+        self.primary_model_primary_stem = primary_model_primary_stem
+        self.is_secondary_model = is_secondary_model
+        self.secondary_model = None
+        self.secondary_model_scale = None
+        self.demucs_4_stem_added_count = 0
+        self.is_demucs_4_stem_secondaries = False
+        self.is_4_stem_ensemble = False
+        self.pre_proc_model = None
+        self.pre_proc_model_activated = False
+        self.is_pre_proc_model = is_pre_proc_model
+        self.is_dry_check = is_dry_check
+        self.model_samplerate = 44100
+        self.model_capacity = 32, 128
+        self.is_vr_51_model = False
+        self.is_demucs_pre_proc_model_inst_mix = False
+        self.manual_download_Button = None
+        self.secondary_model_4_stem = []
+        self.secondary_model_4_stem_scale = []
+        self.secondary_model_4_stem_names = []
+        self.secondary_model_4_stem_model_names_list = []
+        self.all_models = []
+        self.secondary_model_other = None
+        self.secondary_model_scale_other = None
+        self.secondary_model_bass = None
+        self.secondary_model_scale_bass = None
+        self.secondary_model_drums = None
+        self.secondary_model_scale_drums = None
+
+        if selected_process_method == ENSEMBLE_MODE:
+            partitioned_name = model_name.partition(ENSEMBLE_PARTITION)
+            self.process_method = partitioned_name[0]
+            self.model_name = partitioned_name[2]
+            self.model_and_process_tag = model_name
+            self.ensemble_primary_stem, self.ensemble_secondary_stem = root.return_ensemble_stems()
+            self.is_ensemble_mode = True if not is_secondary_model and not is_pre_proc_model else False
+            self.is_4_stem_ensemble = True if root.ensemble_main_stem_var.get() == FOUR_STEM_ENSEMBLE and self.is_ensemble_mode else False
+            self.pre_proc_model_activated = root.is_demucs_pre_proc_model_activate_var.get() if not self.ensemble_primary_stem == VOCAL_STEM else False
+
+        if self.process_method == VR_ARCH_TYPE:
+            self.is_secondary_model_activated = root.vr_is_secondary_model_activate_var.get() if not self.is_secondary_model else False
+            self.aggression_setting = float(int(root.aggression_setting_var.get())/100)
+            self.is_tta = root.is_tta_var.get()
+            self.is_post_process = root.is_post_process_var.get()
+            self.window_size = int(root.window_size_var.get())
+            self.batch_size = 1 if root.batch_size_var.get() == DEF_OPT else int(root.batch_size_var.get())
+            self.crop_size = int(root.crop_size_var.get())
+            self.is_high_end_process = 'mirroring' if root.is_high_end_process_var.get() else 'None'
+            self.post_process_threshold = float(root.post_process_threshold_var.get())
+            self.model_capacity = 32, 128
+            self.model_path = os.path.join(VR_MODELS_DIR, f"{self.model_name}.pth")
+            self.get_model_hash()
+            if self.model_hash:
+                self.model_data = self.get_model_data(VR_HASH_DIR, root.vr_hash_MAPPER) if not self.model_hash == WOOD_INST_MODEL_HASH else WOOD_INST_PARAMS
+                if self.model_data:
+                    vr_model_param = os.path.join(VR_PARAM_DIR, "{}.json".format(self.model_data["vr_model_param"]))
+                    self.primary_stem = self.model_data["primary_stem"]
+                    self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
+                    self.vr_model_param = ModelParameters(vr_model_param)
+                    self.model_samplerate = self.vr_model_param.param['sr']
+                    if "nout" in self.model_data.keys() and "nout_lstm" in self.model_data.keys():
+                        self.model_capacity = self.model_data["nout"], self.model_data["nout_lstm"]
+                        self.is_vr_51_model = True
+                else:
+                    self.model_status = False
+                
+        if self.process_method == MDX_ARCH_TYPE:
+            self.is_secondary_model_activated = root.mdx_is_secondary_model_activate_var.get() if not is_secondary_model else False
+            self.margin = int(root.margin_var.get())
+            self.chunks = root.determine_auto_chunks(root.chunks_var.get(), self.is_gpu_conversion) if root.is_chunk_mdxnet_var.get() else 0
+            self.get_mdx_model_path()
+            self.get_model_hash()
+            if self.model_hash:
+                self.model_data = self.get_model_data(MDX_HASH_DIR, root.mdx_hash_MAPPER)
+                if self.model_data:
+                    self.compensate = self.model_data["compensate"] if root.compensate_var.get() == AUTO_SELECT else float(root.compensate_var.get())
+                    self.mdx_dim_f_set = self.model_data["mdx_dim_f_set"]
+                    self.mdx_dim_t_set = self.model_data["mdx_dim_t_set"]
+                    self.mdx_n_fft_scale_set = self.model_data["mdx_n_fft_scale_set"]
+                    self.primary_stem = self.model_data["primary_stem"]
+                    self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
+                else:
+                    self.model_status = False
+
+        if self.process_method == DEMUCS_ARCH_TYPE:
+            self.is_secondary_model_activated = root.demucs_is_secondary_model_activate_var.get() if not is_secondary_model else False
+            if not self.is_ensemble_mode:
+                self.pre_proc_model_activated = root.is_demucs_pre_proc_model_activate_var.get() if not root.demucs_stems_var.get() in [VOCAL_STEM, INST_STEM] else False
+            self.overlap = float(root.overlap_var.get())
+            self.margin_demucs = int(root.margin_demucs_var.get())
+            self.chunks_demucs = root.determine_auto_chunks(root.chunks_demucs_var.get(), self.is_gpu_conversion)
+            self.shifts = int(root.shifts_var.get())
+            self.is_split_mode = root.is_split_mode_var.get()
+            self.segment = root.segment_var.get()
+            self.is_chunk_demucs = root.is_chunk_demucs_var.get()
+            self.is_demucs_combine_stems = root.is_demucs_combine_stems_var.get()
+            self.is_primary_stem_only = root.is_primary_stem_only_var.get() if self.is_ensemble_mode else root.is_primary_stem_only_Demucs_var.get() 
+            self.is_secondary_stem_only = root.is_secondary_stem_only_var.get() if self.is_ensemble_mode else root.is_secondary_stem_only_Demucs_var.get()
+            self.get_demucs_model_path()
+            self.get_demucs_model_data()
+
+        self.model_basename = os.path.splitext(os.path.basename(self.model_path))[0] if self.model_status else None
+        self.pre_proc_model_activated = self.pre_proc_model_activated if not self.is_secondary_model else False
+        
+        self.is_primary_model_primary_stem_only = is_primary_model_primary_stem_only
+        self.is_primary_model_secondary_stem_only = is_primary_model_secondary_stem_only
+
+        if self.is_secondary_model_activated and self.model_status:
+            if (not self.is_ensemble_mode and root.demucs_stems_var.get() == ALL_STEMS and self.process_method == DEMUCS_ARCH_TYPE) or self.is_4_stem_ensemble:
+                for key in DEMUCS_4_SOURCE_LIST:
+                    self.secondary_model_data(key)
+                    self.secondary_model_4_stem.append(self.secondary_model)
+                    self.secondary_model_4_stem_scale.append(self.secondary_model_scale)
+                    self.secondary_model_4_stem_names.append(key)
+                self.demucs_4_stem_added_count = sum(i is not None for i in self.secondary_model_4_stem)
+                self.is_secondary_model_activated = False if all(i is None for i in self.secondary_model_4_stem) else True
+                self.demucs_4_stem_added_count = self.demucs_4_stem_added_count - 1 if self.is_secondary_model_activated else self.demucs_4_stem_added_count
+                if self.is_secondary_model_activated:
+                    self.secondary_model_4_stem_model_names_list = [None if i is None else i.model_basename for i in self.secondary_model_4_stem]
+                    self.is_demucs_4_stem_secondaries = True 
+            else:
+                primary_stem = self.ensemble_primary_stem if self.is_ensemble_mode and self.process_method == DEMUCS_ARCH_TYPE else self.primary_stem
+                self.secondary_model_data(primary_stem)
+                
+        if self.process_method == DEMUCS_ARCH_TYPE and not is_secondary_model:
+            if self.demucs_stem_count >= 3 and self.pre_proc_model_activated:
+                self.pre_proc_model_activated = True
+                self.pre_proc_model = root.process_determine_demucs_pre_proc_model(self.primary_stem)
+                self.is_demucs_pre_proc_model_inst_mix = root.is_demucs_pre_proc_model_inst_mix_var.get() if self.pre_proc_model else False
+
+    def secondary_model_data(self, primary_stem):
+        secondary_model_data = root.process_determine_secondary_model(self.process_method, primary_stem, self.is_primary_stem_only, self.is_secondary_stem_only)
+        self.secondary_model = secondary_model_data[0]
+        self.secondary_model_scale = secondary_model_data[1]
+        self.is_secondary_model_activated = False if not self.secondary_model else True
+        if self.secondary_model:
+            self.is_secondary_model_activated = False if self.secondary_model.model_basename == self.model_basename else True
+              
+    def get_mdx_model_path(self):
+        
+        if self.model_name.endswith(CKPT):
+            # self.chunks = 0
+            # self.is_mdx_batch_mode = True
+            self.is_mdx_ckpt = True
+            
+        ext = '' if self.is_mdx_ckpt else ONNX
+        
+        for file_name, chosen_mdx_model in root.mdx_name_select_MAPPER.items():
+            if self.model_name in chosen_mdx_model:
+                self.model_path = os.path.join(MDX_MODELS_DIR, f"{file_name}{ext}")
+                break
+        else:
+            self.model_path = os.path.join(MDX_MODELS_DIR, f"{self.model_name}{ext}")
+            
+        self.mixer_path = os.path.join(MDX_MODELS_DIR, f"mixer_val.ckpt")
+    
+    def get_demucs_model_path(self):
+        
+        demucs_newer = [True for x in DEMUCS_NEWER_TAGS if x in self.model_name]
+        demucs_model_dir = DEMUCS_NEWER_REPO_DIR if demucs_newer else DEMUCS_MODELS_DIR
+        
+        for file_name, chosen_model in root.demucs_name_select_MAPPER.items():
+            if self.model_name in chosen_model:
+                self.model_path = os.path.join(demucs_model_dir, file_name)
+                break
+        else:
+            self.model_path = os.path.join(DEMUCS_NEWER_REPO_DIR, f'{self.model_name}.yaml')
+
+    def get_demucs_model_data(self):
+
+        self.demucs_version = DEMUCS_V4
+
+        for key, value in DEMUCS_VERSION_MAPPER.items():
+            if value in self.model_name:
+                self.demucs_version = key
+
+        self.demucs_source_list = DEMUCS_2_SOURCE if DEMUCS_UVR_MODEL in self.model_name else DEMUCS_4_SOURCE
+        self.demucs_source_map = DEMUCS_2_SOURCE_MAPPER if DEMUCS_UVR_MODEL in self.model_name else DEMUCS_4_SOURCE_MAPPER
+        self.demucs_stem_count = 2 if DEMUCS_UVR_MODEL in self.model_name else 4
+        
+        if not self.is_ensemble_mode:
+            self.primary_stem = PRIMARY_STEM if self.demucs_stems == ALL_STEMS else self.demucs_stems
+            self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
+
+    def get_model_data(self, model_hash_dir, hash_mapper):
+
+        model_settings_json = os.path.join(model_hash_dir, "{}.json".format(self.model_hash))
+
+        if os.path.isfile(model_settings_json):
+            return json.load(open(model_settings_json))
+        else:
+            for hash, settings in hash_mapper.items():
+                if self.model_hash in hash:
+                    return settings
+            else:
+                return self.get_model_data_from_popup()
+
+    def get_model_data_from_popup(self):
+        
+        if not self.is_dry_check:
+            confirm = tk.messagebox.askyesno(title=UNRECOGNIZED_MODEL[0],
+                                             message=f"\"{self.model_name}\"{UNRECOGNIZED_MODEL[1]}",
+                                             parent=root)
+                                            
+            if confirm:
+                if self.process_method == VR_ARCH_TYPE:
+                    root.pop_up_vr_param(self.model_hash)
+                    return root.vr_model_params
+                if self.process_method == MDX_ARCH_TYPE:
+                    root.pop_up_mdx_model(self.model_hash, self.model_path)
+                    return root.mdx_model_params
+            else:
+                return None
+        else:
+            return None
+
+    def get_model_hash(self):
+        self.model_hash = None
+        
+        if not os.path.isfile(self.model_path):
+            self.model_status = False
+            self.model_hash is None
+        else:
+            if model_hash_table:
+                for (key, value) in model_hash_table.items():
+                    if self.model_path == key:
+                        self.model_hash = value
+                        break
+                    
+            if not self.model_hash:
+                try:
+                    with open(self.model_path, 'rb') as f:
+                        f.seek(- 10000 * 1024, 2)
+                        self.model_hash = hashlib.md5(f.read()).hexdigest()
+                except:
+                    self.model_hash = hashlib.md5(open(self.model_path,'rb').read()).hexdigest()
+                    
+                table_entry = {self.model_path: self.model_hash}
+                model_hash_table.update(table_entry)

lib/uvr5_pack/demucs/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.

lib/uvr5_pack/demucs/__main__.py

@@ -0,0 +1,272 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import os
+import sys
+import time
+from dataclasses import dataclass, field
+from fractions import Fraction
+
+import torch as th
+from torch import distributed, nn
+from torch.nn.parallel.distributed import DistributedDataParallel
+
+from .augment import FlipChannels, FlipSign, Remix, Shift
+from .compressed import StemsSet, build_musdb_metadata, get_musdb_tracks
+from .model import Demucs
+from .parser import get_name, get_parser
+from .raw import Rawset
+from .tasnet import ConvTasNet
+from .test import evaluate
+from .train import train_model, validate_model
+from .utils import human_seconds, load_model, save_model, sizeof_fmt
+
+
+@dataclass
+class SavedState:
+    metrics: list = field(default_factory=list)
+    last_state: dict = None
+    best_state: dict = None
+    optimizer: dict = None
+
+
+def main():
+    parser = get_parser()
+    args = parser.parse_args()
+    name = get_name(parser, args)
+    print(f"Experiment {name}")
+
+    if args.musdb is None and args.rank == 0:
+        print(
+            "You must provide the path to the MusDB dataset with the --musdb flag. "
+            "To download the MusDB dataset, see https://sigsep.github.io/datasets/musdb.html.",
+            file=sys.stderr)
+        sys.exit(1)
+
+    eval_folder = args.evals / name
+    eval_folder.mkdir(exist_ok=True, parents=True)
+    args.logs.mkdir(exist_ok=True)
+    metrics_path = args.logs / f"{name}.json"
+    eval_folder.mkdir(exist_ok=True, parents=True)
+    args.checkpoints.mkdir(exist_ok=True, parents=True)
+    args.models.mkdir(exist_ok=True, parents=True)
+
+    if args.device is None:
+        device = "cpu"
+        if th.cuda.is_available():
+            device = "cuda"
+    else:
+        device = args.device
+
+    th.manual_seed(args.seed)
+    # Prevents too many threads to be started when running `museval` as it can be quite
+    # inefficient on NUMA architectures.
+    os.environ["OMP_NUM_THREADS"] = "1"
+
+    if args.world_size > 1:
+        if device != "cuda" and args.rank == 0:
+            print("Error: distributed training is only available with cuda device", file=sys.stderr)
+            sys.exit(1)
+        th.cuda.set_device(args.rank % th.cuda.device_count())
+        distributed.init_process_group(backend="nccl",
+                                       init_method="tcp://" + args.master,
+                                       rank=args.rank,
+                                       world_size=args.world_size)
+
+    checkpoint = args.checkpoints / f"{name}.th"
+    checkpoint_tmp = args.checkpoints / f"{name}.th.tmp"
+    if args.restart and checkpoint.exists():
+        checkpoint.unlink()
+
+    if args.test:
+        args.epochs = 1
+        args.repeat = 0
+        model = load_model(args.models / args.test)
+    elif args.tasnet:
+        model = ConvTasNet(audio_channels=args.audio_channels, samplerate=args.samplerate, X=args.X)
+    else:
+        model = Demucs(
+            audio_channels=args.audio_channels,
+            channels=args.channels,
+            context=args.context,
+            depth=args.depth,
+            glu=args.glu,
+            growth=args.growth,
+            kernel_size=args.kernel_size,
+            lstm_layers=args.lstm_layers,
+            rescale=args.rescale,
+            rewrite=args.rewrite,
+            sources=4,
+            stride=args.conv_stride,
+            upsample=args.upsample,
+            samplerate=args.samplerate
+        )
+    model.to(device)
+    if args.show:
+        print(model)
+        size = sizeof_fmt(4 * sum(p.numel() for p in model.parameters()))
+        print(f"Model size {size}")
+        return
+
+    optimizer = th.optim.Adam(model.parameters(), lr=args.lr)
+
+    try:
+        saved = th.load(checkpoint, map_location='cpu')
+    except IOError:
+        saved = SavedState()
+    else:
+        model.load_state_dict(saved.last_state)
+        optimizer.load_state_dict(saved.optimizer)
+
+    if args.save_model:
+        if args.rank == 0:
+            model.to("cpu")
+            model.load_state_dict(saved.best_state)
+            save_model(model, args.models / f"{name}.th")
+        return
+
+    if args.rank == 0:
+        done = args.logs / f"{name}.done"
+        if done.exists():
+            done.unlink()
+
+    if args.augment:
+        augment = nn.Sequential(FlipSign(), FlipChannels(), Shift(args.data_stride),
+                                Remix(group_size=args.remix_group_size)).to(device)
+    else:
+        augment = Shift(args.data_stride)
+
+    if args.mse:
+        criterion = nn.MSELoss()
+    else:
+        criterion = nn.L1Loss()
+
+    # Setting number of samples so that all convolution windows are full.
+    # Prevents hard to debug mistake with the prediction being shifted compared
+    # to the input mixture.
+    samples = model.valid_length(args.samples)
+    print(f"Number of training samples adjusted to {samples}")
+
+    if args.raw:
+        train_set = Rawset(args.raw / "train",
+                           samples=samples + args.data_stride,
+                           channels=args.audio_channels,
+                           streams=[0, 1, 2, 3, 4],
+                           stride=args.data_stride)
+
+        valid_set = Rawset(args.raw / "valid", channels=args.audio_channels)
+    else:
+        if not args.metadata.is_file() and args.rank == 0:
+            build_musdb_metadata(args.metadata, args.musdb, args.workers)
+        if args.world_size > 1:
+            distributed.barrier()
+        metadata = json.load(open(args.metadata))
+        duration = Fraction(samples + args.data_stride, args.samplerate)
+        stride = Fraction(args.data_stride, args.samplerate)
+        train_set = StemsSet(get_musdb_tracks(args.musdb, subsets=["train"], split="train"),
+                             metadata,
+                             duration=duration,
+                             stride=stride,
+                             samplerate=args.samplerate,
+                             channels=args.audio_channels)
+        valid_set = StemsSet(get_musdb_tracks(args.musdb, subsets=["train"], split="valid"),
+                             metadata,
+                             samplerate=args.samplerate,
+                             channels=args.audio_channels)
+
+    best_loss = float("inf")
+    for epoch, metrics in enumerate(saved.metrics):
+        print(f"Epoch {epoch:03d}: "
+              f"train={metrics['train']:.8f} "
+              f"valid={metrics['valid']:.8f} "
+              f"best={metrics['best']:.4f} "
+              f"duration={human_seconds(metrics['duration'])}")
+        best_loss = metrics['best']
+
+    if args.world_size > 1:
+        dmodel = DistributedDataParallel(model,
+                                         device_ids=[th.cuda.current_device()],
+                                         output_device=th.cuda.current_device())
+    else:
+        dmodel = model
+
+    for epoch in range(len(saved.metrics), args.epochs):
+        begin = time.time()
+        model.train()
+        train_loss = train_model(epoch,
+                                 train_set,
+                                 dmodel,
+                                 criterion,
+                                 optimizer,
+                                 augment,
+                                 batch_size=args.batch_size,
+                                 device=device,
+                                 repeat=args.repeat,
+                                 seed=args.seed,
+                                 workers=args.workers,
+                                 world_size=args.world_size)
+        model.eval()
+        valid_loss = validate_model(epoch,
+                                    valid_set,
+                                    model,
+                                    criterion,
+                                    device=device,
+                                    rank=args.rank,
+                                    split=args.split_valid,
+                                    world_size=args.world_size)
+
+        duration = time.time() - begin
+        if valid_loss < best_loss:
+            best_loss = valid_loss
+            saved.best_state = {
+                key: value.to("cpu").clone()
+                for key, value in model.state_dict().items()
+            }
+        saved.metrics.append({
+            "train": train_loss,
+            "valid": valid_loss,
+            "best": best_loss,
+            "duration": duration
+        })
+        if args.rank == 0:
+            json.dump(saved.metrics, open(metrics_path, "w"))
+
+        saved.last_state = model.state_dict()
+        saved.optimizer = optimizer.state_dict()
+        if args.rank == 0 and not args.test:
+            th.save(saved, checkpoint_tmp)
+            checkpoint_tmp.rename(checkpoint)
+
+        print(f"Epoch {epoch:03d}: "
+              f"train={train_loss:.8f} valid={valid_loss:.8f} best={best_loss:.4f} "
+              f"duration={human_seconds(duration)}")
+
+    del dmodel
+    model.load_state_dict(saved.best_state)
+    if args.eval_cpu:
+        device = "cpu"
+        model.to(device)
+    model.eval()
+    evaluate(model,
+             args.musdb,
+             eval_folder,
+             rank=args.rank,
+             world_size=args.world_size,
+             device=device,
+             save=args.save,
+             split=args.split_valid,
+             shifts=args.shifts,
+             workers=args.eval_workers)
+    model.to("cpu")
+    save_model(model, args.models / f"{name}.th")
+    if args.rank == 0:
+        print("done")
+        done.write_text("done")
+
+
+if __name__ == "__main__":
+    main()

lib/uvr5_pack/demucs/apply.py

@@ -0,0 +1,294 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Code to apply a model to a mix. It will handle chunking with overlaps and
+inteprolation between chunks, as well as the "shift trick".
+"""
+from concurrent.futures import ThreadPoolExecutor
+import random
+import typing as tp
+from multiprocessing import Process,Queue,Pipe
+
+import torch as th
+from torch import nn
+from torch.nn import functional as F
+import tqdm
+import tkinter as tk
+
+from .demucs import Demucs
+from .hdemucs import HDemucs
+from .utils import center_trim, DummyPoolExecutor
+
+Model = tp.Union[Demucs, HDemucs]
+
+progress_bar_num = 0
+
+class BagOfModels(nn.Module):
+    def __init__(self, models: tp.List[Model],
+                 weights: tp.Optional[tp.List[tp.List[float]]] = None,
+                 segment: tp.Optional[float] = None):
+        """
+        Represents a bag of models with specific weights.
+        You should call `apply_model` rather than calling directly the forward here for
+        optimal performance.
+
+        Args:
+            models (list[nn.Module]): list of Demucs/HDemucs models.
+            weights (list[list[float]]): list of weights. If None, assumed to
+                be all ones, otherwise it should be a list of N list (N number of models),
+                each containing S floats (S number of sources).
+            segment (None or float): overrides the `segment` attribute of each model
+                (this is performed inplace, be careful if you reuse the models passed).
+        """
+        
+        super().__init__()
+        assert len(models) > 0
+        first = models[0]
+        for other in models:
+            assert other.sources == first.sources
+            assert other.samplerate == first.samplerate
+            assert other.audio_channels == first.audio_channels
+            if segment is not None:
+                other.segment = segment
+
+        self.audio_channels = first.audio_channels
+        self.samplerate = first.samplerate
+        self.sources = first.sources
+        self.models = nn.ModuleList(models)
+
+        if weights is None:
+            weights = [[1. for _ in first.sources] for _ in models]
+        else:
+            assert len(weights) == len(models)
+            for weight in weights:
+                assert len(weight) == len(first.sources)
+        self.weights = weights
+
+    def forward(self, x):
+        raise NotImplementedError("Call `apply_model` on this.")
+
+class TensorChunk:
+    def __init__(self, tensor, offset=0, length=None):
+        total_length = tensor.shape[-1]
+        assert offset >= 0
+        assert offset < total_length
+
+        if length is None:
+            length = total_length - offset
+        else:
+            length = min(total_length - offset, length)
+
+        if isinstance(tensor, TensorChunk):
+            self.tensor = tensor.tensor
+            self.offset = offset + tensor.offset
+        else:
+            self.tensor = tensor
+            self.offset = offset
+        self.length = length
+        self.device = tensor.device
+
+    @property
+    def shape(self):
+        shape = list(self.tensor.shape)
+        shape[-1] = self.length
+        return shape
+
+    def padded(self, target_length):
+        delta = target_length - self.length
+        total_length = self.tensor.shape[-1]
+        assert delta >= 0
+
+        start = self.offset - delta // 2
+        end = start + target_length
+
+        correct_start = max(0, start)
+        correct_end = min(total_length, end)
+
+        pad_left = correct_start - start
+        pad_right = end - correct_end
+
+        out = F.pad(self.tensor[..., correct_start:correct_end], (pad_left, pad_right))
+        assert out.shape[-1] == target_length
+        return out
+
+def tensor_chunk(tensor_or_chunk):
+    if isinstance(tensor_or_chunk, TensorChunk):
+        return tensor_or_chunk
+    else:
+        assert isinstance(tensor_or_chunk, th.Tensor)
+        return TensorChunk(tensor_or_chunk)
+
+def apply_model(model, mix, shifts=1, split=True, overlap=0.25, transition_power=1., static_shifts=1, set_progress_bar=None, device=None, progress=False, num_workers=0, pool=None): 
+    """
+    Apply model to a given mixture.
+
+    Args:
+        shifts (int): if > 0, will shift in time `mix` by a random amount between 0 and 0.5 sec
+            and apply the oppositve shift to the output. This is repeated `shifts` time and
+            all predictions are averaged. This effectively makes the model time equivariant
+            and improves SDR by up to 0.2 points.
+        split (bool): if True, the input will be broken down in 8 seconds extracts
+            and predictions will be performed individually on each and concatenated.
+            Useful for model with large memory footprint like Tasnet.
+        progress (bool): if True, show a progress bar (requires split=True)
+        device (torch.device, str, or None): if provided, device on which to
+            execute the computation, otherwise `mix.device` is assumed.
+            When `device` is different from `mix.device`, only local computations will
+            be on `device`, while the entire tracks will be stored on `mix.device`.
+    """
+    
+    global fut_length
+    global bag_num
+    global prog_bar
+    
+    if device is None:
+        device = mix.device
+    else:
+        device = th.device(device)
+    if pool is None:
+        if num_workers > 0 and device.type == 'cpu':
+            pool = ThreadPoolExecutor(num_workers)
+        else:
+            pool = DummyPoolExecutor()
+            
+    kwargs = {
+        'shifts': shifts,
+        'split': split,
+        'overlap': overlap,
+        'transition_power': transition_power,
+        'progress': progress,
+        'device': device,
+        'pool': pool,
+        'set_progress_bar': set_progress_bar,
+        'static_shifts': static_shifts,
+    }
+    
+    if isinstance(model, BagOfModels):
+        # Special treatment for bag of model.
+        # We explicitely apply multiple times `apply_model` so that the random shifts
+        # are different for each model.
+
+        estimates = 0
+        totals = [0] * len(model.sources)
+        bag_num = len(model.models)
+        fut_length = 0
+        prog_bar = 0
+        current_model = 0 #(bag_num + 1)
+        for sub_model, weight in zip(model.models, model.weights):
+            original_model_device = next(iter(sub_model.parameters())).device
+            sub_model.to(device)
+            fut_length += fut_length
+            current_model += 1
+            out = apply_model(sub_model, mix, **kwargs)
+            sub_model.to(original_model_device)
+            for k, inst_weight in enumerate(weight):
+                out[:, k, :, :] *= inst_weight
+                totals[k] += inst_weight
+            estimates += out
+            del out
+
+        for k in range(estimates.shape[1]):
+            estimates[:, k, :, :] /= totals[k]
+        return estimates
+
+    model.to(device)
+    model.eval()
+    assert transition_power >= 1, "transition_power < 1 leads to weird behavior."
+    batch, channels, length = mix.shape
+    
+    if shifts:
+        kwargs['shifts'] = 0
+        max_shift = int(0.5 * model.samplerate)
+        mix = tensor_chunk(mix)
+        padded_mix = mix.padded(length + 2 * max_shift)
+        out = 0
+        for _ in range(shifts):
+            offset = random.randint(0, max_shift)
+            shifted = TensorChunk(padded_mix, offset, length + max_shift - offset)
+            shifted_out = apply_model(model, shifted, **kwargs)
+            out += shifted_out[..., max_shift - offset:]
+        out /= shifts
+        return out
+    elif split:
+        kwargs['split'] = False
+        out = th.zeros(batch, len(model.sources), channels, length, device=mix.device)
+        sum_weight = th.zeros(length, device=mix.device)
+        segment = int(model.samplerate * model.segment)
+        stride = int((1 - overlap) * segment)
+        offsets = range(0, length, stride)
+        scale = float(format(stride / model.samplerate, ".2f"))
+        # We start from a triangle shaped weight, with maximal weight in the middle
+        # of the segment. Then we normalize and take to the power `transition_power`.
+        # Large values of transition power will lead to sharper transitions.
+        weight = th.cat([th.arange(1, segment // 2 + 1, device=device),
+                         th.arange(segment - segment // 2, 0, -1, device=device)])
+        assert len(weight) == segment
+        # If the overlap < 50%, this will translate to linear transition when
+        # transition_power is 1.
+        weight = (weight / weight.max())**transition_power
+        futures = []
+        for offset in offsets:
+            chunk = TensorChunk(mix, offset, segment)
+            future = pool.submit(apply_model, model, chunk, **kwargs)
+            futures.append((future, offset))
+            offset += segment
+        if progress:
+            futures = tqdm.tqdm(futures, unit_scale=scale, ncols=120, unit='seconds')
+        for future, offset in futures:
+            if set_progress_bar:
+                fut_length = (len(futures) * bag_num * static_shifts)
+                prog_bar += 1
+                set_progress_bar(0.1, (0.8/fut_length*prog_bar))
+            chunk_out = future.result()
+            chunk_length = chunk_out.shape[-1]
+            out[..., offset:offset + segment] += (weight[:chunk_length] * chunk_out).to(mix.device)
+            sum_weight[offset:offset + segment] += weight[:chunk_length].to(mix.device)
+        assert sum_weight.min() > 0
+        out /= sum_weight
+        return out
+    else:
+        if hasattr(model, 'valid_length'):
+            valid_length = model.valid_length(length)
+        else:
+            valid_length = length
+        mix = tensor_chunk(mix)
+        padded_mix = mix.padded(valid_length).to(device)
+        with th.no_grad():
+            out = model(padded_mix)
+        return center_trim(out, length)
+    
+def demucs_segments(demucs_segment, demucs_model):
+    
+    if demucs_segment == 'Default':
+        segment = None
+        if isinstance(demucs_model, BagOfModels):
+            if segment is not None:
+                for sub in demucs_model.models:
+                    sub.segment = segment
+        else:
+            if segment is not None:
+                sub.segment = segment
+    else:
+        try:
+            segment = int(demucs_segment)
+            if isinstance(demucs_model, BagOfModels):
+                if segment is not None:
+                    for sub in demucs_model.models:
+                        sub.segment = segment
+            else:
+                if segment is not None:
+                    sub.segment = segment
+        except:
+            segment = None
+            if isinstance(demucs_model, BagOfModels):
+                if segment is not None:
+                    for sub in demucs_model.models:
+                        sub.segment = segment
+            else:
+                if segment is not None:
+                    sub.segment = segment
+                    
+    return demucs_model

lib/uvr5_pack/demucs/demucs.py

@@ -0,0 +1,459 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import typing as tp
+
+import julius
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .states import capture_init
+from .utils import center_trim, unfold
+
+
+class BLSTM(nn.Module):
+    """
+    BiLSTM with same hidden units as input dim.
+    If `max_steps` is not None, input will be splitting in overlapping
+    chunks and the LSTM applied separately on each chunk.
+    """
+    def __init__(self, dim, layers=1, max_steps=None, skip=False):
+        super().__init__()
+        assert max_steps is None or max_steps % 4 == 0
+        self.max_steps = max_steps
+        self.lstm = nn.LSTM(bidirectional=True, num_layers=layers, hidden_size=dim, input_size=dim)
+        self.linear = nn.Linear(2 * dim, dim)
+        self.skip = skip
+
+    def forward(self, x):
+        B, C, T = x.shape
+        y = x
+        framed = False
+        if self.max_steps is not None and T > self.max_steps:
+            width = self.max_steps
+            stride = width // 2
+            frames = unfold(x, width, stride)
+            nframes = frames.shape[2]
+            framed = True
+            x = frames.permute(0, 2, 1, 3).reshape(-1, C, width)
+
+        x = x.permute(2, 0, 1)
+
+        x = self.lstm(x)[0]
+        x = self.linear(x)
+        x = x.permute(1, 2, 0)
+        if framed:
+            out = []
+            frames = x.reshape(B, -1, C, width)
+            limit = stride // 2
+            for k in range(nframes):
+                if k == 0:
+                    out.append(frames[:, k, :, :-limit])
+                elif k == nframes - 1:
+                    out.append(frames[:, k, :, limit:])
+                else:
+                    out.append(frames[:, k, :, limit:-limit])
+            out = torch.cat(out, -1)
+            out = out[..., :T]
+            x = out
+        if self.skip:
+            x = x + y
+        return x
+
+
+def rescale_conv(conv, reference):
+    """Rescale initial weight scale. It is unclear why it helps but it certainly does.
+    """
+    std = conv.weight.std().detach()
+    scale = (std / reference)**0.5
+    conv.weight.data /= scale
+    if conv.bias is not None:
+        conv.bias.data /= scale
+
+
+def rescale_module(module, reference):
+    for sub in module.modules():
+        if isinstance(sub, (nn.Conv1d, nn.ConvTranspose1d, nn.Conv2d, nn.ConvTranspose2d)):
+            rescale_conv(sub, reference)
+
+
+class LayerScale(nn.Module):
+    """Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf).
+    This rescales diagonaly residual outputs close to 0 initially, then learnt.
+    """
+    def __init__(self, channels: int, init: float = 0):
+        super().__init__()
+        self.scale = nn.Parameter(torch.zeros(channels, requires_grad=True))
+        self.scale.data[:] = init
+
+    def forward(self, x):
+        return self.scale[:, None] * x
+
+
+class DConv(nn.Module):
+    """
+    New residual branches in each encoder layer.
+    This alternates dilated convolutions, potentially with LSTMs and attention.
+    Also before entering each residual branch, dimension is projected on a smaller subspace,
+    e.g. of dim `channels // compress`.
+    """
+    def __init__(self, channels: int, compress: float = 4, depth: int = 2, init: float = 1e-4,
+                 norm=True, attn=False, heads=4, ndecay=4, lstm=False, gelu=True,
+                 kernel=3, dilate=True):
+        """
+        Args:
+            channels: input/output channels for residual branch.
+            compress: amount of channel compression inside the branch.
+            depth: number of layers in the residual branch. Each layer has its own
+                projection, and potentially LSTM and attention.
+            init: initial scale for LayerNorm.
+            norm: use GroupNorm.
+            attn: use LocalAttention.
+            heads: number of heads for the LocalAttention.
+            ndecay: number of decay controls in the LocalAttention.
+            lstm: use LSTM.
+            gelu: Use GELU activation.
+            kernel: kernel size for the (dilated) convolutions.
+            dilate: if true, use dilation, increasing with the depth.
+        """
+
+        super().__init__()
+        assert kernel % 2 == 1
+        self.channels = channels
+        self.compress = compress
+        self.depth = abs(depth)
+        dilate = depth > 0
+
+        norm_fn: tp.Callable[[int], nn.Module]
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm:
+            norm_fn = lambda d: nn.GroupNorm(1, d)  # noqa
+
+        hidden = int(channels / compress)
+
+        act: tp.Type[nn.Module]
+        if gelu:
+            act = nn.GELU
+        else:
+            act = nn.ReLU
+
+        self.layers = nn.ModuleList([])
+        for d in range(self.depth):
+            dilation = 2 ** d if dilate else 1
+            padding = dilation * (kernel // 2)
+            mods = [
+                nn.Conv1d(channels, hidden, kernel, dilation=dilation, padding=padding),
+                norm_fn(hidden), act(),
+                nn.Conv1d(hidden, 2 * channels, 1),
+                norm_fn(2 * channels), nn.GLU(1),
+                LayerScale(channels, init),
+            ]
+            if attn:
+                mods.insert(3, LocalState(hidden, heads=heads, ndecay=ndecay))
+            if lstm:
+                mods.insert(3, BLSTM(hidden, layers=2, max_steps=200, skip=True))
+            layer = nn.Sequential(*mods)
+            self.layers.append(layer)
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = x + layer(x)
+        return x
+
+
+class LocalState(nn.Module):
+    """Local state allows to have attention based only on data (no positional embedding),
+    but while setting a constraint on the time window (e.g. decaying penalty term).
+
+    Also a failed experiments with trying to provide some frequency based attention.
+    """
+    def __init__(self, channels: int, heads: int = 4, nfreqs: int = 0, ndecay: int = 4):
+        super().__init__()
+        assert channels % heads == 0, (channels, heads)
+        self.heads = heads
+        self.nfreqs = nfreqs
+        self.ndecay = ndecay
+        self.content = nn.Conv1d(channels, channels, 1)
+        self.query = nn.Conv1d(channels, channels, 1)
+        self.key = nn.Conv1d(channels, channels, 1)
+        if nfreqs:
+            self.query_freqs = nn.Conv1d(channels, heads * nfreqs, 1)
+        if ndecay:
+            self.query_decay = nn.Conv1d(channels, heads * ndecay, 1)
+            # Initialize decay close to zero (there is a sigmoid), for maximum initial window.
+            self.query_decay.weight.data *= 0.01
+            assert self.query_decay.bias is not None  # stupid type checker
+            self.query_decay.bias.data[:] = -2
+        self.proj = nn.Conv1d(channels + heads * nfreqs, channels, 1)
+
+    def forward(self, x):
+        B, C, T = x.shape
+        heads = self.heads
+        indexes = torch.arange(T, device=x.device, dtype=x.dtype)
+        # left index are keys, right index are queries
+        delta = indexes[:, None] - indexes[None, :]
+
+        queries = self.query(x).view(B, heads, -1, T)
+        keys = self.key(x).view(B, heads, -1, T)
+        # t are keys, s are queries
+        dots = torch.einsum("bhct,bhcs->bhts", keys, queries)
+        dots /= keys.shape[2]**0.5
+        if self.nfreqs:
+            periods = torch.arange(1, self.nfreqs + 1, device=x.device, dtype=x.dtype)
+            freq_kernel = torch.cos(2 * math.pi * delta / periods.view(-1, 1, 1))
+            freq_q = self.query_freqs(x).view(B, heads, -1, T) / self.nfreqs ** 0.5
+            dots += torch.einsum("fts,bhfs->bhts", freq_kernel, freq_q)
+        if self.ndecay:
+            decays = torch.arange(1, self.ndecay + 1, device=x.device, dtype=x.dtype)
+            decay_q = self.query_decay(x).view(B, heads, -1, T)
+            decay_q = torch.sigmoid(decay_q) / 2
+            decay_kernel = - decays.view(-1, 1, 1) * delta.abs() / self.ndecay**0.5
+            dots += torch.einsum("fts,bhfs->bhts", decay_kernel, decay_q)
+
+        # Kill self reference.
+        dots.masked_fill_(torch.eye(T, device=dots.device, dtype=torch.bool), -100)
+        weights = torch.softmax(dots, dim=2)
+
+        content = self.content(x).view(B, heads, -1, T)
+        result = torch.einsum("bhts,bhct->bhcs", weights, content)
+        if self.nfreqs:
+            time_sig = torch.einsum("bhts,fts->bhfs", weights, freq_kernel)
+            result = torch.cat([result, time_sig], 2)
+        result = result.reshape(B, -1, T)
+        return x + self.proj(result)
+
+
+class Demucs(nn.Module):
+    @capture_init
+    def __init__(self,
+                 sources,
+                 # Channels
+                 audio_channels=2,
+                 channels=64,
+                 growth=2.,
+                 # Main structure
+                 depth=6,
+                 rewrite=True,
+                 lstm_layers=0,
+                 # Convolutions
+                 kernel_size=8,
+                 stride=4,
+                 context=1,
+                 # Activations
+                 gelu=True,
+                 glu=True,
+                 # Normalization
+                 norm_starts=4,
+                 norm_groups=4,
+                 # DConv residual branch
+                 dconv_mode=1,
+                 dconv_depth=2,
+                 dconv_comp=4,
+                 dconv_attn=4,
+                 dconv_lstm=4,
+                 dconv_init=1e-4,
+                 # Pre/post processing
+                 normalize=True,
+                 resample=True,
+                 # Weight init
+                 rescale=0.1,
+                 # Metadata
+                 samplerate=44100,
+                 segment=4 * 10):
+        """
+        Args:
+            sources (list[str]): list of source names
+            audio_channels (int): stereo or mono
+            channels (int): first convolution channels
+            depth (int): number of encoder/decoder layers
+            growth (float): multiply (resp divide) number of channels by that
+                for each layer of the encoder (resp decoder)
+            depth (int): number of layers in the encoder and in the decoder.
+            rewrite (bool): add 1x1 convolution to each layer.
+            lstm_layers (int): number of lstm layers, 0 = no lstm. Deactivated
+                by default, as this is now replaced by the smaller and faster small LSTMs
+                in the DConv branches.
+            kernel_size (int): kernel size for convolutions
+            stride (int): stride for convolutions
+            context (int): kernel size of the convolution in the
+                decoder before the transposed convolution. If > 1,
+                will provide some context from neighboring time steps.
+            gelu: use GELU activation function.
+            glu (bool): use glu instead of ReLU for the 1x1 rewrite conv.
+            norm_starts: layer at which group norm starts being used.
+                decoder layers are numbered in reverse order.
+            norm_groups: number of groups for group norm.
+            dconv_mode: if 1: dconv in encoder only, 2: decoder only, 3: both.
+            dconv_depth: depth of residual DConv branch.
+            dconv_comp: compression of DConv branch.
+            dconv_attn: adds attention layers in DConv branch starting at this layer.
+            dconv_lstm: adds a LSTM layer in DConv branch starting at this layer.
+            dconv_init: initial scale for the DConv branch LayerScale.
+            normalize (bool): normalizes the input audio on the fly, and scales back
+                the output by the same amount.
+            resample (bool): upsample x2 the input and downsample /2 the output.
+            rescale (int): rescale initial weights of convolutions
+                to get their standard deviation closer to `rescale`.
+            samplerate (int): stored as meta information for easing
+                future evaluations of the model.
+            segment (float): duration of the chunks of audio to ideally evaluate the model on.
+                This is used by `demucs.apply.apply_model`.
+        """
+
+        super().__init__()
+        self.audio_channels = audio_channels
+        self.sources = sources
+        self.kernel_size = kernel_size
+        self.context = context
+        self.stride = stride
+        self.depth = depth
+        self.resample = resample
+        self.channels = channels
+        self.normalize = normalize
+        self.samplerate = samplerate
+        self.segment = segment
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+        self.skip_scales = nn.ModuleList()
+
+        if glu:
+            activation = nn.GLU(dim=1)
+            ch_scale = 2
+        else:
+            activation = nn.ReLU()
+            ch_scale = 1
+        if gelu:
+            act2 = nn.GELU
+        else:
+            act2 = nn.ReLU
+
+        in_channels = audio_channels
+        padding = 0
+        for index in range(depth):
+            norm_fn = lambda d: nn.Identity()  # noqa
+            if index >= norm_starts:
+                norm_fn = lambda d: nn.GroupNorm(norm_groups, d)  # noqa
+
+            encode = []
+            encode += [
+                nn.Conv1d(in_channels, channels, kernel_size, stride),
+                norm_fn(channels),
+                act2(),
+            ]
+            attn = index >= dconv_attn
+            lstm = index >= dconv_lstm
+            if dconv_mode & 1:
+                encode += [DConv(channels, depth=dconv_depth, init=dconv_init,
+                                 compress=dconv_comp, attn=attn, lstm=lstm)]
+            if rewrite:
+                encode += [
+                    nn.Conv1d(channels, ch_scale * channels, 1),
+                    norm_fn(ch_scale * channels), activation]
+            self.encoder.append(nn.Sequential(*encode))
+
+            decode = []
+            if index > 0:
+                out_channels = in_channels
+            else:
+                out_channels = len(self.sources) * audio_channels
+            if rewrite:
+                decode += [
+                    nn.Conv1d(channels, ch_scale * channels, 2 * context + 1, padding=context),
+                    norm_fn(ch_scale * channels), activation]
+            if dconv_mode & 2:
+                decode += [DConv(channels, depth=dconv_depth, init=dconv_init,
+                                 compress=dconv_comp, attn=attn, lstm=lstm)]
+            decode += [nn.ConvTranspose1d(channels, out_channels,
+                       kernel_size, stride, padding=padding)]
+            if index > 0:
+                decode += [norm_fn(out_channels), act2()]
+            self.decoder.insert(0, nn.Sequential(*decode))
+            in_channels = channels
+            channels = int(growth * channels)
+
+        channels = in_channels
+        if lstm_layers:
+            self.lstm = BLSTM(channels, lstm_layers)
+        else:
+            self.lstm = None
+
+        if rescale:
+            rescale_module(self, reference=rescale)
+
+    def valid_length(self, length):
+        """
+        Return the nearest valid length to use with the model so that
+        there is no time steps left over in a convolution, e.g. for all
+        layers, size of the input - kernel_size % stride = 0.
+
+        Note that input are automatically padded if necessary to ensure that the output
+        has the same length as the input.
+        """
+        if self.resample:
+            length *= 2
+
+        for _ in range(self.depth):
+            length = math.ceil((length - self.kernel_size) / self.stride) + 1
+            length = max(1, length)
+
+        for idx in range(self.depth):
+            length = (length - 1) * self.stride + self.kernel_size
+
+        if self.resample:
+            length = math.ceil(length / 2)
+        return int(length)
+
+    def forward(self, mix):
+        x = mix
+        length = x.shape[-1]
+
+        if self.normalize:
+            mono = mix.mean(dim=1, keepdim=True)
+            mean = mono.mean(dim=-1, keepdim=True)
+            std = mono.std(dim=-1, keepdim=True)
+            x = (x - mean) / (1e-5 + std)
+        else:
+            mean = 0
+            std = 1
+
+        delta = self.valid_length(length) - length
+        x = F.pad(x, (delta // 2, delta - delta // 2))
+
+        if self.resample:
+            x = julius.resample_frac(x, 1, 2)
+
+        saved = []
+        for encode in self.encoder:
+            x = encode(x)
+            saved.append(x)
+
+        if self.lstm:
+            x = self.lstm(x)
+
+        for decode in self.decoder:
+            skip = saved.pop(-1)
+            skip = center_trim(skip, x)
+            x = decode(x + skip)
+
+        if self.resample:
+            x = julius.resample_frac(x, 2, 1)
+        x = x * std + mean
+        x = center_trim(x, length)
+        x = x.view(x.size(0), len(self.sources), self.audio_channels, x.size(-1))
+        return x
+
+    def load_state_dict(self, state, strict=True):
+        # fix a mismatch with previous generation Demucs models.
+        for idx in range(self.depth):
+            for a in ['encoder', 'decoder']:
+                for b in ['bias', 'weight']:
+                    new = f'{a}.{idx}.3.{b}'
+                    old = f'{a}.{idx}.2.{b}'
+                    if old in state and new not in state:
+                        state[new] = state.pop(old)
+        super().load_state_dict(state, strict=strict)

lib/uvr5_pack/demucs/filtering.py

@@ -0,0 +1,502 @@
+from typing import Optional
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+def atan2(y, x):
+    r"""Element-wise arctangent function of y/x.
+    Returns a new tensor with signed angles in radians.
+    It is an alternative implementation of torch.atan2
+
+    Args:
+        y (Tensor): First input tensor
+        x (Tensor): Second input tensor [shape=y.shape]
+
+    Returns:
+        Tensor: [shape=y.shape].
+    """
+    pi = 2 * torch.asin(torch.tensor(1.0))
+    x += ((x == 0) & (y == 0)) * 1.0
+    out = torch.atan(y / x)
+    out += ((y >= 0) & (x < 0)) * pi
+    out -= ((y < 0) & (x < 0)) * pi
+    out *= 1 - ((y > 0) & (x == 0)) * 1.0
+    out += ((y > 0) & (x == 0)) * (pi / 2)
+    out *= 1 - ((y < 0) & (x == 0)) * 1.0
+    out += ((y < 0) & (x == 0)) * (-pi / 2)
+    return out
+
+
+# Define basic complex operations on torch.Tensor objects whose last dimension
+# consists in the concatenation of the real and imaginary parts.
+
+
+def _norm(x: torch.Tensor) -> torch.Tensor:
+    r"""Computes the norm value of a torch Tensor, assuming that it
+    comes as real and imaginary part in its last dimension.
+
+    Args:
+        x (Tensor): Input Tensor of shape [shape=(..., 2)]
+
+    Returns:
+        Tensor: shape as x excluding the last dimension.
+    """
+    return torch.abs(x[..., 0]) ** 2 + torch.abs(x[..., 1]) ** 2
+
+
+def _mul_add(a: torch.Tensor, b: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Element-wise multiplication of two complex Tensors described
+    through their real and imaginary parts.
+    The result is added to the `out` tensor"""
+
+    # check `out` and allocate it if needed
+    target_shape = torch.Size([max(sa, sb) for (sa, sb) in zip(a.shape, b.shape)])
+    if out is None or out.shape != target_shape:
+        out = torch.zeros(target_shape, dtype=a.dtype, device=a.device)
+    if out is a:
+        real_a = a[..., 0]
+        out[..., 0] = out[..., 0] + (real_a * b[..., 0] - a[..., 1] * b[..., 1])
+        out[..., 1] = out[..., 1] + (real_a * b[..., 1] + a[..., 1] * b[..., 0])
+    else:
+        out[..., 0] = out[..., 0] + (a[..., 0] * b[..., 0] - a[..., 1] * b[..., 1])
+        out[..., 1] = out[..., 1] + (a[..., 0] * b[..., 1] + a[..., 1] * b[..., 0])
+    return out
+
+
+def _mul(a: torch.Tensor, b: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Element-wise multiplication of two complex Tensors described
+    through their real and imaginary parts
+    can work in place in case out is a only"""
+    target_shape = torch.Size([max(sa, sb) for (sa, sb) in zip(a.shape, b.shape)])
+    if out is None or out.shape != target_shape:
+        out = torch.zeros(target_shape, dtype=a.dtype, device=a.device)
+    if out is a:
+        real_a = a[..., 0]
+        out[..., 0] = real_a * b[..., 0] - a[..., 1] * b[..., 1]
+        out[..., 1] = real_a * b[..., 1] + a[..., 1] * b[..., 0]
+    else:
+        out[..., 0] = a[..., 0] * b[..., 0] - a[..., 1] * b[..., 1]
+        out[..., 1] = a[..., 0] * b[..., 1] + a[..., 1] * b[..., 0]
+    return out
+
+
+def _inv(z: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Element-wise multiplicative inverse of a Tensor with complex
+    entries described through their real and imaginary parts.
+    can work in place in case out is z"""
+    ez = _norm(z)
+    if out is None or out.shape != z.shape:
+        out = torch.zeros_like(z)
+    out[..., 0] = z[..., 0] / ez
+    out[..., 1] = -z[..., 1] / ez
+    return out
+
+
+def _conj(z, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Element-wise complex conjugate of a Tensor with complex entries
+    described through their real and imaginary parts.
+    can work in place in case out is z"""
+    if out is None or out.shape != z.shape:
+        out = torch.zeros_like(z)
+    out[..., 0] = z[..., 0]
+    out[..., 1] = -z[..., 1]
+    return out
+
+
+def _invert(M: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """
+    Invert 1x1 or 2x2 matrices
+
+    Will generate errors if the matrices are singular: user must handle this
+    through his own regularization schemes.
+
+    Args:
+        M (Tensor): [shape=(..., nb_channels, nb_channels, 2)]
+            matrices to invert: must be square along dimensions -3 and -2
+
+    Returns:
+        invM (Tensor): [shape=M.shape]
+            inverses of M
+    """
+    nb_channels = M.shape[-2]
+
+    if out is None or out.shape != M.shape:
+        out = torch.empty_like(M)
+
+    if nb_channels == 1:
+        # scalar case
+        out = _inv(M, out)
+    elif nb_channels == 2:
+        # two channels case: analytical expression
+
+        # first compute the determinent
+        det = _mul(M[..., 0, 0, :], M[..., 1, 1, :])
+        det = det - _mul(M[..., 0, 1, :], M[..., 1, 0, :])
+        # invert it
+        invDet = _inv(det)
+
+        # then fill out the matrix with the inverse
+        out[..., 0, 0, :] = _mul(invDet, M[..., 1, 1, :], out[..., 0, 0, :])
+        out[..., 1, 0, :] = _mul(-invDet, M[..., 1, 0, :], out[..., 1, 0, :])
+        out[..., 0, 1, :] = _mul(-invDet, M[..., 0, 1, :], out[..., 0, 1, :])
+        out[..., 1, 1, :] = _mul(invDet, M[..., 0, 0, :], out[..., 1, 1, :])
+    else:
+        raise Exception("Only 2 channels are supported for the torch version.")
+    return out
+
+
+# Now define the signal-processing low-level functions used by the Separator
+
+
+def expectation_maximization(
+    y: torch.Tensor,
+    x: torch.Tensor,
+    iterations: int = 2,
+    eps: float = 1e-10,
+    batch_size: int = 200,
+):
+    r"""Expectation maximization algorithm, for refining source separation
+    estimates.
+
+    This algorithm allows to make source separation results better by
+    enforcing multichannel consistency for the estimates. This usually means
+    a better perceptual quality in terms of spatial artifacts.
+
+    The implementation follows the details presented in [1]_, taking
+    inspiration from the original EM algorithm proposed in [2]_ and its
+    weighted refinement proposed in [3]_, [4]_.
+    It works by iteratively:
+
+     * Re-estimate source parameters (power spectral densities and spatial
+       covariance matrices) through :func:`get_local_gaussian_model`.
+
+     * Separate again the mixture with the new parameters by first computing
+       the new modelled mixture covariance matrices with :func:`get_mix_model`,
+       prepare the Wiener filters through :func:`wiener_gain` and apply them
+       with :func:`apply_filter``.
+
+    References
+    ----------
+    .. [1] S. Uhlich and M. Porcu and F. Giron and M. Enenkl and T. Kemp and
+        N. Takahashi and Y. Mitsufuji, "Improving music source separation based
+        on deep neural networks through data augmentation and network
+        blending." 2017 IEEE International Conference on Acoustics, Speech
+        and Signal Processing (ICASSP). IEEE, 2017.
+
+    .. [2] N.Q. Duong and E. Vincent and R.Gribonval. "Under-determined
+        reverberant audio source separation using a full-rank spatial
+        covariance model." IEEE Transactions on Audio, Speech, and Language
+        Processing 18.7 (2010): 1830-1840.
+
+    .. [3] A. Nugraha and A. Liutkus and E. Vincent. "Multichannel audio source
+        separation with deep neural networks." IEEE/ACM Transactions on Audio,
+        Speech, and Language Processing 24.9 (2016): 1652-1664.
+
+    .. [4] A. Nugraha and A. Liutkus and E. Vincent. "Multichannel music
+        separation with deep neural networks." 2016 24th European Signal
+        Processing Conference (EUSIPCO). IEEE, 2016.
+
+    .. [5] A. Liutkus and R. Badeau and G. Richard "Kernel additive models for
+        source separation." IEEE Transactions on Signal Processing
+        62.16 (2014): 4298-4310.
+
+    Args:
+        y (Tensor): [shape=(nb_frames, nb_bins, nb_channels, 2, nb_sources)]
+            initial estimates for the sources
+        x (Tensor): [shape=(nb_frames, nb_bins, nb_channels, 2)]
+            complex STFT of the mixture signal
+        iterations (int): [scalar]
+            number of iterations for the EM algorithm.
+        eps (float or None): [scalar]
+            The epsilon value to use for regularization and filters.
+
+    Returns:
+        y (Tensor): [shape=(nb_frames, nb_bins, nb_channels, 2, nb_sources)]
+            estimated sources after iterations
+        v (Tensor): [shape=(nb_frames, nb_bins, nb_sources)]
+            estimated power spectral densities
+        R (Tensor): [shape=(nb_bins, nb_channels, nb_channels, 2, nb_sources)]
+            estimated spatial covariance matrices
+
+    Notes:
+        * You need an initial estimate for the sources to apply this
+          algorithm. This is precisely what the :func:`wiener` function does.
+        * This algorithm *is not* an implementation of the "exact" EM
+          proposed in [1]_. In particular, it does compute the posterior
+          covariance matrices the same (exact) way. Instead, it uses the
+          simplified approximate scheme initially proposed in [5]_ and further
+          refined in [3]_, [4]_, that boils down to just take the empirical
+          covariance of the recent source estimates, followed by a weighted
+          average for the update of the spatial covariance matrix. It has been
+          empirically demonstrated that this simplified algorithm is more
+          robust for music separation.
+
+    Warning:
+        It is *very* important to make sure `x.dtype` is `torch.float64`
+        if you want double precision, because this function will **not**
+        do such conversion for you from `torch.complex32`, in case you want the
+        smaller RAM usage on purpose.
+
+        It is usually always better in terms of quality to have double
+        precision, by e.g. calling :func:`expectation_maximization`
+        with ``x.to(torch.float64)``.
+    """
+    # dimensions
+    (nb_frames, nb_bins, nb_channels) = x.shape[:-1]
+    nb_sources = y.shape[-1]
+
+    regularization = torch.cat(
+        (
+            torch.eye(nb_channels, dtype=x.dtype, device=x.device)[..., None],
+            torch.zeros((nb_channels, nb_channels, 1), dtype=x.dtype, device=x.device),
+        ),
+        dim=2,
+    )
+    regularization = torch.sqrt(torch.as_tensor(eps)) * (
+        regularization[None, None, ...].expand((-1, nb_bins, -1, -1, -1))
+    )
+
+    # allocate the spatial covariance matrices
+    R = [
+        torch.zeros((nb_bins, nb_channels, nb_channels, 2), dtype=x.dtype, device=x.device)
+        for j in range(nb_sources)
+    ]
+    weight: torch.Tensor = torch.zeros((nb_bins,), dtype=x.dtype, device=x.device)
+
+    v: torch.Tensor = torch.zeros((nb_frames, nb_bins, nb_sources), dtype=x.dtype, device=x.device)
+    for it in range(iterations):
+        # constructing the mixture covariance matrix. Doing it with a loop
+        # to avoid storing anytime in RAM the whole 6D tensor
+
+        # update the PSD as the average spectrogram over channels
+        v = torch.mean(torch.abs(y[..., 0, :]) ** 2 + torch.abs(y[..., 1, :]) ** 2, dim=-2)
+
+        # update spatial covariance matrices (weighted update)
+        for j in range(nb_sources):
+            R[j] = torch.tensor(0.0, device=x.device)
+            weight = torch.tensor(eps, device=x.device)
+            pos: int = 0
+            batch_size = batch_size if batch_size else nb_frames
+            while pos < nb_frames:
+                t = torch.arange(pos, min(nb_frames, pos + batch_size))
+                pos = int(t[-1]) + 1
+
+                R[j] = R[j] + torch.sum(_covariance(y[t, ..., j]), dim=0)
+                weight = weight + torch.sum(v[t, ..., j], dim=0)
+            R[j] = R[j] / weight[..., None, None, None]
+            weight = torch.zeros_like(weight)
+
+        # cloning y if we track gradient, because we're going to update it
+        if y.requires_grad:
+            y = y.clone()
+
+        pos = 0
+        while pos < nb_frames:
+            t = torch.arange(pos, min(nb_frames, pos + batch_size))
+            pos = int(t[-1]) + 1
+
+            y[t, ...] = torch.tensor(0.0, device=x.device, dtype=x.dtype)
+
+            # compute mix covariance matrix
+            Cxx = regularization
+            for j in range(nb_sources):
+                Cxx = Cxx + (v[t, ..., j, None, None, None] * R[j][None, ...].clone())
+
+            # invert it
+            inv_Cxx = _invert(Cxx)
+
+            # separate the sources
+            for j in range(nb_sources):
+
+                # create a wiener gain for this source
+                gain = torch.zeros_like(inv_Cxx)
+
+                # computes multichannel Wiener gain as v_j R_j inv_Cxx
+                indices = torch.cartesian_prod(
+                    torch.arange(nb_channels),
+                    torch.arange(nb_channels),
+                    torch.arange(nb_channels),
+                )
+                for index in indices:
+                    gain[:, :, index[0], index[1], :] = _mul_add(
+                        R[j][None, :, index[0], index[2], :].clone(),
+                        inv_Cxx[:, :, index[2], index[1], :],
+                        gain[:, :, index[0], index[1], :],
+                    )
+                gain = gain * v[t, ..., None, None, None, j]
+
+                # apply it to the mixture
+                for i in range(nb_channels):
+                    y[t, ..., j] = _mul_add(gain[..., i, :], x[t, ..., i, None, :], y[t, ..., j])
+
+    return y, v, R
+
+
+def wiener(
+    targets_spectrograms: torch.Tensor,
+    mix_stft: torch.Tensor,
+    iterations: int = 1,
+    softmask: bool = False,
+    residual: bool = False,
+    scale_factor: float = 10.0,
+    eps: float = 1e-10,
+):
+    """Wiener-based separation for multichannel audio.
+
+    The method uses the (possibly multichannel) spectrograms  of the
+    sources to separate the (complex) Short Term Fourier Transform  of the
+    mix. Separation is done in a sequential way by:
+
+    * Getting an initial estimate. This can be done in two ways: either by
+      directly using the spectrograms with the mixture phase, or
+      by using a softmasking strategy. This initial phase is controlled
+      by the `softmask` flag.
+
+    * If required, adding an additional residual target as the mix minus
+      all targets.
+
+    * Refinining these initial estimates through a call to
+      :func:`expectation_maximization` if the number of iterations is nonzero.
+
+    This implementation also allows to specify the epsilon value used for
+    regularization. It is based on [1]_, [2]_, [3]_, [4]_.
+
+    References
+    ----------
+    .. [1] S. Uhlich and M. Porcu and F. Giron and M. Enenkl and T. Kemp and
+        N. Takahashi and Y. Mitsufuji, "Improving music source separation based
+        on deep neural networks through data augmentation and network
+        blending." 2017 IEEE International Conference on Acoustics, Speech
+        and Signal Processing (ICASSP). IEEE, 2017.
+
+    .. [2] A. Nugraha and A. Liutkus and E. Vincent. "Multichannel audio source
+        separation with deep neural networks." IEEE/ACM Transactions on Audio,
+        Speech, and Language Processing 24.9 (2016): 1652-1664.
+
+    .. [3] A. Nugraha and A. Liutkus and E. Vincent. "Multichannel music
+        separation with deep neural networks." 2016 24th European Signal
+        Processing Conference (EUSIPCO). IEEE, 2016.
+
+    .. [4] A. Liutkus and R. Badeau and G. Richard "Kernel additive models for
+        source separation." IEEE Transactions on Signal Processing
+        62.16 (2014): 4298-4310.
+
+    Args:
+        targets_spectrograms (Tensor): spectrograms of the sources
+            [shape=(nb_frames, nb_bins, nb_channels, nb_sources)].
+            This is a nonnegative tensor that is
+            usually the output of the actual separation method of the user. The
+            spectrograms may be mono, but they need to be 4-dimensional in all
+            cases.
+        mix_stft (Tensor): [shape=(nb_frames, nb_bins, nb_channels, complex=2)]
+            STFT of the mixture signal.
+        iterations (int): [scalar]
+            number of iterations for the EM algorithm
+        softmask (bool): Describes how the initial estimates are obtained.
+            * if `False`, then the mixture phase will directly be used with the
+            spectrogram as initial estimates.
+            * if `True`, initial estimates are obtained by multiplying the
+            complex mix element-wise with the ratio of each target spectrogram
+            with the sum of them all. This strategy is better if the model are
+            not really good, and worse otherwise.
+        residual (bool): if `True`, an additional target is created, which is
+            equal to the mixture minus the other targets, before application of
+            expectation maximization
+        eps (float): Epsilon value to use for computing the separations.
+            This is used whenever division with a model energy is
+            performed, i.e. when softmasking and when iterating the EM.
+            It can be understood as the energy of the additional white noise
+            that is taken out when separating.
+
+    Returns:
+        Tensor: shape=(nb_frames, nb_bins, nb_channels, complex=2, nb_sources)
+            STFT of estimated sources
+
+    Notes:
+        * Be careful that you need *magnitude spectrogram estimates* for the
+        case `softmask==False`.
+        * `softmask=False` is recommended
+        * The epsilon value will have a huge impact on performance. If it's
+        large, only the parts of the signal with a significant energy will
+        be kept in the sources. This epsilon then directly controls the
+        energy of the reconstruction error.
+
+    Warning:
+        As in :func:`expectation_maximization`, we recommend converting the
+        mixture `x` to double precision `torch.float64` *before* calling
+        :func:`wiener`.
+    """
+    if softmask:
+        # if we use softmask, we compute the ratio mask for all targets and
+        # multiply by the mix stft
+        y = (
+            mix_stft[..., None]
+            * (
+                targets_spectrograms
+                / (eps + torch.sum(targets_spectrograms, dim=-1, keepdim=True).to(mix_stft.dtype))
+            )[..., None, :]
+        )
+    else:
+        # otherwise, we just multiply the targets spectrograms with mix phase
+        # we tacitly assume that we have magnitude estimates.
+        angle = atan2(mix_stft[..., 1], mix_stft[..., 0])[..., None]
+        nb_sources = targets_spectrograms.shape[-1]
+        y = torch.zeros(
+            mix_stft.shape + (nb_sources,), dtype=mix_stft.dtype, device=mix_stft.device
+        )
+        y[..., 0, :] = targets_spectrograms * torch.cos(angle)
+        y[..., 1, :] = targets_spectrograms * torch.sin(angle)
+
+    if residual:
+        # if required, adding an additional target as the mix minus
+        # available targets
+        y = torch.cat([y, mix_stft[..., None] - y.sum(dim=-1, keepdim=True)], dim=-1)
+
+    if iterations == 0:
+        return y
+
+    # we need to refine the estimates. Scales down the estimates for
+    # numerical stability
+    max_abs = torch.max(
+        torch.as_tensor(1.0, dtype=mix_stft.dtype, device=mix_stft.device),
+        torch.sqrt(_norm(mix_stft)).max() / scale_factor,
+    )
+
+    mix_stft = mix_stft / max_abs
+    y = y / max_abs
+
+    # call expectation maximization
+    y = expectation_maximization(y, mix_stft, iterations, eps=eps)[0]
+
+    # scale estimates up again
+    y = y * max_abs
+    return y
+
+
+def _covariance(y_j):
+    """
+    Compute the empirical covariance for a source.
+
+    Args:
+        y_j (Tensor): complex stft of the source.
+            [shape=(nb_frames, nb_bins, nb_channels, 2)].
+
+    Returns:
+        Cj (Tensor): [shape=(nb_frames, nb_bins, nb_channels, nb_channels, 2)]
+            just y_j * conj(y_j.T): empirical covariance for each TF bin.
+    """
+    (nb_frames, nb_bins, nb_channels) = y_j.shape[:-1]
+    Cj = torch.zeros(
+        (nb_frames, nb_bins, nb_channels, nb_channels, 2),
+        dtype=y_j.dtype,
+        device=y_j.device,
+    )
+    indices = torch.cartesian_prod(torch.arange(nb_channels), torch.arange(nb_channels))
+    for index in indices:
+        Cj[:, :, index[0], index[1], :] = _mul_add(
+            y_j[:, :, index[0], :],
+            _conj(y_j[:, :, index[1], :]),
+            Cj[:, :, index[0], index[1], :],
+        )
+    return Cj

lib/uvr5_pack/demucs/hdemucs.py

@@ -0,0 +1,782 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+This code contains the spectrogram and Hybrid version of Demucs.
+"""
+from copy import deepcopy
+import math
+import typing as tp
+import torch
+from torch import nn
+from torch.nn import functional as F
+from .filtering import wiener
+from .demucs import DConv, rescale_module
+from .states import capture_init
+from .spec import spectro, ispectro
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'constant', value: float = 0.):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen."""
+    x0 = x
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            extra_pad_right = min(padding_right, extra_pad)
+            extra_pad_left = extra_pad - extra_pad_right
+            paddings = (padding_left - extra_pad_left, padding_right - extra_pad_right)
+            x = F.pad(x, (extra_pad_left, extra_pad_right))
+    out = F.pad(x, paddings, mode, value)
+    assert out.shape[-1] == length + padding_left + padding_right
+    assert (out[..., padding_left: padding_left + length] == x0).all()
+    return out
+
+class ScaledEmbedding(nn.Module):
+    """
+    Boost learning rate for embeddings (with `scale`).
+    Also, can make embeddings continuous with `smooth`.
+    """
+    def __init__(self, num_embeddings: int, embedding_dim: int,
+                 scale: float = 10., smooth=False):
+        super().__init__()
+        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
+        if smooth:
+            weight = torch.cumsum(self.embedding.weight.data, dim=0)
+            # when summing gaussian, overscale raises as sqrt(n), so we nornalize by that.
+            weight = weight / torch.arange(1, num_embeddings + 1).to(weight).sqrt()[:, None]
+            self.embedding.weight.data[:] = weight
+        self.embedding.weight.data /= scale
+        self.scale = scale
+
+    @property
+    def weight(self):
+        return self.embedding.weight * self.scale
+
+    def forward(self, x):
+        out = self.embedding(x) * self.scale
+        return out
+
+
+class HEncLayer(nn.Module):
+    def __init__(self, chin, chout, kernel_size=8, stride=4, norm_groups=1, empty=False,
+                 freq=True, dconv=True, norm=True, context=0, dconv_kw={}, pad=True,
+                 rewrite=True):
+        """Encoder layer. This used both by the time and the frequency branch.
+
+        Args:
+            chin: number of input channels.
+            chout: number of output channels.
+            norm_groups: number of groups for group norm.
+            empty: used to make a layer with just the first conv. this is used
+                before merging the time and freq. branches.
+            freq: this is acting on frequencies.
+            dconv: insert DConv residual branches.
+            norm: use GroupNorm.
+            context: context size for the 1x1 conv.
+            dconv_kw: list of kwargs for the DConv class.
+            pad: pad the input. Padding is done so that the output size is
+                always the input size / stride.
+            rewrite: add 1x1 conv at the end of the layer.
+        """
+        super().__init__()
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm:
+            norm_fn = lambda d: nn.GroupNorm(norm_groups, d)  # noqa
+        if pad:
+            pad = kernel_size // 4
+        else:
+            pad = 0
+        klass = nn.Conv1d
+        self.freq = freq
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.empty = empty
+        self.norm = norm
+        self.pad = pad
+        if freq:
+            kernel_size = [kernel_size, 1]
+            stride = [stride, 1]
+            pad = [pad, 0]
+            klass = nn.Conv2d
+        self.conv = klass(chin, chout, kernel_size, stride, pad)
+        if self.empty:
+            return
+        self.norm1 = norm_fn(chout)
+        self.rewrite = None
+        if rewrite:
+            self.rewrite = klass(chout, 2 * chout, 1 + 2 * context, 1, context)
+            self.norm2 = norm_fn(2 * chout)
+
+        self.dconv = None
+        if dconv:
+            self.dconv = DConv(chout, **dconv_kw)
+
+    def forward(self, x, inject=None):
+        """
+        `inject` is used to inject the result from the time branch into the frequency branch,
+        when both have the same stride.
+        """
+        if not self.freq and x.dim() == 4:
+            B, C, Fr, T = x.shape
+            x = x.view(B, -1, T)
+
+        if not self.freq:
+            le = x.shape[-1]
+            if not le % self.stride == 0:
+                x = F.pad(x, (0, self.stride - (le % self.stride)))
+        y = self.conv(x)
+        if self.empty:
+            return y
+        if inject is not None:
+            assert inject.shape[-1] == y.shape[-1], (inject.shape, y.shape)
+            if inject.dim() == 3 and y.dim() == 4:
+                inject = inject[:, :, None]
+            y = y + inject
+        y = F.gelu(self.norm1(y))
+        if self.dconv:
+            if self.freq:
+                B, C, Fr, T = y.shape
+                y = y.permute(0, 2, 1, 3).reshape(-1, C, T)
+            y = self.dconv(y)
+            if self.freq:
+                y = y.view(B, Fr, C, T).permute(0, 2, 1, 3)
+        if self.rewrite:
+            z = self.norm2(self.rewrite(y))
+            z = F.glu(z, dim=1)
+        else:
+            z = y
+        return z
+
+
+class MultiWrap(nn.Module):
+    """
+    Takes one layer and replicate it N times. each replica will act
+    on a frequency band. All is done so that if the N replica have the same weights,
+    then this is exactly equivalent to applying the original module on all frequencies.
+
+    This is a bit over-engineered to avoid edge artifacts when splitting
+    the frequency bands, but it is possible the naive implementation would work as well...
+    """
+    def __init__(self, layer, split_ratios):
+        """
+        Args:
+            layer: module to clone, must be either HEncLayer or HDecLayer.
+            split_ratios: list of float indicating which ratio to keep for each band.
+        """
+        super().__init__()
+        self.split_ratios = split_ratios
+        self.layers = nn.ModuleList()
+        self.conv = isinstance(layer, HEncLayer)
+        assert not layer.norm
+        assert layer.freq
+        assert layer.pad
+        if not self.conv:
+            assert not layer.context_freq
+        for k in range(len(split_ratios) + 1):
+            lay = deepcopy(layer)
+            if self.conv:
+                lay.conv.padding = (0, 0)
+            else:
+                lay.pad = False
+            for m in lay.modules():
+                if hasattr(m, 'reset_parameters'):
+                    m.reset_parameters()
+            self.layers.append(lay)
+
+    def forward(self, x, skip=None, length=None):
+        B, C, Fr, T = x.shape
+
+        ratios = list(self.split_ratios) + [1]
+        start = 0
+        outs = []
+        for ratio, layer in zip(ratios, self.layers):
+            if self.conv:
+                pad = layer.kernel_size // 4
+                if ratio == 1:
+                    limit = Fr
+                    frames = -1
+                else:
+                    limit = int(round(Fr * ratio))
+                    le = limit - start
+                    if start == 0:
+                        le += pad
+                    frames = round((le - layer.kernel_size) / layer.stride + 1)
+                    limit = start + (frames - 1) * layer.stride + layer.kernel_size
+                    if start == 0:
+                        limit -= pad
+                assert limit - start > 0, (limit, start)
+                assert limit <= Fr, (limit, Fr)
+                y = x[:, :, start:limit, :]
+                if start == 0:
+                    y = F.pad(y, (0, 0, pad, 0))
+                if ratio == 1:
+                    y = F.pad(y, (0, 0, 0, pad))
+                outs.append(layer(y))
+                start = limit - layer.kernel_size + layer.stride
+            else:
+                if ratio == 1:
+                    limit = Fr
+                else:
+                    limit = int(round(Fr * ratio))
+                last = layer.last
+                layer.last = True
+
+                y = x[:, :, start:limit]
+                s = skip[:, :, start:limit]
+                out, _ = layer(y, s, None)
+                if outs:
+                    outs[-1][:, :, -layer.stride:] += (
+                        out[:, :, :layer.stride] - layer.conv_tr.bias.view(1, -1, 1, 1))
+                    out = out[:, :, layer.stride:]
+                if ratio == 1:
+                    out = out[:, :, :-layer.stride // 2, :]
+                if start == 0:
+                    out = out[:, :, layer.stride // 2:, :]
+                outs.append(out)
+                layer.last = last
+                start = limit
+        out = torch.cat(outs, dim=2)
+        if not self.conv and not last:
+            out = F.gelu(out)
+        if self.conv:
+            return out
+        else:
+            return out, None
+
+
+class HDecLayer(nn.Module):
+    def __init__(self, chin, chout, last=False, kernel_size=8, stride=4, norm_groups=1, empty=False,
+                 freq=True, dconv=True, norm=True, context=1, dconv_kw={}, pad=True,
+                 context_freq=True, rewrite=True):
+        """
+        Same as HEncLayer but for decoder. See `HEncLayer` for documentation.
+        """
+        super().__init__()
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm:
+            norm_fn = lambda d: nn.GroupNorm(norm_groups, d)  # noqa
+        if pad:
+            pad = kernel_size // 4
+        else:
+            pad = 0
+        self.pad = pad
+        self.last = last
+        self.freq = freq
+        self.chin = chin
+        self.empty = empty
+        self.stride = stride
+        self.kernel_size = kernel_size
+        self.norm = norm
+        self.context_freq = context_freq
+        klass = nn.Conv1d
+        klass_tr = nn.ConvTranspose1d
+        if freq:
+            kernel_size = [kernel_size, 1]
+            stride = [stride, 1]
+            klass = nn.Conv2d
+            klass_tr = nn.ConvTranspose2d
+        self.conv_tr = klass_tr(chin, chout, kernel_size, stride)
+        self.norm2 = norm_fn(chout)
+        if self.empty:
+            return
+        self.rewrite = None
+        if rewrite:
+            if context_freq:
+                self.rewrite = klass(chin, 2 * chin, 1 + 2 * context, 1, context)
+            else:
+                self.rewrite = klass(chin, 2 * chin, [1, 1 + 2 * context], 1,
+                                     [0, context])
+            self.norm1 = norm_fn(2 * chin)
+
+        self.dconv = None
+        if dconv:
+            self.dconv = DConv(chin, **dconv_kw)
+
+    def forward(self, x, skip, length):
+        if self.freq and x.dim() == 3:
+            B, C, T = x.shape
+            x = x.view(B, self.chin, -1, T)
+
+        if not self.empty:
+            x = x + skip
+
+            if self.rewrite:
+                y = F.glu(self.norm1(self.rewrite(x)), dim=1)
+            else:
+                y = x
+            if self.dconv:
+                if self.freq:
+                    B, C, Fr, T = y.shape
+                    y = y.permute(0, 2, 1, 3).reshape(-1, C, T)
+                y = self.dconv(y)
+                if self.freq:
+                    y = y.view(B, Fr, C, T).permute(0, 2, 1, 3)
+        else:
+            y = x
+            assert skip is None
+        z = self.norm2(self.conv_tr(y))
+        if self.freq:
+            if self.pad:
+                z = z[..., self.pad:-self.pad, :]
+        else:
+            z = z[..., self.pad:self.pad + length]
+            assert z.shape[-1] == length, (z.shape[-1], length)
+        if not self.last:
+            z = F.gelu(z)
+        return z, y
+
+
+class HDemucs(nn.Module):
+    """
+    Spectrogram and hybrid Demucs model.
+    The spectrogram model has the same structure as Demucs, except the first few layers are over the
+    frequency axis, until there is only 1 frequency, and then it moves to time convolutions.
+    Frequency layers can still access information across time steps thanks to the DConv residual.
+
+    Hybrid model have a parallel time branch. At some layer, the time branch has the same stride
+    as the frequency branch and then the two are combined. The opposite happens in the decoder.
+
+    Models can either use naive iSTFT from masking, Wiener filtering ([Ulhih et al. 2017]),
+    or complex as channels (CaC) [Choi et al. 2020]. Wiener filtering is based on
+    Open Unmix implementation [Stoter et al. 2019].
+
+    The loss is always on the temporal domain, by backpropagating through the above
+    output methods and iSTFT. This allows to define hybrid models nicely. However, this breaks
+    a bit Wiener filtering, as doing more iteration at test time will change the spectrogram
+    contribution, without changing the one from the waveform, which will lead to worse performance.
+    I tried using the residual option in OpenUnmix Wiener implementation, but it didn't improve.
+    CaC on the other hand provides similar performance for hybrid, and works naturally with
+    hybrid models.
+
+    This model also uses frequency embeddings are used to improve efficiency on convolutions
+    over the freq. axis, following [Isik et al. 2020] (https://arxiv.org/pdf/2008.04470.pdf).
+
+    Unlike classic Demucs, there is no resampling here, and normalization is always applied.
+    """
+    @capture_init
+    def __init__(self,
+                 sources,
+                 # Channels
+                 audio_channels=2,
+                 channels=48,
+                 channels_time=None,
+                 growth=2,
+                 # STFT
+                 nfft=4096,
+                 wiener_iters=0,
+                 end_iters=0,
+                 wiener_residual=False,
+                 cac=True,
+                 # Main structure
+                 depth=6,
+                 rewrite=True,
+                 hybrid=True,
+                 hybrid_old=False,
+                 # Frequency branch
+                 multi_freqs=None,
+                 multi_freqs_depth=2,
+                 freq_emb=0.2,
+                 emb_scale=10,
+                 emb_smooth=True,
+                 # Convolutions
+                 kernel_size=8,
+                 time_stride=2,
+                 stride=4,
+                 context=1,
+                 context_enc=0,
+                 # Normalization
+                 norm_starts=4,
+                 norm_groups=4,
+                 # DConv residual branch
+                 dconv_mode=1,
+                 dconv_depth=2,
+                 dconv_comp=4,
+                 dconv_attn=4,
+                 dconv_lstm=4,
+                 dconv_init=1e-4,
+                 # Weight init
+                 rescale=0.1,
+                 # Metadata
+                 samplerate=44100,
+                 segment=4 * 10):
+        
+        """
+        Args:
+            sources (list[str]): list of source names.
+            audio_channels (int): input/output audio channels.
+            channels (int): initial number of hidden channels.
+            channels_time: if not None, use a different `channels` value for the time branch.
+            growth: increase the number of hidden channels by this factor at each layer.
+            nfft: number of fft bins. Note that changing this require careful computation of
+                various shape parameters and will not work out of the box for hybrid models.
+            wiener_iters: when using Wiener filtering, number of iterations at test time.
+            end_iters: same but at train time. For a hybrid model, must be equal to `wiener_iters`.
+            wiener_residual: add residual source before wiener filtering.
+            cac: uses complex as channels, i.e. complex numbers are 2 channels each
+                in input and output. no further processing is done before ISTFT.
+            depth (int): number of layers in the encoder and in the decoder.
+            rewrite (bool): add 1x1 convolution to each layer.
+            hybrid (bool): make a hybrid time/frequency domain, otherwise frequency only.
+            hybrid_old: some models trained for MDX had a padding bug. This replicates
+                this bug to avoid retraining them.
+            multi_freqs: list of frequency ratios for splitting frequency bands with `MultiWrap`.
+            multi_freqs_depth: how many layers to wrap with `MultiWrap`. Only the outermost
+                layers will be wrapped.
+            freq_emb: add frequency embedding after the first frequency layer if > 0,
+                the actual value controls the weight of the embedding.
+            emb_scale: equivalent to scaling the embedding learning rate
+            emb_smooth: initialize the embedding with a smooth one (with respect to frequencies).
+            kernel_size: kernel_size for encoder and decoder layers.
+            stride: stride for encoder and decoder layers.
+            time_stride: stride for the final time layer, after the merge.
+            context: context for 1x1 conv in the decoder.
+            context_enc: context for 1x1 conv in the encoder.
+            norm_starts: layer at which group norm starts being used.
+                decoder layers are numbered in reverse order.
+            norm_groups: number of groups for group norm.
+            dconv_mode: if 1: dconv in encoder only, 2: decoder only, 3: both.
+            dconv_depth: depth of residual DConv branch.
+            dconv_comp: compression of DConv branch.
+            dconv_attn: adds attention layers in DConv branch starting at this layer.
+            dconv_lstm: adds a LSTM layer in DConv branch starting at this layer.
+            dconv_init: initial scale for the DConv branch LayerScale.
+            rescale: weight recaling trick
+
+        """
+        super().__init__()
+        
+        self.cac = cac
+        self.wiener_residual = wiener_residual
+        self.audio_channels = audio_channels
+        self.sources = sources
+        self.kernel_size = kernel_size
+        self.context = context
+        self.stride = stride
+        self.depth = depth
+        self.channels = channels
+        self.samplerate = samplerate
+        self.segment = segment
+
+        self.nfft = nfft
+        self.hop_length = nfft // 4
+        self.wiener_iters = wiener_iters
+        self.end_iters = end_iters
+        self.freq_emb = None
+        self.hybrid = hybrid
+        self.hybrid_old = hybrid_old
+        if hybrid_old:
+            assert hybrid, "hybrid_old must come with hybrid=True"
+        if hybrid:
+            assert wiener_iters == end_iters
+
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+
+        if hybrid:
+            self.tencoder = nn.ModuleList()
+            self.tdecoder = nn.ModuleList()
+
+        chin = audio_channels
+        chin_z = chin  # number of channels for the freq branch
+        if self.cac:
+            chin_z *= 2
+        chout = channels_time or channels
+        chout_z = channels
+        freqs = nfft // 2
+
+        for index in range(depth):
+            lstm = index >= dconv_lstm
+            attn = index >= dconv_attn
+            norm = index >= norm_starts
+            freq = freqs > 1
+            stri = stride
+            ker = kernel_size
+            if not freq:
+                assert freqs == 1
+                ker = time_stride * 2
+                stri = time_stride
+
+            pad = True
+            last_freq = False
+            if freq and freqs <= kernel_size:
+                ker = freqs
+                pad = False
+                last_freq = True
+
+            kw = {
+                'kernel_size': ker,
+                'stride': stri,
+                'freq': freq,
+                'pad': pad,
+                'norm': norm,
+                'rewrite': rewrite,
+                'norm_groups': norm_groups,
+                'dconv_kw': {
+                    'lstm': lstm,
+                    'attn': attn,
+                    'depth': dconv_depth,
+                    'compress': dconv_comp,
+                    'init': dconv_init,
+                    'gelu': True,
+                }
+            }
+            kwt = dict(kw)
+            kwt['freq'] = 0
+            kwt['kernel_size'] = kernel_size
+            kwt['stride'] = stride
+            kwt['pad'] = True
+            kw_dec = dict(kw)
+            multi = False
+            if multi_freqs and index < multi_freqs_depth:
+                multi = True
+                kw_dec['context_freq'] = False
+
+            if last_freq:
+                chout_z = max(chout, chout_z)
+                chout = chout_z
+
+            enc = HEncLayer(chin_z, chout_z,
+                            dconv=dconv_mode & 1, context=context_enc, **kw)
+            if hybrid and freq:
+                tenc = HEncLayer(chin, chout, dconv=dconv_mode & 1, context=context_enc,
+                                 empty=last_freq, **kwt)
+                self.tencoder.append(tenc)
+
+            if multi:
+                enc = MultiWrap(enc, multi_freqs)
+            self.encoder.append(enc)
+            if index == 0:
+                chin = self.audio_channels * len(self.sources)
+                chin_z = chin
+                if self.cac:
+                    chin_z *= 2
+            dec = HDecLayer(chout_z, chin_z, dconv=dconv_mode & 2,
+                            last=index == 0, context=context, **kw_dec)
+            if multi:
+                dec = MultiWrap(dec, multi_freqs)
+            if hybrid and freq:
+                tdec = HDecLayer(chout, chin, dconv=dconv_mode & 2, empty=last_freq,
+                                 last=index == 0, context=context, **kwt)
+                self.tdecoder.insert(0, tdec)
+            self.decoder.insert(0, dec)
+
+            chin = chout
+            chin_z = chout_z
+            chout = int(growth * chout)
+            chout_z = int(growth * chout_z)
+            if freq:
+                if freqs <= kernel_size:
+                    freqs = 1
+                else:
+                    freqs //= stride
+            if index == 0 and freq_emb:
+                self.freq_emb = ScaledEmbedding(
+                    freqs, chin_z, smooth=emb_smooth, scale=emb_scale)
+                self.freq_emb_scale = freq_emb
+
+        if rescale:
+            rescale_module(self, reference=rescale)
+
+    def _spec(self, x):
+        hl = self.hop_length
+        nfft = self.nfft
+        x0 = x  # noqa
+
+        if self.hybrid:
+            # We re-pad the signal in order to keep the property
+            # that the size of the output is exactly the size of the input
+            # divided by the stride (here hop_length), when divisible.
+            # This is achieved by padding by 1/4th of the kernel size (here nfft).
+            # which is not supported by torch.stft.
+            # Having all convolution operations follow this convention allow to easily
+            # align the time and frequency branches later on.
+            assert hl == nfft // 4
+            le = int(math.ceil(x.shape[-1] / hl))
+            pad = hl // 2 * 3
+            if not self.hybrid_old:
+                x = pad1d(x, (pad, pad + le * hl - x.shape[-1]), mode='reflect')
+            else:
+                x = pad1d(x, (pad, pad + le * hl - x.shape[-1]))
+
+        z = spectro(x, nfft, hl)[..., :-1, :]
+        if self.hybrid:
+            assert z.shape[-1] == le + 4, (z.shape, x.shape, le)
+            z = z[..., 2:2+le]
+        return z
+
+    def _ispec(self, z, length=None, scale=0):
+        hl = self.hop_length // (4 ** scale)
+        z = F.pad(z, (0, 0, 0, 1))
+        if self.hybrid:
+            z = F.pad(z, (2, 2))
+            pad = hl // 2 * 3
+            if not self.hybrid_old:
+                le = hl * int(math.ceil(length / hl)) + 2 * pad
+            else:
+                le = hl * int(math.ceil(length / hl))
+            x = ispectro(z, hl, length=le)
+            if not self.hybrid_old:
+                x = x[..., pad:pad + length]
+            else:
+                x = x[..., :length]
+        else:
+            x = ispectro(z, hl, length)
+        return x
+
+    def _magnitude(self, z):
+        # return the magnitude of the spectrogram, except when cac is True,
+        # in which case we just move the complex dimension to the channel one.
+        if self.cac:
+            B, C, Fr, T = z.shape
+            m = torch.view_as_real(z).permute(0, 1, 4, 2, 3)
+            m = m.reshape(B, C * 2, Fr, T)
+        else:
+            m = z.abs()
+        return m
+
+    def _mask(self, z, m):
+        # Apply masking given the mixture spectrogram `z` and the estimated mask `m`.
+        # If `cac` is True, `m` is actually a full spectrogram and `z` is ignored.
+        niters = self.wiener_iters
+        if self.cac:
+            B, S, C, Fr, T = m.shape
+            out = m.view(B, S, -1, 2, Fr, T).permute(0, 1, 2, 4, 5, 3)
+            out = torch.view_as_complex(out.contiguous())
+            return out
+        if self.training:
+            niters = self.end_iters
+        if niters < 0:
+            z = z[:, None]
+            return z / (1e-8 + z.abs()) * m
+        else:
+            return self._wiener(m, z, niters)
+
+    def _wiener(self, mag_out, mix_stft, niters):
+        # apply wiener filtering from OpenUnmix.
+        init = mix_stft.dtype
+        wiener_win_len = 300
+        residual = self.wiener_residual
+
+        B, S, C, Fq, T = mag_out.shape
+        mag_out = mag_out.permute(0, 4, 3, 2, 1)
+        mix_stft = torch.view_as_real(mix_stft.permute(0, 3, 2, 1))
+
+        outs = []
+        for sample in range(B):
+            pos = 0
+            out = []
+            for pos in range(0, T, wiener_win_len):
+                frame = slice(pos, pos + wiener_win_len)
+                z_out = wiener(
+                    mag_out[sample, frame], mix_stft[sample, frame], niters,
+                    residual=residual)
+                out.append(z_out.transpose(-1, -2))
+            outs.append(torch.cat(out, dim=0))
+        out = torch.view_as_complex(torch.stack(outs, 0))
+        out = out.permute(0, 4, 3, 2, 1).contiguous()
+        if residual:
+            out = out[:, :-1]
+        assert list(out.shape) == [B, S, C, Fq, T]
+        return out.to(init)
+
+    def forward(self, mix):
+        x = mix
+        length = x.shape[-1]
+
+        z = self._spec(mix)
+        mag = self._magnitude(z)
+        x = mag
+
+        B, C, Fq, T = x.shape
+
+        # unlike previous Demucs, we always normalize because it is easier.
+        mean = x.mean(dim=(1, 2, 3), keepdim=True)
+        std = x.std(dim=(1, 2, 3), keepdim=True)
+        x = (x - mean) / (1e-5 + std)
+        # x will be the freq. branch input.
+
+        if self.hybrid:
+            # Prepare the time branch input.
+            xt = mix
+            meant = xt.mean(dim=(1, 2), keepdim=True)
+            stdt = xt.std(dim=(1, 2), keepdim=True)
+            xt = (xt - meant) / (1e-5 + stdt)
+
+        # okay, this is a giant mess I know...
+        saved = []  # skip connections, freq.
+        saved_t = []  # skip connections, time.
+        lengths = []  # saved lengths to properly remove padding, freq branch.
+        lengths_t = []  # saved lengths for time branch.
+        for idx, encode in enumerate(self.encoder):
+            lengths.append(x.shape[-1])
+            inject = None
+            if self.hybrid and idx < len(self.tencoder):
+                # we have not yet merged branches.
+                lengths_t.append(xt.shape[-1])
+                tenc = self.tencoder[idx]
+                xt = tenc(xt)
+                if not tenc.empty:
+                    # save for skip connection
+                    saved_t.append(xt)
+                else:
+                    # tenc contains just the first conv., so that now time and freq.
+                    # branches have the same shape and can be merged.
+                    inject = xt
+            x = encode(x, inject)
+            if idx == 0 and self.freq_emb is not None:
+                # add frequency embedding to allow for non equivariant convolutions
+                # over the frequency axis.
+                frs = torch.arange(x.shape[-2], device=x.device)
+                emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x)
+                x = x + self.freq_emb_scale * emb
+
+            saved.append(x)
+
+        x = torch.zeros_like(x)
+        if self.hybrid:
+            xt = torch.zeros_like(x)
+        # initialize everything to zero (signal will go through u-net skips).
+
+        for idx, decode in enumerate(self.decoder):
+            skip = saved.pop(-1)
+            x, pre = decode(x, skip, lengths.pop(-1))
+            # `pre` contains the output just before final transposed convolution,
+            # which is used when the freq. and time branch separate.
+
+            if self.hybrid:
+                offset = self.depth - len(self.tdecoder)
+            if self.hybrid and idx >= offset:
+                tdec = self.tdecoder[idx - offset]
+                length_t = lengths_t.pop(-1)
+                if tdec.empty:
+                    assert pre.shape[2] == 1, pre.shape
+                    pre = pre[:, :, 0]
+                    xt, _ = tdec(pre, None, length_t)
+                else:
+                    skip = saved_t.pop(-1)
+                    xt, _ = tdec(xt, skip, length_t)
+
+        # Let's make sure we used all stored skip connections.
+        assert len(saved) == 0
+        assert len(lengths_t) == 0
+        assert len(saved_t) == 0
+
+        S = len(self.sources)
+        x = x.view(B, S, -1, Fq, T)
+        x = x * std[:, None] + mean[:, None]
+
+        zout = self._mask(z, x)
+        x = self._ispec(zout, length)
+
+        if self.hybrid:
+            xt = xt.view(B, S, -1, length)
+            xt = xt * stdt[:, None] + meant[:, None]
+            x = xt + x
+        return x
+
+

lib/uvr5_pack/demucs/htdemucs.py

@@ -0,0 +1,648 @@
+# Copyright (c) Meta, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# First author is Simon Rouard.
+"""
+This code contains the spectrogram and Hybrid version of Demucs.
+"""
+import math
+
+from .filtering import wiener
+import torch
+from torch import nn
+from torch.nn import functional as F
+from fractions import Fraction
+from einops import rearrange
+
+from .transformer import CrossTransformerEncoder
+
+from .demucs import rescale_module
+from .states import capture_init
+from .spec import spectro, ispectro
+from .hdemucs import pad1d, ScaledEmbedding, HEncLayer, MultiWrap, HDecLayer
+
+
+class HTDemucs(nn.Module):
+    """
+    Spectrogram and hybrid Demucs model.
+    The spectrogram model has the same structure as Demucs, except the first few layers are over the
+    frequency axis, until there is only 1 frequency, and then it moves to time convolutions.
+    Frequency layers can still access information across time steps thanks to the DConv residual.
+
+    Hybrid model have a parallel time branch. At some layer, the time branch has the same stride
+    as the frequency branch and then the two are combined. The opposite happens in the decoder.
+
+    Models can either use naive iSTFT from masking, Wiener filtering ([Ulhih et al. 2017]),
+    or complex as channels (CaC) [Choi et al. 2020]. Wiener filtering is based on
+    Open Unmix implementation [Stoter et al. 2019].
+
+    The loss is always on the temporal domain, by backpropagating through the above
+    output methods and iSTFT. This allows to define hybrid models nicely. However, this breaks
+    a bit Wiener filtering, as doing more iteration at test time will change the spectrogram
+    contribution, without changing the one from the waveform, which will lead to worse performance.
+    I tried using the residual option in OpenUnmix Wiener implementation, but it didn't improve.
+    CaC on the other hand provides similar performance for hybrid, and works naturally with
+    hybrid models.
+
+    This model also uses frequency embeddings are used to improve efficiency on convolutions
+    over the freq. axis, following [Isik et al. 2020] (https://arxiv.org/pdf/2008.04470.pdf).
+
+    Unlike classic Demucs, there is no resampling here, and normalization is always applied.
+    """
+
+    @capture_init
+    def __init__(
+        self,
+        sources,
+        # Channels
+        audio_channels=2,
+        channels=48,
+        channels_time=None,
+        growth=2,
+        # STFT
+        nfft=4096,
+        wiener_iters=0,
+        end_iters=0,
+        wiener_residual=False,
+        cac=True,
+        # Main structure
+        depth=4,
+        rewrite=True,
+        # Frequency branch
+        multi_freqs=None,
+        multi_freqs_depth=3,
+        freq_emb=0.2,
+        emb_scale=10,
+        emb_smooth=True,
+        # Convolutions
+        kernel_size=8,
+        time_stride=2,
+        stride=4,
+        context=1,
+        context_enc=0,
+        # Normalization
+        norm_starts=4,
+        norm_groups=4,
+        # DConv residual branch
+        dconv_mode=1,
+        dconv_depth=2,
+        dconv_comp=8,
+        dconv_init=1e-3,
+        # Before the Transformer
+        bottom_channels=0,
+        # Transformer
+        t_layers=5,
+        t_emb="sin",
+        t_hidden_scale=4.0,
+        t_heads=8,
+        t_dropout=0.0,
+        t_max_positions=10000,
+        t_norm_in=True,
+        t_norm_in_group=False,
+        t_group_norm=False,
+        t_norm_first=True,
+        t_norm_out=True,
+        t_max_period=10000.0,
+        t_weight_decay=0.0,
+        t_lr=None,
+        t_layer_scale=True,
+        t_gelu=True,
+        t_weight_pos_embed=1.0,
+        t_sin_random_shift=0,
+        t_cape_mean_normalize=True,
+        t_cape_augment=True,
+        t_cape_glob_loc_scale=[5000.0, 1.0, 1.4],
+        t_sparse_self_attn=False,
+        t_sparse_cross_attn=False,
+        t_mask_type="diag",
+        t_mask_random_seed=42,
+        t_sparse_attn_window=500,
+        t_global_window=100,
+        t_sparsity=0.95,
+        t_auto_sparsity=False,
+        # ------ Particuliar parameters
+        t_cross_first=False,
+        # Weight init
+        rescale=0.1,
+        # Metadata
+        samplerate=44100,
+        segment=10,
+        use_train_segment=True,
+    ):
+        """
+        Args:
+            sources (list[str]): list of source names.
+            audio_channels (int): input/output audio channels.
+            channels (int): initial number of hidden channels.
+            channels_time: if not None, use a different `channels` value for the time branch.
+            growth: increase the number of hidden channels by this factor at each layer.
+            nfft: number of fft bins. Note that changing this require careful computation of
+                various shape parameters and will not work out of the box for hybrid models.
+            wiener_iters: when using Wiener filtering, number of iterations at test time.
+            end_iters: same but at train time. For a hybrid model, must be equal to `wiener_iters`.
+            wiener_residual: add residual source before wiener filtering.
+            cac: uses complex as channels, i.e. complex numbers are 2 channels each
+                in input and output. no further processing is done before ISTFT.
+            depth (int): number of layers in the encoder and in the decoder.
+            rewrite (bool): add 1x1 convolution to each layer.
+            multi_freqs: list of frequency ratios for splitting frequency bands with `MultiWrap`.
+            multi_freqs_depth: how many layers to wrap with `MultiWrap`. Only the outermost
+                layers will be wrapped.
+            freq_emb: add frequency embedding after the first frequency layer if > 0,
+                the actual value controls the weight of the embedding.
+            emb_scale: equivalent to scaling the embedding learning rate
+            emb_smooth: initialize the embedding with a smooth one (with respect to frequencies).
+            kernel_size: kernel_size for encoder and decoder layers.
+            stride: stride for encoder and decoder layers.
+            time_stride: stride for the final time layer, after the merge.
+            context: context for 1x1 conv in the decoder.
+            context_enc: context for 1x1 conv in the encoder.
+            norm_starts: layer at which group norm starts being used.
+                decoder layers are numbered in reverse order.
+            norm_groups: number of groups for group norm.
+            dconv_mode: if 1: dconv in encoder only, 2: decoder only, 3: both.
+            dconv_depth: depth of residual DConv branch.
+            dconv_comp: compression of DConv branch.
+            dconv_attn: adds attention layers in DConv branch starting at this layer.
+            dconv_lstm: adds a LSTM layer in DConv branch starting at this layer.
+            dconv_init: initial scale for the DConv branch LayerScale.
+            bottom_channels: if >0 it adds a linear layer (1x1 Conv) before and after the
+                transformer in order to change the number of channels
+            t_layers: number of layers in each branch (waveform and spec) of the transformer
+            t_emb: "sin", "cape" or "scaled"
+            t_hidden_scale: the hidden scale of the Feedforward parts of the transformer
+                for instance if C = 384 (the number of channels in the transformer) and
+                t_hidden_scale = 4.0 then the intermediate layer of the FFN has dimension
+                384 * 4 = 1536
+            t_heads: number of heads for the transformer
+            t_dropout: dropout in the transformer
+            t_max_positions: max_positions for the "scaled" positional embedding, only
+                useful if t_emb="scaled"
+            t_norm_in: (bool) norm before addinf positional embedding and getting into the
+                transformer layers
+            t_norm_in_group: (bool) if True while t_norm_in=True, the norm is on all the
+                timesteps (GroupNorm with group=1)
+            t_group_norm: (bool) if True, the norms of the Encoder Layers are on all the
+                timesteps (GroupNorm with group=1)
+            t_norm_first: (bool) if True the norm is before the attention and before the FFN
+            t_norm_out: (bool) if True, there is a GroupNorm (group=1) at the end of each layer
+            t_max_period: (float) denominator in the sinusoidal embedding expression
+            t_weight_decay: (float) weight decay for the transformer
+            t_lr: (float) specific learning rate for the transformer
+            t_layer_scale: (bool) Layer Scale for the transformer
+            t_gelu: (bool) activations of the transformer are GeLU if True, ReLU else
+            t_weight_pos_embed: (float) weighting of the positional embedding
+            t_cape_mean_normalize: (bool) if t_emb="cape", normalisation of positional embeddings
+                see: https://arxiv.org/abs/2106.03143
+            t_cape_augment: (bool) if t_emb="cape", must be True during training and False
+                during the inference, see: https://arxiv.org/abs/2106.03143
+            t_cape_glob_loc_scale: (list of 3 floats) if t_emb="cape", CAPE parameters
+                see: https://arxiv.org/abs/2106.03143
+            t_sparse_self_attn: (bool) if True, the self attentions are sparse
+            t_sparse_cross_attn: (bool) if True, the cross-attentions are sparse (don't use it
+                unless you designed really specific masks)
+            t_mask_type: (str) can be "diag", "jmask", "random", "global" or any combination
+                with '_' between: i.e. "diag_jmask_random" (note that this is permutation
+                invariant i.e. "diag_jmask_random" is equivalent to "jmask_random_diag")
+            t_mask_random_seed: (int) if "random" is in t_mask_type, controls the seed
+                that generated the random part of the mask
+            t_sparse_attn_window: (int) if "diag" is in t_mask_type, for a query (i), and
+                a key (j), the mask is True id |i-j|<=t_sparse_attn_window
+            t_global_window: (int) if "global" is in t_mask_type, mask[:t_global_window, :]
+                and mask[:, :t_global_window] will be True
+            t_sparsity: (float) if "random" is in t_mask_type, t_sparsity is the sparsity
+                level of the random part of the mask.
+            t_cross_first: (bool) if True cross attention is the first layer of the
+                transformer (False seems to be better)
+            rescale: weight rescaling trick
+            use_train_segment: (bool) if True, the actual size that is used during the
+                training is used during inference.
+        """
+        super().__init__()
+        self.cac = cac
+        self.wiener_residual = wiener_residual
+        self.audio_channels = audio_channels
+        self.sources = sources
+        self.kernel_size = kernel_size
+        self.context = context
+        self.stride = stride
+        self.depth = depth
+        self.bottom_channels = bottom_channels
+        self.channels = channels
+        self.samplerate = samplerate
+        self.segment = segment
+        self.use_train_segment = use_train_segment
+        self.nfft = nfft
+        self.hop_length = nfft // 4
+        self.wiener_iters = wiener_iters
+        self.end_iters = end_iters
+        self.freq_emb = None
+        assert wiener_iters == end_iters
+
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+
+        self.tencoder = nn.ModuleList()
+        self.tdecoder = nn.ModuleList()
+
+        chin = audio_channels
+        chin_z = chin  # number of channels for the freq branch
+        if self.cac:
+            chin_z *= 2
+        chout = channels_time or channels
+        chout_z = channels
+        freqs = nfft // 2
+
+        for index in range(depth):
+            norm = index >= norm_starts
+            freq = freqs > 1
+            stri = stride
+            ker = kernel_size
+            if not freq:
+                assert freqs == 1
+                ker = time_stride * 2
+                stri = time_stride
+
+            pad = True
+            last_freq = False
+            if freq and freqs <= kernel_size:
+                ker = freqs
+                pad = False
+                last_freq = True
+
+            kw = {
+                "kernel_size": ker,
+                "stride": stri,
+                "freq": freq,
+                "pad": pad,
+                "norm": norm,
+                "rewrite": rewrite,
+                "norm_groups": norm_groups,
+                "dconv_kw": {
+                    "depth": dconv_depth,
+                    "compress": dconv_comp,
+                    "init": dconv_init,
+                    "gelu": True,
+                },
+            }
+            kwt = dict(kw)
+            kwt["freq"] = 0
+            kwt["kernel_size"] = kernel_size
+            kwt["stride"] = stride
+            kwt["pad"] = True
+            kw_dec = dict(kw)
+            multi = False
+            if multi_freqs and index < multi_freqs_depth:
+                multi = True
+                kw_dec["context_freq"] = False
+
+            if last_freq:
+                chout_z = max(chout, chout_z)
+                chout = chout_z
+
+            enc = HEncLayer(
+                chin_z, chout_z, dconv=dconv_mode & 1, context=context_enc, **kw
+            )
+            if freq:
+                tenc = HEncLayer(
+                    chin,
+                    chout,
+                    dconv=dconv_mode & 1,
+                    context=context_enc,
+                    empty=last_freq,
+                    **kwt
+                )
+                self.tencoder.append(tenc)
+
+            if multi:
+                enc = MultiWrap(enc, multi_freqs)
+            self.encoder.append(enc)
+            if index == 0:
+                chin = self.audio_channels * len(self.sources)
+                chin_z = chin
+                if self.cac:
+                    chin_z *= 2
+            dec = HDecLayer(
+                chout_z,
+                chin_z,
+                dconv=dconv_mode & 2,
+                last=index == 0,
+                context=context,
+                **kw_dec
+            )
+            if multi:
+                dec = MultiWrap(dec, multi_freqs)
+            if freq:
+                tdec = HDecLayer(
+                    chout,
+                    chin,
+                    dconv=dconv_mode & 2,
+                    empty=last_freq,
+                    last=index == 0,
+                    context=context,
+                    **kwt
+                )
+                self.tdecoder.insert(0, tdec)
+            self.decoder.insert(0, dec)
+
+            chin = chout
+            chin_z = chout_z
+            chout = int(growth * chout)
+            chout_z = int(growth * chout_z)
+            if freq:
+                if freqs <= kernel_size:
+                    freqs = 1
+                else:
+                    freqs //= stride
+            if index == 0 and freq_emb:
+                self.freq_emb = ScaledEmbedding(
+                    freqs, chin_z, smooth=emb_smooth, scale=emb_scale
+                )
+                self.freq_emb_scale = freq_emb
+
+        if rescale:
+            rescale_module(self, reference=rescale)
+
+        transformer_channels = channels * growth ** (depth - 1)
+        if bottom_channels:
+            self.channel_upsampler = nn.Conv1d(transformer_channels, bottom_channels, 1)
+            self.channel_downsampler = nn.Conv1d(
+                bottom_channels, transformer_channels, 1
+            )
+            self.channel_upsampler_t = nn.Conv1d(
+                transformer_channels, bottom_channels, 1
+            )
+            self.channel_downsampler_t = nn.Conv1d(
+                bottom_channels, transformer_channels, 1
+            )
+
+            transformer_channels = bottom_channels
+
+        if t_layers > 0:
+            self.crosstransformer = CrossTransformerEncoder(
+                dim=transformer_channels,
+                emb=t_emb,
+                hidden_scale=t_hidden_scale,
+                num_heads=t_heads,
+                num_layers=t_layers,
+                cross_first=t_cross_first,
+                dropout=t_dropout,
+                max_positions=t_max_positions,
+                norm_in=t_norm_in,
+                norm_in_group=t_norm_in_group,
+                group_norm=t_group_norm,
+                norm_first=t_norm_first,
+                norm_out=t_norm_out,
+                max_period=t_max_period,
+                weight_decay=t_weight_decay,
+                lr=t_lr,
+                layer_scale=t_layer_scale,
+                gelu=t_gelu,
+                sin_random_shift=t_sin_random_shift,
+                weight_pos_embed=t_weight_pos_embed,
+                cape_mean_normalize=t_cape_mean_normalize,
+                cape_augment=t_cape_augment,
+                cape_glob_loc_scale=t_cape_glob_loc_scale,
+                sparse_self_attn=t_sparse_self_attn,
+                sparse_cross_attn=t_sparse_cross_attn,
+                mask_type=t_mask_type,
+                mask_random_seed=t_mask_random_seed,
+                sparse_attn_window=t_sparse_attn_window,
+                global_window=t_global_window,
+                sparsity=t_sparsity,
+                auto_sparsity=t_auto_sparsity,
+            )
+        else:
+            self.crosstransformer = None
+
+    def _spec(self, x):
+        hl = self.hop_length
+        nfft = self.nfft
+        x0 = x  # noqa
+
+        # We re-pad the signal in order to keep the property
+        # that the size of the output is exactly the size of the input
+        # divided by the stride (here hop_length), when divisible.
+        # This is achieved by padding by 1/4th of the kernel size (here nfft).
+        # which is not supported by torch.stft.
+        # Having all convolution operations follow this convention allow to easily
+        # align the time and frequency branches later on.
+        assert hl == nfft // 4
+        le = int(math.ceil(x.shape[-1] / hl))
+        pad = hl // 2 * 3
+        x = pad1d(x, (pad, pad + le * hl - x.shape[-1]), mode="reflect")
+
+        z = spectro(x, nfft, hl)[..., :-1, :]
+        assert z.shape[-1] == le + 4, (z.shape, x.shape, le)
+        z = z[..., 2: 2 + le]
+        return z
+
+    def _ispec(self, z, length=None, scale=0):
+        hl = self.hop_length // (4**scale)
+        z = F.pad(z, (0, 0, 0, 1))
+        z = F.pad(z, (2, 2))
+        pad = hl // 2 * 3
+        le = hl * int(math.ceil(length / hl)) + 2 * pad
+        x = ispectro(z, hl, length=le)
+        x = x[..., pad: pad + length]
+        return x
+
+    def _magnitude(self, z):
+        # return the magnitude of the spectrogram, except when cac is True,
+        # in which case we just move the complex dimension to the channel one.
+        if self.cac:
+            B, C, Fr, T = z.shape
+            m = torch.view_as_real(z).permute(0, 1, 4, 2, 3)
+            m = m.reshape(B, C * 2, Fr, T)
+        else:
+            m = z.abs()
+        return m
+
+    def _mask(self, z, m):
+        # Apply masking given the mixture spectrogram `z` and the estimated mask `m`.
+        # If `cac` is True, `m` is actually a full spectrogram and `z` is ignored.
+        niters = self.wiener_iters
+        if self.cac:
+            B, S, C, Fr, T = m.shape
+            out = m.view(B, S, -1, 2, Fr, T).permute(0, 1, 2, 4, 5, 3)
+            out = torch.view_as_complex(out.contiguous())
+            return out
+        if self.training:
+            niters = self.end_iters
+        if niters < 0:
+            z = z[:, None]
+            return z / (1e-8 + z.abs()) * m
+        else:
+            return self._wiener(m, z, niters)
+
+    def _wiener(self, mag_out, mix_stft, niters):
+        # apply wiener filtering from OpenUnmix.
+        init = mix_stft.dtype
+        wiener_win_len = 300
+        residual = self.wiener_residual
+
+        B, S, C, Fq, T = mag_out.shape
+        mag_out = mag_out.permute(0, 4, 3, 2, 1)
+        mix_stft = torch.view_as_real(mix_stft.permute(0, 3, 2, 1))
+
+        outs = []
+        for sample in range(B):
+            pos = 0
+            out = []
+            for pos in range(0, T, wiener_win_len):
+                frame = slice(pos, pos + wiener_win_len)
+                z_out = wiener(
+                    mag_out[sample, frame],
+                    mix_stft[sample, frame],
+                    niters,
+                    residual=residual,
+                )
+                out.append(z_out.transpose(-1, -2))
+            outs.append(torch.cat(out, dim=0))
+        out = torch.view_as_complex(torch.stack(outs, 0))
+        out = out.permute(0, 4, 3, 2, 1).contiguous()
+        if residual:
+            out = out[:, :-1]
+        assert list(out.shape) == [B, S, C, Fq, T]
+        return out.to(init)
+
+    def valid_length(self, length: int):
+        """
+        Return a length that is appropriate for evaluation.
+        In our case, always return the training length, unless
+        it is smaller than the given length, in which case this
+        raises an error.
+        """
+        if not self.use_train_segment:
+            return length
+        training_length = int(self.segment * self.samplerate)
+        if training_length < length:
+            raise ValueError(
+                    f"Given length {length} is longer than "
+                    f"training length {training_length}")
+        return training_length
+
+    def forward(self, mix):
+        length = mix.shape[-1]
+        length_pre_pad = None
+        if self.use_train_segment:
+            if self.training:
+                self.segment = Fraction(mix.shape[-1], self.samplerate)
+            else:
+                training_length = int(self.segment * self.samplerate)
+                if mix.shape[-1] < training_length:
+                    length_pre_pad = mix.shape[-1]
+                    mix = F.pad(mix, (0, training_length - length_pre_pad))
+        z = self._spec(mix)
+        mag = self._magnitude(z)
+        x = mag
+
+        B, C, Fq, T = x.shape
+
+        # unlike previous Demucs, we always normalize because it is easier.
+        mean = x.mean(dim=(1, 2, 3), keepdim=True)
+        std = x.std(dim=(1, 2, 3), keepdim=True)
+        x = (x - mean) / (1e-5 + std)
+        # x will be the freq. branch input.
+
+        # Prepare the time branch input.
+        xt = mix
+        meant = xt.mean(dim=(1, 2), keepdim=True)
+        stdt = xt.std(dim=(1, 2), keepdim=True)
+        xt = (xt - meant) / (1e-5 + stdt)
+
+        # okay, this is a giant mess I know...
+        saved = []  # skip connections, freq.
+        saved_t = []  # skip connections, time.
+        lengths = []  # saved lengths to properly remove padding, freq branch.
+        lengths_t = []  # saved lengths for time branch.
+        for idx, encode in enumerate(self.encoder):
+            lengths.append(x.shape[-1])
+            inject = None
+            if idx < len(self.tencoder):
+                # we have not yet merged branches.
+                lengths_t.append(xt.shape[-1])
+                tenc = self.tencoder[idx]
+                xt = tenc(xt)
+                if not tenc.empty:
+                    # save for skip connection
+                    saved_t.append(xt)
+                else:
+                    # tenc contains just the first conv., so that now time and freq.
+                    # branches have the same shape and can be merged.
+                    inject = xt
+            x = encode(x, inject)
+            if idx == 0 and self.freq_emb is not None:
+                # add frequency embedding to allow for non equivariant convolutions
+                # over the frequency axis.
+                frs = torch.arange(x.shape[-2], device=x.device)
+                emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x)
+                x = x + self.freq_emb_scale * emb
+
+            saved.append(x)
+        if self.crosstransformer:
+            if self.bottom_channels:
+                b, c, f, t = x.shape
+                x = rearrange(x, "b c f t-> b c (f t)")
+                x = self.channel_upsampler(x)
+                x = rearrange(x, "b c (f t)-> b c f t", f=f)
+                xt = self.channel_upsampler_t(xt)
+
+            x, xt = self.crosstransformer(x, xt)
+
+            if self.bottom_channels:
+                x = rearrange(x, "b c f t-> b c (f t)")
+                x = self.channel_downsampler(x)
+                x = rearrange(x, "b c (f t)-> b c f t", f=f)
+                xt = self.channel_downsampler_t(xt)
+
+        for idx, decode in enumerate(self.decoder):
+            skip = saved.pop(-1)
+            x, pre = decode(x, skip, lengths.pop(-1))
+            # `pre` contains the output just before final transposed convolution,
+            # which is used when the freq. and time branch separate.
+
+            offset = self.depth - len(self.tdecoder)
+            if idx >= offset:
+                tdec = self.tdecoder[idx - offset]
+                length_t = lengths_t.pop(-1)
+                if tdec.empty:
+                    assert pre.shape[2] == 1, pre.shape
+                    pre = pre[:, :, 0]
+                    xt, _ = tdec(pre, None, length_t)
+                else:
+                    skip = saved_t.pop(-1)
+                    xt, _ = tdec(xt, skip, length_t)
+
+        # Let's make sure we used all stored skip connections.
+        assert len(saved) == 0
+        assert len(lengths_t) == 0
+        assert len(saved_t) == 0
+
+        S = len(self.sources)
+        x = x.view(B, S, -1, Fq, T)
+        x = x * std[:, None] + mean[:, None]
+
+        zout = self._mask(z, x)
+        if self.use_train_segment:
+            if self.training:
+                x = self._ispec(zout, length)
+            else:
+                x = self._ispec(zout, training_length)
+        else:
+            x = self._ispec(zout, length)
+
+        if self.use_train_segment:
+            if self.training:
+                xt = xt.view(B, S, -1, length)
+            else:
+                xt = xt.view(B, S, -1, training_length)
+        else:
+            xt = xt.view(B, S, -1, length)
+        xt = xt * stdt[:, None] + meant[:, None]
+        x = xt + x
+        if length_pre_pad:
+            x = x[..., :length_pre_pad]
+        return x

lib/uvr5_pack/demucs/model.py

@@ -0,0 +1,218 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+import torch as th
+from torch import nn
+
+from .utils import capture_init, center_trim
+
+
+class BLSTM(nn.Module):
+    def __init__(self, dim, layers=1):
+        super().__init__()
+        self.lstm = nn.LSTM(bidirectional=True, num_layers=layers, hidden_size=dim, input_size=dim)
+        self.linear = nn.Linear(2 * dim, dim)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        x = self.lstm(x)[0]
+        x = self.linear(x)
+        x = x.permute(1, 2, 0)
+        return x
+
+
+def rescale_conv(conv, reference):
+    std = conv.weight.std().detach()
+    scale = (std / reference)**0.5
+    conv.weight.data /= scale
+    if conv.bias is not None:
+        conv.bias.data /= scale
+
+
+def rescale_module(module, reference):
+    for sub in module.modules():
+        if isinstance(sub, (nn.Conv1d, nn.ConvTranspose1d)):
+            rescale_conv(sub, reference)
+
+
+def upsample(x, stride):
+    """
+    Linear upsampling, the output will be `stride` times longer.
+    """
+    batch, channels, time = x.size()
+    weight = th.arange(stride, device=x.device, dtype=th.float) / stride
+    x = x.view(batch, channels, time, 1)
+    out = x[..., :-1, :] * (1 - weight) + x[..., 1:, :] * weight
+    return out.reshape(batch, channels, -1)
+
+
+def downsample(x, stride):
+    """
+    Downsample x by decimation.
+    """
+    return x[:, :, ::stride]
+
+
+class Demucs(nn.Module):
+    @capture_init
+    def __init__(self,
+                 sources=4,
+                 audio_channels=2,
+                 channels=64,
+                 depth=6,
+                 rewrite=True,
+                 glu=True,
+                 upsample=False,
+                 rescale=0.1,
+                 kernel_size=8,
+                 stride=4,
+                 growth=2.,
+                 lstm_layers=2,
+                 context=3,
+                 samplerate=44100):
+        """
+        Args:
+            sources (int): number of sources to separate
+            audio_channels (int): stereo or mono
+            channels (int): first convolution channels
+            depth (int): number of encoder/decoder layers
+            rewrite (bool): add 1x1 convolution to each encoder layer
+                and a convolution to each decoder layer.
+                For the decoder layer, `context` gives the kernel size.
+            glu (bool): use glu instead of ReLU
+            upsample (bool): use linear upsampling with convolutions
+                Wave-U-Net style, instead of transposed convolutions
+            rescale (int): rescale initial weights of convolutions
+                to get their standard deviation closer to `rescale`
+            kernel_size (int): kernel size for convolutions
+            stride (int): stride for convolutions
+            growth (float): multiply (resp divide) number of channels by that
+                for each layer of the encoder (resp decoder)
+            lstm_layers (int): number of lstm layers, 0 = no lstm
+            context (int): kernel size of the convolution in the
+                decoder before the transposed convolution. If > 1,
+                will provide some context from neighboring time
+                steps.
+        """
+
+        super().__init__()
+        self.audio_channels = audio_channels
+        self.sources = sources
+        self.kernel_size = kernel_size
+        self.context = context
+        self.stride = stride
+        self.depth = depth
+        self.upsample = upsample
+        self.channels = channels
+        self.samplerate = samplerate
+
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+
+        self.final = None
+        if upsample:
+            self.final = nn.Conv1d(channels + audio_channels, sources * audio_channels, 1)
+            stride = 1
+
+        if glu:
+            activation = nn.GLU(dim=1)
+            ch_scale = 2
+        else:
+            activation = nn.ReLU()
+            ch_scale = 1
+        in_channels = audio_channels
+        for index in range(depth):
+            encode = []
+            encode += [nn.Conv1d(in_channels, channels, kernel_size, stride), nn.ReLU()]
+            if rewrite:
+                encode += [nn.Conv1d(channels, ch_scale * channels, 1), activation]
+            self.encoder.append(nn.Sequential(*encode))
+
+            decode = []
+            if index > 0:
+                out_channels = in_channels
+            else:
+                if upsample:
+                    out_channels = channels
+                else:
+                    out_channels = sources * audio_channels
+            if rewrite:
+                decode += [nn.Conv1d(channels, ch_scale * channels, context), activation]
+            if upsample:
+                decode += [
+                    nn.Conv1d(channels, out_channels, kernel_size, stride=1),
+                ]
+            else:
+                decode += [nn.ConvTranspose1d(channels, out_channels, kernel_size, stride)]
+            if index > 0:
+                decode.append(nn.ReLU())
+            self.decoder.insert(0, nn.Sequential(*decode))
+            in_channels = channels
+            channels = int(growth * channels)
+
+        channels = in_channels
+
+        if lstm_layers:
+            self.lstm = BLSTM(channels, lstm_layers)
+        else:
+            self.lstm = None
+
+        if rescale:
+            rescale_module(self, reference=rescale)
+
+    def valid_length(self, length):
+        """
+        Return the nearest valid length to use with the model so that
+        there is no time steps left over in a convolutions, e.g. for all
+        layers, size of the input - kernel_size % stride = 0.
+
+        If the mixture has a valid length, the estimated sources
+        will have exactly the same length when context = 1. If context > 1,
+        the two signals can be center trimmed to match.
+
+        For training, extracts should have a valid length.For evaluation
+        on full tracks we recommend passing `pad = True` to :method:`forward`.
+        """
+        for _ in range(self.depth):
+            if self.upsample:
+                length = math.ceil(length / self.stride) + self.kernel_size - 1
+            else:
+                length = math.ceil((length - self.kernel_size) / self.stride) + 1
+            length = max(1, length)
+            length += self.context - 1
+        for _ in range(self.depth):
+            if self.upsample:
+                length = length * self.stride + self.kernel_size - 1
+            else:
+                length = (length - 1) * self.stride + self.kernel_size
+
+        return int(length)
+
+    def forward(self, mix):
+        x = mix
+        saved = [x]
+        for encode in self.encoder:
+            x = encode(x)
+            saved.append(x)
+            if self.upsample:
+                x = downsample(x, self.stride)
+        if self.lstm:
+            x = self.lstm(x)
+        for decode in self.decoder:
+            if self.upsample:
+                x = upsample(x, stride=self.stride)
+            skip = center_trim(saved.pop(-1), x)
+            x = x + skip
+            x = decode(x)
+        if self.final:
+            skip = center_trim(saved.pop(-1), x)
+            x = th.cat([x, skip], dim=1)
+            x = self.final(x)
+
+        x = x.view(x.size(0), self.sources, self.audio_channels, x.size(-1))
+        return x

lib/uvr5_pack/demucs/model_v2.py

@@ -0,0 +1,218 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+import julius
+from torch import nn
+from .tasnet_v2 import ConvTasNet
+
+from .utils import capture_init, center_trim
+
+
+class BLSTM(nn.Module):
+    def __init__(self, dim, layers=1):
+        super().__init__()
+        self.lstm = nn.LSTM(bidirectional=True, num_layers=layers, hidden_size=dim, input_size=dim)
+        self.linear = nn.Linear(2 * dim, dim)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        x = self.lstm(x)[0]
+        x = self.linear(x)
+        x = x.permute(1, 2, 0)
+        return x
+
+
+def rescale_conv(conv, reference):
+    std = conv.weight.std().detach()
+    scale = (std / reference)**0.5
+    conv.weight.data /= scale
+    if conv.bias is not None:
+        conv.bias.data /= scale
+
+
+def rescale_module(module, reference):
+    for sub in module.modules():
+        if isinstance(sub, (nn.Conv1d, nn.ConvTranspose1d)):
+            rescale_conv(sub, reference)
+
+def auto_load_demucs_model_v2(sources, demucs_model_name):
+    
+    if '48' in demucs_model_name:
+        channels=48
+    elif 'unittest' in demucs_model_name:
+        channels=4
+    else:
+        channels=64
+    
+    if 'tasnet' in demucs_model_name:
+        init_demucs_model = ConvTasNet(sources, X=10)
+    else:
+        init_demucs_model = Demucs(sources, channels=channels)
+        
+    return init_demucs_model
+
+class Demucs(nn.Module):
+    @capture_init
+    def __init__(self,
+                 sources,
+                 audio_channels=2,
+                 channels=64,
+                 depth=6,
+                 rewrite=True,
+                 glu=True,
+                 rescale=0.1,
+                 resample=True,
+                 kernel_size=8,
+                 stride=4,
+                 growth=2.,
+                 lstm_layers=2,
+                 context=3,
+                 normalize=False,
+                 samplerate=44100,
+                 segment_length=4 * 10 * 44100):
+        """
+        Args:
+            sources (list[str]): list of source names
+            audio_channels (int): stereo or mono
+            channels (int): first convolution channels
+            depth (int): number of encoder/decoder layers
+            rewrite (bool): add 1x1 convolution to each encoder layer
+                and a convolution to each decoder layer.
+                For the decoder layer, `context` gives the kernel size.
+            glu (bool): use glu instead of ReLU
+            resample_input (bool): upsample x2 the input and downsample /2 the output.
+            rescale (int): rescale initial weights of convolutions
+                to get their standard deviation closer to `rescale`
+            kernel_size (int): kernel size for convolutions
+            stride (int): stride for convolutions
+            growth (float): multiply (resp divide) number of channels by that
+                for each layer of the encoder (resp decoder)
+            lstm_layers (int): number of lstm layers, 0 = no lstm
+            context (int): kernel size of the convolution in the
+                decoder before the transposed convolution. If > 1,
+                will provide some context from neighboring time
+                steps.
+            samplerate (int): stored as meta information for easing
+                future evaluations of the model.
+            segment_length (int): stored as meta information for easing
+                future evaluations of the model. Length of the segments on which
+                the model was trained.
+        """
+
+        super().__init__()
+        self.audio_channels = audio_channels
+        self.sources = sources
+        self.kernel_size = kernel_size
+        self.context = context
+        self.stride = stride
+        self.depth = depth
+        self.resample = resample
+        self.channels = channels
+        self.normalize = normalize
+        self.samplerate = samplerate
+        self.segment_length = segment_length
+
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+
+        if glu:
+            activation = nn.GLU(dim=1)
+            ch_scale = 2
+        else:
+            activation = nn.ReLU()
+            ch_scale = 1
+        in_channels = audio_channels
+        for index in range(depth):
+            encode = []
+            encode += [nn.Conv1d(in_channels, channels, kernel_size, stride), nn.ReLU()]
+            if rewrite:
+                encode += [nn.Conv1d(channels, ch_scale * channels, 1), activation]
+            self.encoder.append(nn.Sequential(*encode))
+
+            decode = []
+            if index > 0:
+                out_channels = in_channels
+            else:
+                out_channels = len(self.sources) * audio_channels
+            if rewrite:
+                decode += [nn.Conv1d(channels, ch_scale * channels, context), activation]
+            decode += [nn.ConvTranspose1d(channels, out_channels, kernel_size, stride)]
+            if index > 0:
+                decode.append(nn.ReLU())
+            self.decoder.insert(0, nn.Sequential(*decode))
+            in_channels = channels
+            channels = int(growth * channels)
+
+        channels = in_channels
+
+        if lstm_layers:
+            self.lstm = BLSTM(channels, lstm_layers)
+        else:
+            self.lstm = None
+
+        if rescale:
+            rescale_module(self, reference=rescale)
+
+    def valid_length(self, length):
+        """
+        Return the nearest valid length to use with the model so that
+        there is no time steps left over in a convolutions, e.g. for all
+        layers, size of the input - kernel_size % stride = 0.
+
+        If the mixture has a valid length, the estimated sources
+        will have exactly the same length when context = 1. If context > 1,
+        the two signals can be center trimmed to match.
+
+        For training, extracts should have a valid length.For evaluation
+        on full tracks we recommend passing `pad = True` to :method:`forward`.
+        """
+        if self.resample:
+            length *= 2
+        for _ in range(self.depth):
+            length = math.ceil((length - self.kernel_size) / self.stride) + 1
+            length = max(1, length)
+            length += self.context - 1
+        for _ in range(self.depth):
+            length = (length - 1) * self.stride + self.kernel_size
+
+        if self.resample:
+            length = math.ceil(length / 2)
+        return int(length)
+
+    def forward(self, mix):
+        x = mix
+
+        if self.normalize:
+            mono = mix.mean(dim=1, keepdim=True)
+            mean = mono.mean(dim=-1, keepdim=True)
+            std = mono.std(dim=-1, keepdim=True)
+        else:
+            mean = 0
+            std = 1
+
+        x = (x - mean) / (1e-5 + std)
+
+        if self.resample:
+            x = julius.resample_frac(x, 1, 2)
+
+        saved = []
+        for encode in self.encoder:
+            x = encode(x)
+            saved.append(x)
+        if self.lstm:
+            x = self.lstm(x)
+        for decode in self.decoder:
+            skip = center_trim(saved.pop(-1), x)
+            x = x + skip
+            x = decode(x)
+
+        if self.resample:
+            x = julius.resample_frac(x, 2, 1)
+        x = x * std + mean
+        x = x.view(x.size(0), len(self.sources), self.audio_channels, x.size(-1))
+        return x

lib/uvr5_pack/demucs/pretrained.py

@@ -0,0 +1,180 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Loading pretrained models.
+"""
+
+import logging
+from pathlib import Path
+import typing as tp
+
+#from dora.log import fatal
+
+import logging
+
+from diffq import DiffQuantizer
+import torch.hub
+
+from .model import Demucs
+from .tasnet_v2 import ConvTasNet
+from .utils import set_state
+
+from .hdemucs import HDemucs
+from .repo import RemoteRepo, LocalRepo, ModelOnlyRepo, BagOnlyRepo, AnyModelRepo, ModelLoadingError  # noqa
+
+logger = logging.getLogger(__name__)
+ROOT_URL = "https://dl.fbaipublicfiles.com/demucs/mdx_final/"
+REMOTE_ROOT = Path(__file__).parent / 'remote'
+
+SOURCES = ["drums", "bass", "other", "vocals"]
+
+
+def demucs_unittest():
+    model = HDemucs(channels=4, sources=SOURCES)
+    return model
+
+
+def add_model_flags(parser):
+    group = parser.add_mutually_exclusive_group(required=False)
+    group.add_argument("-s", "--sig", help="Locally trained XP signature.")
+    group.add_argument("-n", "--name", default="mdx_extra_q",
+                       help="Pretrained model name or signature. Default is mdx_extra_q.")
+    parser.add_argument("--repo", type=Path,
+                        help="Folder containing all pre-trained models for use with -n.")
+
+
+def _parse_remote_files(remote_file_list) -> tp.Dict[str, str]:
+    root: str = ''
+    models: tp.Dict[str, str] = {}
+    for line in remote_file_list.read_text().split('\n'):
+        line = line.strip()
+        if line.startswith('#'):
+            continue
+        elif line.startswith('root:'):
+            root = line.split(':', 1)[1].strip()
+        else:
+            sig = line.split('-', 1)[0]
+            assert sig not in models
+            models[sig] = ROOT_URL + root + line
+    return models
+
+def get_model(name: str,
+              repo: tp.Optional[Path] = None):
+    """`name` must be a bag of models name or a pretrained signature
+    from the remote AWS model repo or the specified local repo if `repo` is not None.
+    """
+    if name == 'demucs_unittest':
+        return demucs_unittest()
+    model_repo: ModelOnlyRepo
+    if repo is None:
+        models = _parse_remote_files(REMOTE_ROOT / 'files.txt')
+        model_repo = RemoteRepo(models)
+        bag_repo = BagOnlyRepo(REMOTE_ROOT, model_repo)
+    else:
+        if not repo.is_dir():
+            fatal(f"{repo} must exist and be a directory.")
+        model_repo = LocalRepo(repo)
+        bag_repo = BagOnlyRepo(repo, model_repo)
+    any_repo = AnyModelRepo(model_repo, bag_repo)
+    model = any_repo.get_model(name)
+    model.eval()
+    return model
+
+def get_model_from_args(args):
+    """
+    Load local model package or pre-trained model.
+    """
+    return get_model(name=args.name, repo=args.repo)
+
+logger = logging.getLogger(__name__)
+ROOT = "https://dl.fbaipublicfiles.com/demucs/v3.0/"
+
+PRETRAINED_MODELS = {
+    'demucs': 'e07c671f',
+    'demucs48_hq': '28a1282c',
+    'demucs_extra': '3646af93',
+    'demucs_quantized': '07afea75',
+    'tasnet': 'beb46fac',
+    'tasnet_extra': 'df3777b2',
+    'demucs_unittest': '09ebc15f',
+}
+
+SOURCES = ["drums", "bass", "other", "vocals"]
+
+
+def get_url(name):
+    sig = PRETRAINED_MODELS[name]
+    return ROOT + name + "-" + sig[:8] + ".th"
+
+def is_pretrained(name):
+    return name in PRETRAINED_MODELS
+
+
+def load_pretrained(name):
+    if name == "demucs":
+        return demucs(pretrained=True)
+    elif name == "demucs48_hq":
+        return demucs(pretrained=True, hq=True, channels=48)
+    elif name == "demucs_extra":
+        return demucs(pretrained=True, extra=True)
+    elif name == "demucs_quantized":
+        return demucs(pretrained=True, quantized=True)
+    elif name == "demucs_unittest":
+        return demucs_unittest(pretrained=True)
+    elif name == "tasnet":
+        return tasnet(pretrained=True)
+    elif name == "tasnet_extra":
+        return tasnet(pretrained=True, extra=True)
+    else:
+        raise ValueError(f"Invalid pretrained name {name}")
+
+
+def _load_state(name, model, quantizer=None):
+    url = get_url(name)
+    state = torch.hub.load_state_dict_from_url(url, map_location='cpu', check_hash=True)
+    set_state(model, quantizer, state)
+    if quantizer:
+        quantizer.detach()
+
+
+def demucs_unittest(pretrained=True):
+    model = Demucs(channels=4, sources=SOURCES)
+    if pretrained:
+        _load_state('demucs_unittest', model)
+    return model
+
+
+def demucs(pretrained=True, extra=False, quantized=False, hq=False, channels=64):
+    if not pretrained and (extra or quantized or hq):
+        raise ValueError("if extra or quantized is True, pretrained must be True.")
+    model = Demucs(sources=SOURCES, channels=channels)
+    if pretrained:
+        name = 'demucs'
+        if channels != 64:
+            name += str(channels)
+        quantizer = None
+        if sum([extra, quantized, hq]) > 1:
+            raise ValueError("Only one of extra, quantized, hq, can be True.")
+        if quantized:
+            quantizer = DiffQuantizer(model, group_size=8, min_size=1)
+            name += '_quantized'
+        if extra:
+            name += '_extra'
+        if hq:
+            name += '_hq'
+        _load_state(name, model, quantizer)
+    return model
+
+
+def tasnet(pretrained=True, extra=False):
+    if not pretrained and extra:
+        raise ValueError("if extra is True, pretrained must be True.")
+    model = ConvTasNet(X=10, sources=SOURCES)
+    if pretrained:
+        name = 'tasnet'
+        if extra:
+            name = 'tasnet_extra'
+        _load_state(name, model)
+    return model

lib/uvr5_pack/demucs/repo.py

@@ -0,0 +1,148 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Represents a model repository, including pre-trained models and bags of models.
+A repo can either be the main remote repository stored in AWS, or a local repository
+with your own models.
+"""
+
+from hashlib import sha256
+from pathlib import Path
+import typing as tp
+
+import torch
+import yaml
+
+from .apply import BagOfModels, Model
+from .states import load_model
+
+
+AnyModel = tp.Union[Model, BagOfModels]
+
+
+class ModelLoadingError(RuntimeError):
+    pass
+
+
+def check_checksum(path: Path, checksum: str):
+    sha = sha256()
+    with open(path, 'rb') as file:
+        while True:
+            buf = file.read(2**20)
+            if not buf:
+                break
+            sha.update(buf)
+    actual_checksum = sha.hexdigest()[:len(checksum)]
+    if actual_checksum != checksum:
+        raise ModelLoadingError(f'Invalid checksum for file {path}, '
+                                f'expected {checksum} but got {actual_checksum}')
+
+class ModelOnlyRepo:
+    """Base class for all model only repos.
+    """
+    def has_model(self, sig: str) -> bool:
+        raise NotImplementedError()
+
+    def get_model(self, sig: str) -> Model:
+        raise NotImplementedError()
+
+
+class RemoteRepo(ModelOnlyRepo):
+    def __init__(self, models: tp.Dict[str, str]):
+        self._models = models
+
+    def has_model(self, sig: str) -> bool:
+        return sig in self._models
+
+    def get_model(self, sig: str) -> Model:
+        try:
+            url = self._models[sig]
+        except KeyError:
+            raise ModelLoadingError(f'Could not find a pre-trained model with signature {sig}.')
+        pkg = torch.hub.load_state_dict_from_url(url, map_location='cpu', check_hash=True)
+        return load_model(pkg)
+
+
+class LocalRepo(ModelOnlyRepo):
+    def __init__(self, root: Path):
+        self.root = root
+        self.scan()
+
+    def scan(self):
+        self._models = {}
+        self._checksums = {}
+        for file in self.root.iterdir():
+            if file.suffix == '.th':
+                if '-' in file.stem:
+                    xp_sig, checksum = file.stem.split('-')
+                    self._checksums[xp_sig] = checksum
+                else:
+                    xp_sig = file.stem
+                if xp_sig in self._models:
+                    print('Whats xp? ', xp_sig)
+                    raise ModelLoadingError(
+                        f'Duplicate pre-trained model exist for signature {xp_sig}. '
+                        'Please delete all but one.')
+                self._models[xp_sig] = file
+
+    def has_model(self, sig: str) -> bool:
+        return sig in self._models
+
+    def get_model(self, sig: str) -> Model:
+        try:
+            file = self._models[sig]
+        except KeyError:
+            raise ModelLoadingError(f'Could not find pre-trained model with signature {sig}.')
+        if sig in self._checksums:
+            check_checksum(file, self._checksums[sig])
+        return load_model(file)
+
+
+class BagOnlyRepo:
+    """Handles only YAML files containing bag of models, leaving the actual
+    model loading to some Repo.
+    """
+    def __init__(self, root: Path, model_repo: ModelOnlyRepo):
+        self.root = root
+        self.model_repo = model_repo
+        self.scan()
+
+    def scan(self):
+        self._bags = {}
+        for file in self.root.iterdir():
+            if file.suffix == '.yaml':
+                self._bags[file.stem] = file
+
+    def has_model(self, name: str) -> bool:
+        return name in self._bags
+
+    def get_model(self, name: str) -> BagOfModels:
+        try:
+            yaml_file = self._bags[name]
+        except KeyError:
+            raise ModelLoadingError(f'{name} is neither a single pre-trained model or '
+                                    'a bag of models.')
+        bag = yaml.safe_load(open(yaml_file))
+        signatures = bag['models']
+        models = [self.model_repo.get_model(sig) for sig in signatures]
+        weights = bag.get('weights')
+        segment = bag.get('segment')
+        return BagOfModels(models, weights, segment)
+
+
+class AnyModelRepo:
+    def __init__(self, model_repo: ModelOnlyRepo, bag_repo: BagOnlyRepo):
+        self.model_repo = model_repo
+        self.bag_repo = bag_repo
+
+    def has_model(self, name_or_sig: str) -> bool:
+        return self.model_repo.has_model(name_or_sig) or self.bag_repo.has_model(name_or_sig)
+
+    def get_model(self, name_or_sig: str) -> AnyModel:
+        print('name_or_sig: ', name_or_sig)
+        if self.model_repo.has_model(name_or_sig):
+            return self.model_repo.get_model(name_or_sig)
+        else:
+            return self.bag_repo.get_model(name_or_sig)

lib/uvr5_pack/demucs/spec.py

@@ -0,0 +1,41 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Conveniance wrapper to perform STFT and iSTFT"""
+
+import torch as th
+
+
+def spectro(x, n_fft=512, hop_length=None, pad=0):
+    *other, length = x.shape
+    x = x.reshape(-1, length)
+    z = th.stft(x,
+                n_fft * (1 + pad),
+                hop_length or n_fft // 4,
+                window=th.hann_window(n_fft).to(x),
+                win_length=n_fft,
+                normalized=True,
+                center=True,
+                return_complex=True,
+                pad_mode='reflect')
+    _, freqs, frame = z.shape
+    return z.view(*other, freqs, frame)
+
+
+def ispectro(z, hop_length=None, length=None, pad=0):
+    *other, freqs, frames = z.shape
+    n_fft = 2 * freqs - 2
+    z = z.view(-1, freqs, frames)
+    win_length = n_fft // (1 + pad)
+    x = th.istft(z,
+                 n_fft,
+                 hop_length,
+                 window=th.hann_window(win_length).to(z.real),
+                 win_length=win_length,
+                 normalized=True,
+                 length=length,
+                 center=True)
+    _, length = x.shape
+    return x.view(*other, length)

lib/uvr5_pack/demucs/states.py

@@ -0,0 +1,148 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Utilities to save and load models.
+"""
+from contextlib import contextmanager
+
+import functools
+import hashlib
+import inspect
+import io
+from pathlib import Path
+import warnings
+
+from omegaconf import OmegaConf
+from diffq import DiffQuantizer, UniformQuantizer, restore_quantized_state
+import torch
+
+
+def get_quantizer(model, args, optimizer=None):
+    """Return the quantizer given the XP quantization args."""
+    quantizer = None
+    if args.diffq:
+        quantizer = DiffQuantizer(
+            model, min_size=args.min_size, group_size=args.group_size)
+        if optimizer is not None:
+            quantizer.setup_optimizer(optimizer)
+    elif args.qat:
+        quantizer = UniformQuantizer(
+                model, bits=args.qat, min_size=args.min_size)
+    return quantizer
+
+
+def load_model(path_or_package, strict=False):
+    """Load a model from the given serialized model, either given as a dict (already loaded)
+    or a path to a file on disk."""
+    if isinstance(path_or_package, dict):
+        package = path_or_package
+    elif isinstance(path_or_package, (str, Path)):
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            path = path_or_package
+            package = torch.load(path, 'cpu')
+    else:
+        raise ValueError(f"Invalid type for {path_or_package}.")
+
+    klass = package["klass"]
+    args = package["args"]
+    kwargs = package["kwargs"]
+
+    if strict:
+        model = klass(*args, **kwargs)
+    else:
+        sig = inspect.signature(klass)
+        for key in list(kwargs):
+            if key not in sig.parameters:
+                warnings.warn("Dropping inexistant parameter " + key)
+                del kwargs[key]
+        model = klass(*args, **kwargs)
+
+    state = package["state"]
+
+    set_state(model, state)
+    return model
+
+
+def get_state(model, quantizer, half=False):
+    """Get the state from a model, potentially with quantization applied.
+    If `half` is True, model are stored as half precision, which shouldn't impact performance
+    but half the state size."""
+    if quantizer is None:
+        dtype = torch.half if half else None
+        state = {k: p.data.to(device='cpu', dtype=dtype) for k, p in model.state_dict().items()}
+    else:
+        state = quantizer.get_quantized_state()
+        state['__quantized'] = True
+    return state
+
+
+def set_state(model, state, quantizer=None):
+    """Set the state on a given model."""
+    if state.get('__quantized'):
+        if quantizer is not None:
+            quantizer.restore_quantized_state(model, state['quantized'])
+        else:
+            restore_quantized_state(model, state)
+    else:
+        model.load_state_dict(state)
+    return state
+
+
+def save_with_checksum(content, path):
+    """Save the given value on disk, along with a sha256 hash.
+    Should be used with the output of either `serialize_model` or `get_state`."""
+    buf = io.BytesIO()
+    torch.save(content, buf)
+    sig = hashlib.sha256(buf.getvalue()).hexdigest()[:8]
+
+    path = path.parent / (path.stem + "-" + sig + path.suffix)
+    path.write_bytes(buf.getvalue())
+
+
+def serialize_model(model, training_args, quantizer=None, half=True):
+    args, kwargs = model._init_args_kwargs
+    klass = model.__class__
+
+    state = get_state(model, quantizer, half)
+    return {
+        'klass': klass,
+        'args': args,
+        'kwargs': kwargs,
+        'state': state,
+        'training_args': OmegaConf.to_container(training_args, resolve=True),
+    }
+
+
+def copy_state(state):
+    return {k: v.cpu().clone() for k, v in state.items()}
+
+
+@contextmanager
+def swap_state(model, state):
+    """
+    Context manager that swaps the state of a model, e.g:
+
+        # model is in old state
+        with swap_state(model, new_state):
+            # model in new state
+        # model back to old state
+    """
+    old_state = copy_state(model.state_dict())
+    model.load_state_dict(state, strict=False)
+    try:
+        yield
+    finally:
+        model.load_state_dict(old_state)
+
+
+def capture_init(init):
+    @functools.wraps(init)
+    def __init__(self, *args, **kwargs):
+        self._init_args_kwargs = (args, kwargs)
+        init(self, *args, **kwargs)
+
+    return __init__

lib/uvr5_pack/demucs/tasnet.py

@@ -0,0 +1,447 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Created on 2018/12
+# Author: Kaituo XU
+# Modified on 2019/11 by Alexandre Defossez, added support for multiple output channels
+# Here is the original license:
+# The MIT License (MIT)
+#
+# Copyright (c) 2018 Kaituo XU
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import capture_init
+
+EPS = 1e-8
+
+
+def overlap_and_add(signal, frame_step):
+    outer_dimensions = signal.size()[:-2]
+    frames, frame_length = signal.size()[-2:]
+
+    subframe_length = math.gcd(frame_length, frame_step)  # gcd=Greatest Common Divisor
+    subframe_step = frame_step // subframe_length
+    subframes_per_frame = frame_length // subframe_length
+    output_size = frame_step * (frames - 1) + frame_length
+    output_subframes = output_size // subframe_length
+
+    subframe_signal = signal.view(*outer_dimensions, -1, subframe_length)
+
+    frame = torch.arange(0, output_subframes,
+                         device=signal.device).unfold(0, subframes_per_frame, subframe_step)
+    frame = frame.long()  # signal may in GPU or CPU
+    frame = frame.contiguous().view(-1)
+
+    result = signal.new_zeros(*outer_dimensions, output_subframes, subframe_length)
+    result.index_add_(-2, frame, subframe_signal)
+    result = result.view(*outer_dimensions, -1)
+    return result
+
+
+class ConvTasNet(nn.Module):
+    @capture_init
+    def __init__(self,
+                 N=256,
+                 L=20,
+                 B=256,
+                 H=512,
+                 P=3,
+                 X=8,
+                 R=4,
+                 C=4,
+                 audio_channels=1,
+                 samplerate=44100,
+                 norm_type="gLN",
+                 causal=False,
+                 mask_nonlinear='relu'):
+        """
+        Args:
+            N: Number of filters in autoencoder
+            L: Length of the filters (in samples)
+            B: Number of channels in bottleneck 1 × 1-conv block
+            H: Number of channels in convolutional blocks
+            P: Kernel size in convolutional blocks
+            X: Number of convolutional blocks in each repeat
+            R: Number of repeats
+            C: Number of speakers
+            norm_type: BN, gLN, cLN
+            causal: causal or non-causal
+            mask_nonlinear: use which non-linear function to generate mask
+        """
+        super(ConvTasNet, self).__init__()
+        # Hyper-parameter
+        self.N, self.L, self.B, self.H, self.P, self.X, self.R, self.C = N, L, B, H, P, X, R, C
+        self.norm_type = norm_type
+        self.causal = causal
+        self.mask_nonlinear = mask_nonlinear
+        self.audio_channels = audio_channels
+        self.samplerate = samplerate
+        # Components
+        self.encoder = Encoder(L, N, audio_channels)
+        self.separator = TemporalConvNet(N, B, H, P, X, R, C, norm_type, causal, mask_nonlinear)
+        self.decoder = Decoder(N, L, audio_channels)
+        # init
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_normal_(p)
+
+    def valid_length(self, length):
+        return length
+
+    def forward(self, mixture):
+        """
+        Args:
+            mixture: [M, T], M is batch size, T is #samples
+        Returns:
+            est_source: [M, C, T]
+        """
+        mixture_w = self.encoder(mixture)
+        est_mask = self.separator(mixture_w)
+        est_source = self.decoder(mixture_w, est_mask)
+
+        # T changed after conv1d in encoder, fix it here
+        T_origin = mixture.size(-1)
+        T_conv = est_source.size(-1)
+        est_source = F.pad(est_source, (0, T_origin - T_conv))
+        return est_source
+
+
+class Encoder(nn.Module):
+    """Estimation of the nonnegative mixture weight by a 1-D conv layer.
+    """
+    def __init__(self, L, N, audio_channels):
+        super(Encoder, self).__init__()
+        # Hyper-parameter
+        self.L, self.N = L, N
+        # Components
+        # 50% overlap
+        self.conv1d_U = nn.Conv1d(audio_channels, N, kernel_size=L, stride=L // 2, bias=False)
+
+    def forward(self, mixture):
+        """
+        Args:
+            mixture: [M, T], M is batch size, T is #samples
+        Returns:
+            mixture_w: [M, N, K], where K = (T-L)/(L/2)+1 = 2T/L-1
+        """
+        mixture_w = F.relu(self.conv1d_U(mixture))  # [M, N, K]
+        return mixture_w
+
+
+class Decoder(nn.Module):
+    def __init__(self, N, L, audio_channels):
+        super(Decoder, self).__init__()
+        # Hyper-parameter
+        self.N, self.L = N, L
+        self.audio_channels = audio_channels
+        # Components
+        self.basis_signals = nn.Linear(N, audio_channels * L, bias=False)
+
+    def forward(self, mixture_w, est_mask):
+        """
+        Args:
+            mixture_w: [M, N, K]
+            est_mask: [M, C, N, K]
+        Returns:
+            est_source: [M, C, T]
+        """
+        # D = W * M
+        source_w = torch.unsqueeze(mixture_w, 1) * est_mask  # [M, C, N, K]
+        source_w = torch.transpose(source_w, 2, 3)  # [M, C, K, N]
+        # S = DV
+        est_source = self.basis_signals(source_w)  # [M, C, K, ac * L]
+        m, c, k, _ = est_source.size()
+        est_source = est_source.view(m, c, k, self.audio_channels, -1).transpose(2, 3).contiguous()
+        est_source = overlap_and_add(est_source, self.L // 2)  # M x C x ac x T
+        return est_source
+
+
+class TemporalConvNet(nn.Module):
+    def __init__(self, N, B, H, P, X, R, C, norm_type="gLN", causal=False, mask_nonlinear='relu'):
+        """
+        Args:
+            N: Number of filters in autoencoder
+            B: Number of channels in bottleneck 1 × 1-conv block
+            H: Number of channels in convolutional blocks
+            P: Kernel size in convolutional blocks
+            X: Number of convolutional blocks in each repeat
+            R: Number of repeats
+            C: Number of speakers
+            norm_type: BN, gLN, cLN
+            causal: causal or non-causal
+            mask_nonlinear: use which non-linear function to generate mask
+        """
+        super(TemporalConvNet, self).__init__()
+        # Hyper-parameter
+        self.C = C
+        self.mask_nonlinear = mask_nonlinear
+        # Components
+        # [M, N, K] -> [M, N, K]
+        layer_norm = ChannelwiseLayerNorm(N)
+        # [M, N, K] -> [M, B, K]
+        bottleneck_conv1x1 = nn.Conv1d(N, B, 1, bias=False)
+        # [M, B, K] -> [M, B, K]
+        repeats = []
+        for r in range(R):
+            blocks = []
+            for x in range(X):
+                dilation = 2**x
+                padding = (P - 1) * dilation if causal else (P - 1) * dilation // 2
+                blocks += [
+                    TemporalBlock(B,
+                                  H,
+                                  P,
+                                  stride=1,
+                                  padding=padding,
+                                  dilation=dilation,
+                                  norm_type=norm_type,
+                                  causal=causal)
+                ]
+            repeats += [nn.Sequential(*blocks)]
+        temporal_conv_net = nn.Sequential(*repeats)
+        # [M, B, K] -> [M, C*N, K]
+        mask_conv1x1 = nn.Conv1d(B, C * N, 1, bias=False)
+        # Put together
+        self.network = nn.Sequential(layer_norm, bottleneck_conv1x1, temporal_conv_net,
+                                     mask_conv1x1)
+
+    def forward(self, mixture_w):
+        """
+        Keep this API same with TasNet
+        Args:
+            mixture_w: [M, N, K], M is batch size
+        returns:
+            est_mask: [M, C, N, K]
+        """
+        M, N, K = mixture_w.size()
+        score = self.network(mixture_w)  # [M, N, K] -> [M, C*N, K]
+        score = score.view(M, self.C, N, K)  # [M, C*N, K] -> [M, C, N, K]
+        if self.mask_nonlinear == 'softmax':
+            est_mask = F.softmax(score, dim=1)
+        elif self.mask_nonlinear == 'relu':
+            est_mask = F.relu(score)
+        else:
+            raise ValueError("Unsupported mask non-linear function")
+        return est_mask
+
+
+class TemporalBlock(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 dilation,
+                 norm_type="gLN",
+                 causal=False):
+        super(TemporalBlock, self).__init__()
+        # [M, B, K] -> [M, H, K]
+        conv1x1 = nn.Conv1d(in_channels, out_channels, 1, bias=False)
+        prelu = nn.PReLU()
+        norm = chose_norm(norm_type, out_channels)
+        # [M, H, K] -> [M, B, K]
+        dsconv = DepthwiseSeparableConv(out_channels, in_channels, kernel_size, stride, padding,
+                                        dilation, norm_type, causal)
+        # Put together
+        self.net = nn.Sequential(conv1x1, prelu, norm, dsconv)
+
+    def forward(self, x):
+        """
+        Args:
+            x: [M, B, K]
+        Returns:
+            [M, B, K]
+        """
+        residual = x
+        out = self.net(x)
+        # TODO: when P = 3 here works fine, but when P = 2 maybe need to pad?
+        return out + residual  # look like w/o F.relu is better than w/ F.relu
+        # return F.relu(out + residual)
+
+
+class DepthwiseSeparableConv(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 dilation,
+                 norm_type="gLN",
+                 causal=False):
+        super(DepthwiseSeparableConv, self).__init__()
+        # Use `groups` option to implement depthwise convolution
+        # [M, H, K] -> [M, H, K]
+        depthwise_conv = nn.Conv1d(in_channels,
+                                   in_channels,
+                                   kernel_size,
+                                   stride=stride,
+                                   padding=padding,
+                                   dilation=dilation,
+                                   groups=in_channels,
+                                   bias=False)
+        if causal:
+            chomp = Chomp1d(padding)
+        prelu = nn.PReLU()
+        norm = chose_norm(norm_type, in_channels)
+        # [M, H, K] -> [M, B, K]
+        pointwise_conv = nn.Conv1d(in_channels, out_channels, 1, bias=False)
+        # Put together
+        if causal:
+            self.net = nn.Sequential(depthwise_conv, chomp, prelu, norm, pointwise_conv)
+        else:
+            self.net = nn.Sequential(depthwise_conv, prelu, norm, pointwise_conv)
+
+    def forward(self, x):
+        """
+        Args:
+            x: [M, H, K]
+        Returns:
+            result: [M, B, K]
+        """
+        return self.net(x)
+
+
+class Chomp1d(nn.Module):
+    """To ensure the output length is the same as the input.
+    """
+    def __init__(self, chomp_size):
+        super(Chomp1d, self).__init__()
+        self.chomp_size = chomp_size
+
+    def forward(self, x):
+        """
+        Args:
+            x: [M, H, Kpad]
+        Returns:
+            [M, H, K]
+        """
+        return x[:, :, :-self.chomp_size].contiguous()
+
+
+def chose_norm(norm_type, channel_size):
+    """The input of normlization will be (M, C, K), where M is batch size,
+       C is channel size and K is sequence length.
+    """
+    if norm_type == "gLN":
+        return GlobalLayerNorm(channel_size)
+    elif norm_type == "cLN":
+        return ChannelwiseLayerNorm(channel_size)
+    elif norm_type == "id":
+        return nn.Identity()
+    else:  # norm_type == "BN":
+        # Given input (M, C, K), nn.BatchNorm1d(C) will accumulate statics
+        # along M and K, so this BN usage is right.
+        return nn.BatchNorm1d(channel_size)
+
+
+# TODO: Use nn.LayerNorm to impl cLN to speed up
+class ChannelwiseLayerNorm(nn.Module):
+    """Channel-wise Layer Normalization (cLN)"""
+    def __init__(self, channel_size):
+        super(ChannelwiseLayerNorm, self).__init__()
+        self.gamma = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
+        self.beta = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.gamma.data.fill_(1)
+        self.beta.data.zero_()
+
+    def forward(self, y):
+        """
+        Args:
+            y: [M, N, K], M is batch size, N is channel size, K is length
+        Returns:
+            cLN_y: [M, N, K]
+        """
+        mean = torch.mean(y, dim=1, keepdim=True)  # [M, 1, K]
+        var = torch.var(y, dim=1, keepdim=True, unbiased=False)  # [M, 1, K]
+        cLN_y = self.gamma * (y - mean) / torch.pow(var + EPS, 0.5) + self.beta
+        return cLN_y
+
+
+class GlobalLayerNorm(nn.Module):
+    """Global Layer Normalization (gLN)"""
+    def __init__(self, channel_size):
+        super(GlobalLayerNorm, self).__init__()
+        self.gamma = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
+        self.beta = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.gamma.data.fill_(1)
+        self.beta.data.zero_()
+
+    def forward(self, y):
+        """
+        Args:
+            y: [M, N, K], M is batch size, N is channel size, K is length
+        Returns:
+            gLN_y: [M, N, K]
+        """
+        # TODO: in torch 1.0, torch.mean() support dim list
+        mean = y.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True)  # [M, 1, 1]
+        var = (torch.pow(y - mean, 2)).mean(dim=1, keepdim=True).mean(dim=2, keepdim=True)
+        gLN_y = self.gamma * (y - mean) / torch.pow(var + EPS, 0.5) + self.beta
+        return gLN_y
+
+
+if __name__ == "__main__":
+    torch.manual_seed(123)
+    M, N, L, T = 2, 3, 4, 12
+    K = 2 * T // L - 1
+    B, H, P, X, R, C, norm_type, causal = 2, 3, 3, 3, 2, 2, "gLN", False
+    mixture = torch.randint(3, (M, T))
+    # test Encoder
+    encoder = Encoder(L, N)
+    encoder.conv1d_U.weight.data = torch.randint(2, encoder.conv1d_U.weight.size())
+    mixture_w = encoder(mixture)
+    print('mixture', mixture)
+    print('U', encoder.conv1d_U.weight)
+    print('mixture_w', mixture_w)
+    print('mixture_w size', mixture_w.size())
+
+    # test TemporalConvNet
+    separator = TemporalConvNet(N, B, H, P, X, R, C, norm_type=norm_type, causal=causal)
+    est_mask = separator(mixture_w)
+    print('est_mask', est_mask)
+
+    # test Decoder
+    decoder = Decoder(N, L)
+    est_mask = torch.randint(2, (B, K, C, N))
+    est_source = decoder(mixture_w, est_mask)
+    print('est_source', est_source)
+
+    # test Conv-TasNet
+    conv_tasnet = ConvTasNet(N, L, B, H, P, X, R, C, norm_type=norm_type)
+    est_source = conv_tasnet(mixture)
+    print('est_source', est_source)
+    print('est_source size', est_source.size())

lib/uvr5_pack/demucs/tasnet_v2.py

@@ -0,0 +1,452 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Created on 2018/12
+# Author: Kaituo XU
+# Modified on 2019/11 by Alexandre Defossez, added support for multiple output channels
+# Here is the original license:
+# The MIT License (MIT)
+#
+# Copyright (c) 2018 Kaituo XU
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import capture_init
+
+EPS = 1e-8
+
+
+def overlap_and_add(signal, frame_step):
+    outer_dimensions = signal.size()[:-2]
+    frames, frame_length = signal.size()[-2:]
+
+    subframe_length = math.gcd(frame_length, frame_step)  # gcd=Greatest Common Divisor
+    subframe_step = frame_step // subframe_length
+    subframes_per_frame = frame_length // subframe_length
+    output_size = frame_step * (frames - 1) + frame_length
+    output_subframes = output_size // subframe_length
+
+    subframe_signal = signal.view(*outer_dimensions, -1, subframe_length)
+
+    frame = torch.arange(0, output_subframes,
+                         device=signal.device).unfold(0, subframes_per_frame, subframe_step)
+    frame = frame.long()  # signal may in GPU or CPU
+    frame = frame.contiguous().view(-1)
+
+    result = signal.new_zeros(*outer_dimensions, output_subframes, subframe_length)
+    result.index_add_(-2, frame, subframe_signal)
+    result = result.view(*outer_dimensions, -1)
+    return result
+
+
+class ConvTasNet(nn.Module):
+    @capture_init
+    def __init__(self,
+                 sources,
+                 N=256,
+                 L=20,
+                 B=256,
+                 H=512,
+                 P=3,
+                 X=8,
+                 R=4,
+                 audio_channels=2,
+                 norm_type="gLN",
+                 causal=False,
+                 mask_nonlinear='relu',
+                 samplerate=44100,
+                 segment_length=44100 * 2 * 4):
+        """
+        Args:
+            sources: list of sources
+            N: Number of filters in autoencoder
+            L: Length of the filters (in samples)
+            B: Number of channels in bottleneck 1 × 1-conv block
+            H: Number of channels in convolutional blocks
+            P: Kernel size in convolutional blocks
+            X: Number of convolutional blocks in each repeat
+            R: Number of repeats
+            norm_type: BN, gLN, cLN
+            causal: causal or non-causal
+            mask_nonlinear: use which non-linear function to generate mask
+        """
+        super(ConvTasNet, self).__init__()
+        # Hyper-parameter
+        self.sources = sources
+        self.C = len(sources)
+        self.N, self.L, self.B, self.H, self.P, self.X, self.R = N, L, B, H, P, X, R
+        self.norm_type = norm_type
+        self.causal = causal
+        self.mask_nonlinear = mask_nonlinear
+        self.audio_channels = audio_channels
+        self.samplerate = samplerate
+        self.segment_length = segment_length
+        # Components
+        self.encoder = Encoder(L, N, audio_channels)
+        self.separator = TemporalConvNet(
+            N, B, H, P, X, R, self.C, norm_type, causal, mask_nonlinear)
+        self.decoder = Decoder(N, L, audio_channels)
+        # init
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_normal_(p)
+
+    def valid_length(self, length):
+        return length
+
+    def forward(self, mixture):
+        """
+        Args:
+            mixture: [M, T], M is batch size, T is #samples
+        Returns:
+            est_source: [M, C, T]
+        """
+        mixture_w = self.encoder(mixture)
+        est_mask = self.separator(mixture_w)
+        est_source = self.decoder(mixture_w, est_mask)
+
+        # T changed after conv1d in encoder, fix it here
+        T_origin = mixture.size(-1)
+        T_conv = est_source.size(-1)
+        est_source = F.pad(est_source, (0, T_origin - T_conv))
+        return est_source
+
+
+class Encoder(nn.Module):
+    """Estimation of the nonnegative mixture weight by a 1-D conv layer.
+    """
+    def __init__(self, L, N, audio_channels):
+        super(Encoder, self).__init__()
+        # Hyper-parameter
+        self.L, self.N = L, N
+        # Components
+        # 50% overlap
+        self.conv1d_U = nn.Conv1d(audio_channels, N, kernel_size=L, stride=L // 2, bias=False)
+
+    def forward(self, mixture):
+        """
+        Args:
+            mixture: [M, T], M is batch size, T is #samples
+        Returns:
+            mixture_w: [M, N, K], where K = (T-L)/(L/2)+1 = 2T/L-1
+        """
+        mixture_w = F.relu(self.conv1d_U(mixture))  # [M, N, K]
+        return mixture_w
+
+
+class Decoder(nn.Module):
+    def __init__(self, N, L, audio_channels):
+        super(Decoder, self).__init__()
+        # Hyper-parameter
+        self.N, self.L = N, L
+        self.audio_channels = audio_channels
+        # Components
+        self.basis_signals = nn.Linear(N, audio_channels * L, bias=False)
+
+    def forward(self, mixture_w, est_mask):
+        """
+        Args:
+            mixture_w: [M, N, K]
+            est_mask: [M, C, N, K]
+        Returns:
+            est_source: [M, C, T]
+        """
+        # D = W * M
+        source_w = torch.unsqueeze(mixture_w, 1) * est_mask  # [M, C, N, K]
+        source_w = torch.transpose(source_w, 2, 3)  # [M, C, K, N]
+        # S = DV
+        est_source = self.basis_signals(source_w)  # [M, C, K, ac * L]
+        m, c, k, _ = est_source.size()
+        est_source = est_source.view(m, c, k, self.audio_channels, -1).transpose(2, 3).contiguous()
+        est_source = overlap_and_add(est_source, self.L // 2)  # M x C x ac x T
+        return est_source
+
+
+class TemporalConvNet(nn.Module):
+    def __init__(self, N, B, H, P, X, R, C, norm_type="gLN", causal=False, mask_nonlinear='relu'):
+        """
+        Args:
+            N: Number of filters in autoencoder
+            B: Number of channels in bottleneck 1 × 1-conv block
+            H: Number of channels in convolutional blocks
+            P: Kernel size in convolutional blocks
+            X: Number of convolutional blocks in each repeat
+            R: Number of repeats
+            C: Number of speakers
+            norm_type: BN, gLN, cLN
+            causal: causal or non-causal
+            mask_nonlinear: use which non-linear function to generate mask
+        """
+        super(TemporalConvNet, self).__init__()
+        # Hyper-parameter
+        self.C = C
+        self.mask_nonlinear = mask_nonlinear
+        # Components
+        # [M, N, K] -> [M, N, K]
+        layer_norm = ChannelwiseLayerNorm(N)
+        # [M, N, K] -> [M, B, K]
+        bottleneck_conv1x1 = nn.Conv1d(N, B, 1, bias=False)
+        # [M, B, K] -> [M, B, K]
+        repeats = []
+        for r in range(R):
+            blocks = []
+            for x in range(X):
+                dilation = 2**x
+                padding = (P - 1) * dilation if causal else (P - 1) * dilation // 2
+                blocks += [
+                    TemporalBlock(B,
+                                  H,
+                                  P,
+                                  stride=1,
+                                  padding=padding,
+                                  dilation=dilation,
+                                  norm_type=norm_type,
+                                  causal=causal)
+                ]
+            repeats += [nn.Sequential(*blocks)]
+        temporal_conv_net = nn.Sequential(*repeats)
+        # [M, B, K] -> [M, C*N, K]
+        mask_conv1x1 = nn.Conv1d(B, C * N, 1, bias=False)
+        # Put together
+        self.network = nn.Sequential(layer_norm, bottleneck_conv1x1, temporal_conv_net,
+                                     mask_conv1x1)
+
+    def forward(self, mixture_w):
+        """
+        Keep this API same with TasNet
+        Args:
+            mixture_w: [M, N, K], M is batch size
+        returns:
+            est_mask: [M, C, N, K]
+        """
+        M, N, K = mixture_w.size()
+        score = self.network(mixture_w)  # [M, N, K] -> [M, C*N, K]
+        score = score.view(M, self.C, N, K)  # [M, C*N, K] -> [M, C, N, K]
+        if self.mask_nonlinear == 'softmax':
+            est_mask = F.softmax(score, dim=1)
+        elif self.mask_nonlinear == 'relu':
+            est_mask = F.relu(score)
+        else:
+            raise ValueError("Unsupported mask non-linear function")
+        return est_mask
+
+
+class TemporalBlock(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 dilation,
+                 norm_type="gLN",
+                 causal=False):
+        super(TemporalBlock, self).__init__()
+        # [M, B, K] -> [M, H, K]
+        conv1x1 = nn.Conv1d(in_channels, out_channels, 1, bias=False)
+        prelu = nn.PReLU()
+        norm = chose_norm(norm_type, out_channels)
+        # [M, H, K] -> [M, B, K]
+        dsconv = DepthwiseSeparableConv(out_channels, in_channels, kernel_size, stride, padding,
+                                        dilation, norm_type, causal)
+        # Put together
+        self.net = nn.Sequential(conv1x1, prelu, norm, dsconv)
+
+    def forward(self, x):
+        """
+        Args:
+            x: [M, B, K]
+        Returns:
+            [M, B, K]
+        """
+        residual = x
+        out = self.net(x)
+        # TODO: when P = 3 here works fine, but when P = 2 maybe need to pad?
+        return out + residual  # look like w/o F.relu is better than w/ F.relu
+        # return F.relu(out + residual)
+
+
+class DepthwiseSeparableConv(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 dilation,
+                 norm_type="gLN",
+                 causal=False):
+        super(DepthwiseSeparableConv, self).__init__()
+        # Use `groups` option to implement depthwise convolution
+        # [M, H, K] -> [M, H, K]
+        depthwise_conv = nn.Conv1d(in_channels,
+                                   in_channels,
+                                   kernel_size,
+                                   stride=stride,
+                                   padding=padding,
+                                   dilation=dilation,
+                                   groups=in_channels,
+                                   bias=False)
+        if causal:
+            chomp = Chomp1d(padding)
+        prelu = nn.PReLU()
+        norm = chose_norm(norm_type, in_channels)
+        # [M, H, K] -> [M, B, K]
+        pointwise_conv = nn.Conv1d(in_channels, out_channels, 1, bias=False)
+        # Put together
+        if causal:
+            self.net = nn.Sequential(depthwise_conv, chomp, prelu, norm, pointwise_conv)
+        else:
+            self.net = nn.Sequential(depthwise_conv, prelu, norm, pointwise_conv)
+
+    def forward(self, x):
+        """
+        Args:
+            x: [M, H, K]
+        Returns:
+            result: [M, B, K]
+        """
+        return self.net(x)
+
+
+class Chomp1d(nn.Module):
+    """To ensure the output length is the same as the input.
+    """
+    def __init__(self, chomp_size):
+        super(Chomp1d, self).__init__()
+        self.chomp_size = chomp_size
+
+    def forward(self, x):
+        """
+        Args:
+            x: [M, H, Kpad]
+        Returns:
+            [M, H, K]
+        """
+        return x[:, :, :-self.chomp_size].contiguous()
+
+
+def chose_norm(norm_type, channel_size):
+    """The input of normlization will be (M, C, K), where M is batch size,
+       C is channel size and K is sequence length.
+    """
+    if norm_type == "gLN":
+        return GlobalLayerNorm(channel_size)
+    elif norm_type == "cLN":
+        return ChannelwiseLayerNorm(channel_size)
+    elif norm_type == "id":
+        return nn.Identity()
+    else:  # norm_type == "BN":
+        # Given input (M, C, K), nn.BatchNorm1d(C) will accumulate statics
+        # along M and K, so this BN usage is right.
+        return nn.BatchNorm1d(channel_size)
+
+
+# TODO: Use nn.LayerNorm to impl cLN to speed up
+class ChannelwiseLayerNorm(nn.Module):
+    """Channel-wise Layer Normalization (cLN)"""
+    def __init__(self, channel_size):
+        super(ChannelwiseLayerNorm, self).__init__()
+        self.gamma = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
+        self.beta = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.gamma.data.fill_(1)
+        self.beta.data.zero_()
+
+    def forward(self, y):
+        """
+        Args:
+            y: [M, N, K], M is batch size, N is channel size, K is length
+        Returns:
+            cLN_y: [M, N, K]
+        """
+        mean = torch.mean(y, dim=1, keepdim=True)  # [M, 1, K]
+        var = torch.var(y, dim=1, keepdim=True, unbiased=False)  # [M, 1, K]
+        cLN_y = self.gamma * (y - mean) / torch.pow(var + EPS, 0.5) + self.beta
+        return cLN_y
+
+
+class GlobalLayerNorm(nn.Module):
+    """Global Layer Normalization (gLN)"""
+    def __init__(self, channel_size):
+        super(GlobalLayerNorm, self).__init__()
+        self.gamma = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
+        self.beta = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.gamma.data.fill_(1)
+        self.beta.data.zero_()
+
+    def forward(self, y):
+        """
+        Args:
+            y: [M, N, K], M is batch size, N is channel size, K is length
+        Returns:
+            gLN_y: [M, N, K]
+        """
+        # TODO: in torch 1.0, torch.mean() support dim list
+        mean = y.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True)  # [M, 1, 1]
+        var = (torch.pow(y - mean, 2)).mean(dim=1, keepdim=True).mean(dim=2, keepdim=True)
+        gLN_y = self.gamma * (y - mean) / torch.pow(var + EPS, 0.5) + self.beta
+        return gLN_y
+
+
+if __name__ == "__main__":
+    torch.manual_seed(123)
+    M, N, L, T = 2, 3, 4, 12
+    K = 2 * T // L - 1
+    B, H, P, X, R, C, norm_type, causal = 2, 3, 3, 3, 2, 2, "gLN", False
+    mixture = torch.randint(3, (M, T))
+    # test Encoder
+    encoder = Encoder(L, N)
+    encoder.conv1d_U.weight.data = torch.randint(2, encoder.conv1d_U.weight.size())
+    mixture_w = encoder(mixture)
+    print('mixture', mixture)
+    print('U', encoder.conv1d_U.weight)
+    print('mixture_w', mixture_w)
+    print('mixture_w size', mixture_w.size())
+
+    # test TemporalConvNet
+    separator = TemporalConvNet(N, B, H, P, X, R, C, norm_type=norm_type, causal=causal)
+    est_mask = separator(mixture_w)
+    print('est_mask', est_mask)
+
+    # test Decoder
+    decoder = Decoder(N, L)
+    est_mask = torch.randint(2, (B, K, C, N))
+    est_source = decoder(mixture_w, est_mask)
+    print('est_source', est_source)
+
+    # test Conv-TasNet
+    conv_tasnet = ConvTasNet(N, L, B, H, P, X, R, C, norm_type=norm_type)
+    est_source = conv_tasnet(mixture)
+    print('est_source', est_source)
+    print('est_source size', est_source.size())

lib/uvr5_pack/demucs/transformer.py

@@ -0,0 +1,839 @@
+# Copyright (c) 2019-present, Meta, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# First author is Simon Rouard.
+
+import random
+import typing as tp
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+from einops import rearrange
+
+
+def create_sin_embedding(
+    length: int, dim: int, shift: int = 0, device="cpu", max_period=10000
+):
+    # We aim for TBC format
+    assert dim % 2 == 0
+    pos = shift + torch.arange(length, device=device).view(-1, 1, 1)
+    half_dim = dim // 2
+    adim = torch.arange(dim // 2, device=device).view(1, 1, -1)
+    phase = pos / (max_period ** (adim / (half_dim - 1)))
+    return torch.cat(
+        [
+            torch.cos(phase),
+            torch.sin(phase),
+        ],
+        dim=-1,
+    )
+
+
+def create_2d_sin_embedding(d_model, height, width, device="cpu", max_period=10000):
+    """
+    :param d_model: dimension of the model
+    :param height: height of the positions
+    :param width: width of the positions
+    :return: d_model*height*width position matrix
+    """
+    if d_model % 4 != 0:
+        raise ValueError(
+            "Cannot use sin/cos positional encoding with "
+            "odd dimension (got dim={:d})".format(d_model)
+        )
+    pe = torch.zeros(d_model, height, width)
+    # Each dimension use half of d_model
+    d_model = int(d_model / 2)
+    div_term = torch.exp(
+        torch.arange(0.0, d_model, 2) * -(math.log(max_period) / d_model)
+    )
+    pos_w = torch.arange(0.0, width).unsqueeze(1)
+    pos_h = torch.arange(0.0, height).unsqueeze(1)
+    pe[0:d_model:2, :, :] = (
+        torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
+    )
+    pe[1:d_model:2, :, :] = (
+        torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
+    )
+    pe[d_model::2, :, :] = (
+        torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
+    )
+    pe[d_model + 1:: 2, :, :] = (
+        torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
+    )
+
+    return pe[None, :].to(device)
+
+
+def create_sin_embedding_cape(
+    length: int,
+    dim: int,
+    batch_size: int,
+    mean_normalize: bool,
+    augment: bool,  # True during training
+    max_global_shift: float = 0.0,  # delta max
+    max_local_shift: float = 0.0,  # epsilon max
+    max_scale: float = 1.0,
+    device: str = "cpu",
+    max_period: float = 10000.0,
+):
+    # We aim for TBC format
+    assert dim % 2 == 0
+    pos = 1.0 * torch.arange(length).view(-1, 1, 1)  # (length, 1, 1)
+    pos = pos.repeat(1, batch_size, 1)  # (length, batch_size, 1)
+    if mean_normalize:
+        pos -= torch.nanmean(pos, dim=0, keepdim=True)
+
+    if augment:
+        delta = np.random.uniform(
+            -max_global_shift, +max_global_shift, size=[1, batch_size, 1]
+        )
+        delta_local = np.random.uniform(
+            -max_local_shift, +max_local_shift, size=[length, batch_size, 1]
+        )
+        log_lambdas = np.random.uniform(
+            -np.log(max_scale), +np.log(max_scale), size=[1, batch_size, 1]
+        )
+        pos = (pos + delta + delta_local) * np.exp(log_lambdas)
+
+    pos = pos.to(device)
+
+    half_dim = dim // 2
+    adim = torch.arange(dim // 2, device=device).view(1, 1, -1)
+    phase = pos / (max_period ** (adim / (half_dim - 1)))
+    return torch.cat(
+        [
+            torch.cos(phase),
+            torch.sin(phase),
+        ],
+        dim=-1,
+    ).float()
+
+
+def get_causal_mask(length):
+    pos = torch.arange(length)
+    return pos > pos[:, None]
+
+
+def get_elementary_mask(
+    T1,
+    T2,
+    mask_type,
+    sparse_attn_window,
+    global_window,
+    mask_random_seed,
+    sparsity,
+    device,
+):
+    """
+    When the input of the Decoder has length T1 and the output T2
+    The mask matrix has shape (T2, T1)
+    """
+    assert mask_type in ["diag", "jmask", "random", "global"]
+
+    if mask_type == "global":
+        mask = torch.zeros(T2, T1, dtype=torch.bool)
+        mask[:, :global_window] = True
+        line_window = int(global_window * T2 / T1)
+        mask[:line_window, :] = True
+
+    if mask_type == "diag":
+
+        mask = torch.zeros(T2, T1, dtype=torch.bool)
+        rows = torch.arange(T2)[:, None]
+        cols = (
+            (T1 / T2 * rows + torch.arange(-sparse_attn_window, sparse_attn_window + 1))
+            .long()
+            .clamp(0, T1 - 1)
+        )
+        mask.scatter_(1, cols, torch.ones(1, dtype=torch.bool).expand_as(cols))
+
+    elif mask_type == "jmask":
+        mask = torch.zeros(T2 + 2, T1 + 2, dtype=torch.bool)
+        rows = torch.arange(T2 + 2)[:, None]
+        t = torch.arange(0, int((2 * T1) ** 0.5 + 1))
+        t = (t * (t + 1) / 2).int()
+        t = torch.cat([-t.flip(0)[:-1], t])
+        cols = (T1 / T2 * rows + t).long().clamp(0, T1 + 1)
+        mask.scatter_(1, cols, torch.ones(1, dtype=torch.bool).expand_as(cols))
+        mask = mask[1:-1, 1:-1]
+
+    elif mask_type == "random":
+        gene = torch.Generator(device=device)
+        gene.manual_seed(mask_random_seed)
+        mask = (
+            torch.rand(T1 * T2, generator=gene, device=device).reshape(T2, T1)
+            > sparsity
+        )
+
+    mask = mask.to(device)
+    return mask
+
+
+def get_mask(
+    T1,
+    T2,
+    mask_type,
+    sparse_attn_window,
+    global_window,
+    mask_random_seed,
+    sparsity,
+    device,
+):
+    """
+    Return a SparseCSRTensor mask that is a combination of elementary masks
+    mask_type can be a combination of multiple masks: for instance "diag_jmask_random"
+    """
+    from xformers.sparse import SparseCSRTensor
+    # create a list
+    mask_types = mask_type.split("_")
+
+    all_masks = [
+        get_elementary_mask(
+            T1,
+            T2,
+            mask,
+            sparse_attn_window,
+            global_window,
+            mask_random_seed,
+            sparsity,
+            device,
+        )
+        for mask in mask_types
+    ]
+
+    final_mask = torch.stack(all_masks).sum(axis=0) > 0
+
+    return SparseCSRTensor.from_dense(final_mask[None])
+
+
+class ScaledEmbedding(nn.Module):
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        scale: float = 1.0,
+        boost: float = 3.0,
+    ):
+        super().__init__()
+        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
+        self.embedding.weight.data *= scale / boost
+        self.boost = boost
+
+    @property
+    def weight(self):
+        return self.embedding.weight * self.boost
+
+    def forward(self, x):
+        return self.embedding(x) * self.boost
+
+
+class LayerScale(nn.Module):
+    """Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf).
+    This rescales diagonaly residual outputs close to 0 initially, then learnt.
+    """
+
+    def __init__(self, channels: int, init: float = 0, channel_last=False):
+        """
+        channel_last = False corresponds to (B, C, T) tensors
+        channel_last = True corresponds to (T, B, C) tensors
+        """
+        super().__init__()
+        self.channel_last = channel_last
+        self.scale = nn.Parameter(torch.zeros(channels, requires_grad=True))
+        self.scale.data[:] = init
+
+    def forward(self, x):
+        if self.channel_last:
+            return self.scale * x
+        else:
+            return self.scale[:, None] * x
+
+
+class MyGroupNorm(nn.GroupNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x):
+        """
+        x: (B, T, C)
+        if num_groups=1: Normalisation on all T and C together for each B
+        """
+        x = x.transpose(1, 2)
+        return super().forward(x).transpose(1, 2)
+
+
+class MyTransformerEncoderLayer(nn.TransformerEncoderLayer):
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        activation=F.relu,
+        group_norm=0,
+        norm_first=False,
+        norm_out=False,
+        layer_norm_eps=1e-5,
+        layer_scale=False,
+        init_values=1e-4,
+        device=None,
+        dtype=None,
+        sparse=False,
+        mask_type="diag",
+        mask_random_seed=42,
+        sparse_attn_window=500,
+        global_window=50,
+        auto_sparsity=False,
+        sparsity=0.95,
+        batch_first=False,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation=activation,
+            layer_norm_eps=layer_norm_eps,
+            batch_first=batch_first,
+            norm_first=norm_first,
+            device=device,
+            dtype=dtype,
+        )
+        self.sparse = sparse
+        self.auto_sparsity = auto_sparsity
+        if sparse:
+            if not auto_sparsity:
+                self.mask_type = mask_type
+                self.sparse_attn_window = sparse_attn_window
+                self.global_window = global_window
+            self.sparsity = sparsity
+        if group_norm:
+            self.norm1 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm2 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+
+        self.norm_out = None
+        if self.norm_first & norm_out:
+            self.norm_out = MyGroupNorm(num_groups=int(norm_out), num_channels=d_model)
+        self.gamma_1 = (
+            LayerScale(d_model, init_values, True) if layer_scale else nn.Identity()
+        )
+        self.gamma_2 = (
+            LayerScale(d_model, init_values, True) if layer_scale else nn.Identity()
+        )
+
+        if sparse:
+            self.self_attn = MultiheadAttention(
+                d_model, nhead, dropout=dropout, batch_first=batch_first,
+                auto_sparsity=sparsity if auto_sparsity else 0,
+            )
+            self.__setattr__("src_mask", torch.zeros(1, 1))
+            self.mask_random_seed = mask_random_seed
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        """
+        if batch_first = False, src shape is (T, B, C)
+        the case where batch_first=True is not covered
+        """
+        device = src.device
+        x = src
+        T, B, C = x.shape
+        if self.sparse and not self.auto_sparsity:
+            assert src_mask is None
+            src_mask = self.src_mask
+            if src_mask.shape[-1] != T:
+                src_mask = get_mask(
+                    T,
+                    T,
+                    self.mask_type,
+                    self.sparse_attn_window,
+                    self.global_window,
+                    self.mask_random_seed,
+                    self.sparsity,
+                    device,
+                )
+                self.__setattr__("src_mask", src_mask)
+
+        if self.norm_first:
+            x = x + self.gamma_1(
+                self._sa_block(self.norm1(x), src_mask, src_key_padding_mask)
+            )
+            x = x + self.gamma_2(self._ff_block(self.norm2(x)))
+
+            if self.norm_out:
+                x = self.norm_out(x)
+        else:
+            x = self.norm1(
+                x + self.gamma_1(self._sa_block(x, src_mask, src_key_padding_mask))
+            )
+            x = self.norm2(x + self.gamma_2(self._ff_block(x)))
+
+        return x
+
+
+class CrossTransformerEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation=F.relu,
+        layer_norm_eps: float = 1e-5,
+        layer_scale: bool = False,
+        init_values: float = 1e-4,
+        norm_first: bool = False,
+        group_norm: bool = False,
+        norm_out: bool = False,
+        sparse=False,
+        mask_type="diag",
+        mask_random_seed=42,
+        sparse_attn_window=500,
+        global_window=50,
+        sparsity=0.95,
+        auto_sparsity=None,
+        device=None,
+        dtype=None,
+        batch_first=False,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        self.sparse = sparse
+        self.auto_sparsity = auto_sparsity
+        if sparse:
+            if not auto_sparsity:
+                self.mask_type = mask_type
+                self.sparse_attn_window = sparse_attn_window
+                self.global_window = global_window
+            self.sparsity = sparsity
+
+        self.cross_attn: nn.Module
+        self.cross_attn = nn.MultiheadAttention(
+            d_model, nhead, dropout=dropout, batch_first=batch_first)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward, **factory_kwargs)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model, **factory_kwargs)
+
+        self.norm_first = norm_first
+        self.norm1: nn.Module
+        self.norm2: nn.Module
+        self.norm3: nn.Module
+        if group_norm:
+            self.norm1 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm2 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm3 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+        else:
+            self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm3 = nn.LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+
+        self.norm_out = None
+        if self.norm_first & norm_out:
+            self.norm_out = MyGroupNorm(num_groups=int(norm_out), num_channels=d_model)
+
+        self.gamma_1 = (
+            LayerScale(d_model, init_values, True) if layer_scale else nn.Identity()
+        )
+        self.gamma_2 = (
+            LayerScale(d_model, init_values, True) if layer_scale else nn.Identity()
+        )
+
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            self.activation = self._get_activation_fn(activation)
+        else:
+            self.activation = activation
+
+        if sparse:
+            self.cross_attn = MultiheadAttention(
+                d_model, nhead, dropout=dropout, batch_first=batch_first,
+                auto_sparsity=sparsity if auto_sparsity else 0)
+            if not auto_sparsity:
+                self.__setattr__("mask", torch.zeros(1, 1))
+                self.mask_random_seed = mask_random_seed
+
+    def forward(self, q, k, mask=None):
+        """
+        Args:
+            q: tensor of shape (T, B, C)
+            k: tensor of shape (S, B, C)
+            mask: tensor of shape (T, S)
+
+        """
+        device = q.device
+        T, B, C = q.shape
+        S, B, C = k.shape
+        if self.sparse and not self.auto_sparsity:
+            assert mask is None
+            mask = self.mask
+            if mask.shape[-1] != S or mask.shape[-2] != T:
+                mask = get_mask(
+                    S,
+                    T,
+                    self.mask_type,
+                    self.sparse_attn_window,
+                    self.global_window,
+                    self.mask_random_seed,
+                    self.sparsity,
+                    device,
+                )
+                self.__setattr__("mask", mask)
+
+        if self.norm_first:
+            x = q + self.gamma_1(self._ca_block(self.norm1(q), self.norm2(k), mask))
+            x = x + self.gamma_2(self._ff_block(self.norm3(x)))
+            if self.norm_out:
+                x = self.norm_out(x)
+        else:
+            x = self.norm1(q + self.gamma_1(self._ca_block(q, k, mask)))
+            x = self.norm2(x + self.gamma_2(self._ff_block(x)))
+
+        return x
+
+    # self-attention block
+    def _ca_block(self, q, k, attn_mask=None):
+        x = self.cross_attn(q, k, k, attn_mask=attn_mask, need_weights=False)[0]
+        return self.dropout1(x)
+
+    # feed forward block
+    def _ff_block(self, x):
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout2(x)
+
+    def _get_activation_fn(self, activation):
+        if activation == "relu":
+            return F.relu
+        elif activation == "gelu":
+            return F.gelu
+
+        raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
+
+
+# ----------------- MULTI-BLOCKS MODELS: -----------------------
+
+
+class CrossTransformerEncoder(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        emb: str = "sin",
+        hidden_scale: float = 4.0,
+        num_heads: int = 8,
+        num_layers: int = 6,
+        cross_first: bool = False,
+        dropout: float = 0.0,
+        max_positions: int = 1000,
+        norm_in: bool = True,
+        norm_in_group: bool = False,
+        group_norm: int = False,
+        norm_first: bool = False,
+        norm_out: bool = False,
+        max_period: float = 10000.0,
+        weight_decay: float = 0.0,
+        lr: tp.Optional[float] = None,
+        layer_scale: bool = False,
+        gelu: bool = True,
+        sin_random_shift: int = 0,
+        weight_pos_embed: float = 1.0,
+        cape_mean_normalize: bool = True,
+        cape_augment: bool = True,
+        cape_glob_loc_scale: list = [5000.0, 1.0, 1.4],
+        sparse_self_attn: bool = False,
+        sparse_cross_attn: bool = False,
+        mask_type: str = "diag",
+        mask_random_seed: int = 42,
+        sparse_attn_window: int = 500,
+        global_window: int = 50,
+        auto_sparsity: bool = False,
+        sparsity: float = 0.95,
+    ):
+        super().__init__()
+        """
+        """
+        assert dim % num_heads == 0
+
+        hidden_dim = int(dim * hidden_scale)
+
+        self.num_layers = num_layers
+        # classic parity = 1 means that if idx%2 == 1 there is a
+        # classical encoder else there is a cross encoder
+        self.classic_parity = 1 if cross_first else 0
+        self.emb = emb
+        self.max_period = max_period
+        self.weight_decay = weight_decay
+        self.weight_pos_embed = weight_pos_embed
+        self.sin_random_shift = sin_random_shift
+        if emb == "cape":
+            self.cape_mean_normalize = cape_mean_normalize
+            self.cape_augment = cape_augment
+            self.cape_glob_loc_scale = cape_glob_loc_scale
+        if emb == "scaled":
+            self.position_embeddings = ScaledEmbedding(max_positions, dim, scale=0.2)
+
+        self.lr = lr
+
+        activation: tp.Any = F.gelu if gelu else F.relu
+
+        self.norm_in: nn.Module
+        self.norm_in_t: nn.Module
+        if norm_in:
+            self.norm_in = nn.LayerNorm(dim)
+            self.norm_in_t = nn.LayerNorm(dim)
+        elif norm_in_group:
+            self.norm_in = MyGroupNorm(int(norm_in_group), dim)
+            self.norm_in_t = MyGroupNorm(int(norm_in_group), dim)
+        else:
+            self.norm_in = nn.Identity()
+            self.norm_in_t = nn.Identity()
+
+        # spectrogram layers
+        self.layers = nn.ModuleList()
+        # temporal layers
+        self.layers_t = nn.ModuleList()
+
+        kwargs_common = {
+            "d_model": dim,
+            "nhead": num_heads,
+            "dim_feedforward": hidden_dim,
+            "dropout": dropout,
+            "activation": activation,
+            "group_norm": group_norm,
+            "norm_first": norm_first,
+            "norm_out": norm_out,
+            "layer_scale": layer_scale,
+            "mask_type": mask_type,
+            "mask_random_seed": mask_random_seed,
+            "sparse_attn_window": sparse_attn_window,
+            "global_window": global_window,
+            "sparsity": sparsity,
+            "auto_sparsity": auto_sparsity,
+            "batch_first": True,
+        }
+
+        kwargs_classic_encoder = dict(kwargs_common)
+        kwargs_classic_encoder.update({
+            "sparse": sparse_self_attn,
+        })
+        kwargs_cross_encoder = dict(kwargs_common)
+        kwargs_cross_encoder.update({
+            "sparse": sparse_cross_attn,
+        })
+
+        for idx in range(num_layers):
+            if idx % 2 == self.classic_parity:
+
+                self.layers.append(MyTransformerEncoderLayer(**kwargs_classic_encoder))
+                self.layers_t.append(
+                    MyTransformerEncoderLayer(**kwargs_classic_encoder)
+                )
+
+            else:
+                self.layers.append(CrossTransformerEncoderLayer(**kwargs_cross_encoder))
+
+                self.layers_t.append(
+                    CrossTransformerEncoderLayer(**kwargs_cross_encoder)
+                )
+
+    def forward(self, x, xt):
+        B, C, Fr, T1 = x.shape
+        pos_emb_2d = create_2d_sin_embedding(
+            C, Fr, T1, x.device, self.max_period
+        )  # (1, C, Fr, T1)
+        pos_emb_2d = rearrange(pos_emb_2d, "b c fr t1 -> b (t1 fr) c")
+        x = rearrange(x, "b c fr t1 -> b (t1 fr) c")
+        x = self.norm_in(x)
+        x = x + self.weight_pos_embed * pos_emb_2d
+
+        B, C, T2 = xt.shape
+        xt = rearrange(xt, "b c t2 -> b t2 c")  # now T2, B, C
+        pos_emb = self._get_pos_embedding(T2, B, C, x.device)
+        pos_emb = rearrange(pos_emb, "t2 b c -> b t2 c")
+        xt = self.norm_in_t(xt)
+        xt = xt + self.weight_pos_embed * pos_emb
+
+        for idx in range(self.num_layers):
+            if idx % 2 == self.classic_parity:
+                x = self.layers[idx](x)
+                xt = self.layers_t[idx](xt)
+            else:
+                old_x = x
+                x = self.layers[idx](x, xt)
+                xt = self.layers_t[idx](xt, old_x)
+
+        x = rearrange(x, "b (t1 fr) c -> b c fr t1", t1=T1)
+        xt = rearrange(xt, "b t2 c -> b c t2")
+        return x, xt
+
+    def _get_pos_embedding(self, T, B, C, device):
+        if self.emb == "sin":
+            shift = random.randrange(self.sin_random_shift + 1)
+            pos_emb = create_sin_embedding(
+                T, C, shift=shift, device=device, max_period=self.max_period
+            )
+        elif self.emb == "cape":
+            if self.training:
+                pos_emb = create_sin_embedding_cape(
+                    T,
+                    C,
+                    B,
+                    device=device,
+                    max_period=self.max_period,
+                    mean_normalize=self.cape_mean_normalize,
+                    augment=self.cape_augment,
+                    max_global_shift=self.cape_glob_loc_scale[0],
+                    max_local_shift=self.cape_glob_loc_scale[1],
+                    max_scale=self.cape_glob_loc_scale[2],
+                )
+            else:
+                pos_emb = create_sin_embedding_cape(
+                    T,
+                    C,
+                    B,
+                    device=device,
+                    max_period=self.max_period,
+                    mean_normalize=self.cape_mean_normalize,
+                    augment=False,
+                )
+
+        elif self.emb == "scaled":
+            pos = torch.arange(T, device=device)
+            pos_emb = self.position_embeddings(pos)[:, None]
+
+        return pos_emb
+
+    def make_optim_group(self):
+        group = {"params": list(self.parameters()), "weight_decay": self.weight_decay}
+        if self.lr is not None:
+            group["lr"] = self.lr
+        return group
+
+
+# Attention Modules
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_attn=False,
+        kdim=None,
+        vdim=None,
+        batch_first=False,
+        auto_sparsity=None,
+    ):
+        super().__init__()
+        assert auto_sparsity is not None, "sanity check"
+        self.num_heads = num_heads
+        self.q = torch.nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.k = torch.nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.v = torch.nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.attn_drop = torch.nn.Dropout(dropout)
+        self.proj = torch.nn.Linear(embed_dim, embed_dim, bias)
+        self.proj_drop = torch.nn.Dropout(dropout)
+        self.batch_first = batch_first
+        self.auto_sparsity = auto_sparsity
+
+    def forward(
+        self,
+        query,
+        key,
+        value,
+        key_padding_mask=None,
+        need_weights=True,
+        attn_mask=None,
+        average_attn_weights=True,
+    ):
+
+        if not self.batch_first:  # N, B, C
+            query = query.permute(1, 0, 2)  # B, N_q, C
+            key = key.permute(1, 0, 2)  # B, N_k, C
+            value = value.permute(1, 0, 2)  # B, N_k, C
+        B, N_q, C = query.shape
+        B, N_k, C = key.shape
+
+        q = (
+            self.q(query)
+            .reshape(B, N_q, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+        )
+        q = q.flatten(0, 1)
+        k = (
+            self.k(key)
+            .reshape(B, N_k, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+        )
+        k = k.flatten(0, 1)
+        v = (
+            self.v(value)
+            .reshape(B, N_k, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+        )
+        v = v.flatten(0, 1)
+
+        if self.auto_sparsity:
+            assert attn_mask is None
+            x = dynamic_sparse_attention(q, k, v, sparsity=self.auto_sparsity)
+        else:
+            x = scaled_dot_product_attention(q, k, v, attn_mask, dropout=self.attn_drop)
+        x = x.reshape(B, self.num_heads, N_q, C // self.num_heads)
+
+        x = x.transpose(1, 2).reshape(B, N_q, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        if not self.batch_first:
+            x = x.permute(1, 0, 2)
+        return x, None
+
+
+def scaled_query_key_softmax(q, k, att_mask):
+    from xformers.ops import masked_matmul
+    q = q / (k.size(-1)) ** 0.5
+    att = masked_matmul(q, k.transpose(-2, -1), att_mask)
+    att = torch.nn.functional.softmax(att, -1)
+    return att
+
+
+def scaled_dot_product_attention(q, k, v, att_mask, dropout):
+    att = scaled_query_key_softmax(q, k, att_mask=att_mask)
+    att = dropout(att)
+    y = att @ v
+    return y
+
+
+def _compute_buckets(x, R):
+    qq = torch.einsum('btf,bfhi->bhti', x, R)
+    qq = torch.cat([qq, -qq], dim=-1)
+    buckets = qq.argmax(dim=-1)
+
+    return buckets.permute(0, 2, 1).byte().contiguous()
+
+
+def dynamic_sparse_attention(query, key, value, sparsity, infer_sparsity=True, attn_bias=None):
+    # assert False, "The code for the custom sparse kernel is not ready for release yet."
+    from xformers.ops import find_locations, sparse_memory_efficient_attention
+    n_hashes = 32
+    proj_size = 4
+    query, key, value = [x.contiguous() for x in [query, key, value]]
+    with torch.no_grad():
+        R = torch.randn(1, query.shape[-1], n_hashes, proj_size // 2, device=query.device)
+        bucket_query = _compute_buckets(query, R)
+        bucket_key = _compute_buckets(key, R)
+        row_offsets, column_indices = find_locations(
+            bucket_query, bucket_key, sparsity, infer_sparsity)
+    return sparse_memory_efficient_attention(
+        query, key, value, row_offsets, column_indices, attn_bias)

lib/uvr5_pack/demucs/utils.py

@@ -0,0 +1,502 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections import defaultdict
+from contextlib import contextmanager
+import math
+import os
+import tempfile
+import typing as tp
+
+import errno
+import functools
+import hashlib
+import inspect
+import io
+import os
+import random
+import socket
+import tempfile
+import warnings
+import zlib
+import tkinter as tk
+
+from diffq import UniformQuantizer, DiffQuantizer
+import torch as th
+import tqdm
+from torch import distributed
+from torch.nn import functional as F
+
+import torch
+
+def unfold(a, kernel_size, stride):
+    """Given input of size [*OT, T], output Tensor of size [*OT, F, K]
+    with K the kernel size, by extracting frames with the given stride.
+
+    This will pad the input so that `F = ceil(T / K)`.
+
+    see https://github.com/pytorch/pytorch/issues/60466
+    """
+    *shape, length = a.shape
+    n_frames = math.ceil(length / stride)
+    tgt_length = (n_frames - 1) * stride + kernel_size
+    a = F.pad(a, (0, tgt_length - length))
+    strides = list(a.stride())
+    assert strides[-1] == 1, 'data should be contiguous'
+    strides = strides[:-1] + [stride, 1]
+    return a.as_strided([*shape, n_frames, kernel_size], strides)
+
+
+def center_trim(tensor: torch.Tensor, reference: tp.Union[torch.Tensor, int]):
+    """
+    Center trim `tensor` with respect to `reference`, along the last dimension.
+    `reference` can also be a number, representing the length to trim to.
+    If the size difference != 0 mod 2, the extra sample is removed on the right side.
+    """
+    ref_size: int
+    if isinstance(reference, torch.Tensor):
+        ref_size = reference.size(-1)
+    else:
+        ref_size = reference
+    delta = tensor.size(-1) - ref_size
+    if delta < 0:
+        raise ValueError("tensor must be larger than reference. " f"Delta is {delta}.")
+    if delta:
+        tensor = tensor[..., delta // 2:-(delta - delta // 2)]
+    return tensor
+
+
+def pull_metric(history: tp.List[dict], name: str):
+    out = []
+    for metrics in history:
+        metric = metrics
+        for part in name.split("."):
+            metric = metric[part]
+        out.append(metric)
+    return out
+
+
+def EMA(beta: float = 1):
+    """
+    Exponential Moving Average callback.
+    Returns a single function that can be called to repeatidly update the EMA
+    with a dict of metrics. The callback will return
+    the new averaged dict of metrics.
+
+    Note that for `beta=1`, this is just plain averaging.
+    """
+    fix: tp.Dict[str, float] = defaultdict(float)
+    total: tp.Dict[str, float] = defaultdict(float)
+
+    def _update(metrics: dict, weight: float = 1) -> dict:
+        nonlocal total, fix
+        for key, value in metrics.items():
+            total[key] = total[key] * beta + weight * float(value)
+            fix[key] = fix[key] * beta + weight
+        return {key: tot / fix[key] for key, tot in total.items()}
+    return _update
+
+
+def sizeof_fmt(num: float, suffix: str = 'B'):
+    """
+    Given `num` bytes, return human readable size.
+    Taken from https://stackoverflow.com/a/1094933
+    """
+    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
+        if abs(num) < 1024.0:
+            return "%3.1f%s%s" % (num, unit, suffix)
+        num /= 1024.0
+    return "%.1f%s%s" % (num, 'Yi', suffix)
+
+
+@contextmanager
+def temp_filenames(count: int, delete=True):
+    names = []
+    try:
+        for _ in range(count):
+            names.append(tempfile.NamedTemporaryFile(delete=False).name)
+        yield names
+    finally:
+        if delete:
+            for name in names:
+                os.unlink(name)
+
+def average_metric(metric, count=1.):
+    """
+    Average `metric` which should be a float across all hosts. `count` should be
+    the weight for this particular host (i.e. number of examples).
+    """
+    metric = th.tensor([count, count * metric], dtype=th.float32, device='cuda')
+    distributed.all_reduce(metric, op=distributed.ReduceOp.SUM)
+    return metric[1].item() / metric[0].item()
+
+
+def free_port(host='', low=20000, high=40000):
+    """
+    Return a port number that is most likely free.
+    This could suffer from a race condition although
+    it should be quite rare.
+    """
+    sock = socket.socket()
+    while True:
+        port = random.randint(low, high)
+        try:
+            sock.bind((host, port))
+        except OSError as error:
+            if error.errno == errno.EADDRINUSE:
+                continue
+            raise
+        return port
+
+
+def sizeof_fmt(num, suffix='B'):
+    """
+    Given `num` bytes, return human readable size.
+    Taken from https://stackoverflow.com/a/1094933
+    """
+    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
+        if abs(num) < 1024.0:
+            return "%3.1f%s%s" % (num, unit, suffix)
+        num /= 1024.0
+    return "%.1f%s%s" % (num, 'Yi', suffix)
+
+
+def human_seconds(seconds, display='.2f'):
+    """
+    Given `seconds` seconds, return human readable duration.
+    """
+    value = seconds * 1e6
+    ratios = [1e3, 1e3, 60, 60, 24]
+    names = ['us', 'ms', 's', 'min', 'hrs', 'days']
+    last = names.pop(0)
+    for name, ratio in zip(names, ratios):
+        if value / ratio < 0.3:
+            break
+        value /= ratio
+        last = name
+    return f"{format(value, display)} {last}"
+
+
+class TensorChunk:
+    def __init__(self, tensor, offset=0, length=None):
+        total_length = tensor.shape[-1]
+        assert offset >= 0
+        assert offset < total_length
+
+        if length is None:
+            length = total_length - offset
+        else:
+            length = min(total_length - offset, length)
+
+        self.tensor = tensor
+        self.offset = offset
+        self.length = length
+        self.device = tensor.device
+
+    @property
+    def shape(self):
+        shape = list(self.tensor.shape)
+        shape[-1] = self.length
+        return shape
+
+    def padded(self, target_length):
+        delta = target_length - self.length
+        total_length = self.tensor.shape[-1]
+        assert delta >= 0
+
+        start = self.offset - delta // 2
+        end = start + target_length
+
+        correct_start = max(0, start)
+        correct_end = min(total_length, end)
+
+        pad_left = correct_start - start
+        pad_right = end - correct_end
+
+        out = F.pad(self.tensor[..., correct_start:correct_end], (pad_left, pad_right))
+        assert out.shape[-1] == target_length
+        return out
+
+
+def tensor_chunk(tensor_or_chunk):
+    if isinstance(tensor_or_chunk, TensorChunk):
+        return tensor_or_chunk
+    else:
+        assert isinstance(tensor_or_chunk, th.Tensor)
+        return TensorChunk(tensor_or_chunk)
+
+
+def apply_model_v1(model, mix, shifts=None, split=False, progress=False, set_progress_bar=None):
+    """
+    Apply model to a given mixture.
+
+    Args:
+        shifts (int): if > 0, will shift in time `mix` by a random amount between 0 and 0.5 sec
+            and apply the oppositve shift to the output. This is repeated `shifts` time and
+            all predictions are averaged. This effectively makes the model time equivariant
+            and improves SDR by up to 0.2 points.
+        split (bool): if True, the input will be broken down in 8 seconds extracts
+            and predictions will be performed individually on each and concatenated.
+            Useful for model with large memory footprint like Tasnet.
+        progress (bool): if True, show a progress bar (requires split=True)
+    """
+
+    channels, length = mix.size()
+    device = mix.device
+    progress_value = 0
+    
+    if split:
+        out = th.zeros(4, channels, length, device=device)
+        shift = model.samplerate * 10
+        offsets = range(0, length, shift)
+        scale = 10
+        if progress:
+            offsets = tqdm.tqdm(offsets, unit_scale=scale, ncols=120, unit='seconds')
+        for offset in offsets:
+            chunk = mix[..., offset:offset + shift]
+            if set_progress_bar:
+                progress_value += 1
+                set_progress_bar(0.1, (0.8/len(offsets)*progress_value))
+                chunk_out = apply_model_v1(model, chunk, shifts=shifts, set_progress_bar=set_progress_bar)
+            else:
+                chunk_out = apply_model_v1(model, chunk, shifts=shifts)
+            out[..., offset:offset + shift] = chunk_out
+            offset += shift
+        return out
+    elif shifts:
+        max_shift = int(model.samplerate / 2)
+        mix = F.pad(mix, (max_shift, max_shift))
+        offsets = list(range(max_shift))
+        random.shuffle(offsets)
+        out = 0
+        for offset in offsets[:shifts]:
+            shifted = mix[..., offset:offset + length + max_shift]
+            if set_progress_bar:
+                shifted_out = apply_model_v1(model, shifted, set_progress_bar=set_progress_bar)
+            else:
+                shifted_out = apply_model_v1(model, shifted)
+            out += shifted_out[..., max_shift - offset:max_shift - offset + length]
+        out /= shifts
+        return out
+    else:
+        valid_length = model.valid_length(length)
+        delta = valid_length - length
+        padded = F.pad(mix, (delta // 2, delta - delta // 2))
+        with th.no_grad():
+            out = model(padded.unsqueeze(0))[0]
+        return center_trim(out, mix)
+
+def apply_model_v2(model, mix, shifts=None, split=False,
+                overlap=0.25, transition_power=1., progress=False, set_progress_bar=None): 
+    """
+    Apply model to a given mixture.
+
+    Args:
+        shifts (int): if > 0, will shift in time `mix` by a random amount between 0 and 0.5 sec
+            and apply the oppositve shift to the output. This is repeated `shifts` time and
+            all predictions are averaged. This effectively makes the model time equivariant
+            and improves SDR by up to 0.2 points.
+        split (bool): if True, the input will be broken down in 8 seconds extracts
+            and predictions will be performed individually on each and concatenated.
+            Useful for model with large memory footprint like Tasnet.
+        progress (bool): if True, show a progress bar (requires split=True)
+    """
+    
+    assert transition_power >= 1, "transition_power < 1 leads to weird behavior."
+    device = mix.device
+    channels, length = mix.shape
+    progress_value = 0
+    
+    if split:
+        out = th.zeros(len(model.sources), channels, length, device=device)
+        sum_weight = th.zeros(length, device=device)
+        segment = model.segment_length
+        stride = int((1 - overlap) * segment)
+        offsets = range(0, length, stride)
+        scale = stride / model.samplerate
+        if progress:
+            offsets = tqdm.tqdm(offsets, unit_scale=scale, ncols=120, unit='seconds')
+        # We start from a triangle shaped weight, with maximal weight in the middle
+        # of the segment. Then we normalize and take to the power `transition_power`.
+        # Large values of transition power will lead to sharper transitions.
+        weight = th.cat([th.arange(1, segment // 2 + 1),
+                         th.arange(segment - segment // 2, 0, -1)]).to(device)
+        assert len(weight) == segment
+        # If the overlap < 50%, this will translate to linear transition when
+        # transition_power is 1.
+        weight = (weight / weight.max())**transition_power
+        for offset in offsets:
+            chunk = TensorChunk(mix, offset, segment)
+            if set_progress_bar:
+                progress_value += 1
+                set_progress_bar(0.1, (0.8/len(offsets)*progress_value))
+                chunk_out = apply_model_v2(model, chunk, shifts=shifts, set_progress_bar=set_progress_bar)
+            else:
+                chunk_out = apply_model_v2(model, chunk, shifts=shifts)
+            chunk_length = chunk_out.shape[-1]
+            out[..., offset:offset + segment] += weight[:chunk_length] * chunk_out
+            sum_weight[offset:offset + segment] += weight[:chunk_length]
+            offset += segment
+        assert sum_weight.min() > 0
+        out /= sum_weight
+        return out
+    elif shifts:
+        max_shift = int(0.5 * model.samplerate)
+        mix = tensor_chunk(mix)
+        padded_mix = mix.padded(length + 2 * max_shift)
+        out = 0
+        for _ in range(shifts):
+            offset = random.randint(0, max_shift)
+            shifted = TensorChunk(padded_mix, offset, length + max_shift - offset)
+            
+            if set_progress_bar:
+                progress_value += 1
+                shifted_out = apply_model_v2(model, shifted, set_progress_bar=set_progress_bar)
+            else:
+                shifted_out = apply_model_v2(model, shifted)
+            out += shifted_out[..., max_shift - offset:]
+        out /= shifts
+        return out
+    else:
+        valid_length = model.valid_length(length)
+        mix = tensor_chunk(mix)
+        padded_mix = mix.padded(valid_length)
+        with th.no_grad():
+            out = model(padded_mix.unsqueeze(0))[0]
+        return center_trim(out, length)
+
+
+@contextmanager
+def temp_filenames(count, delete=True):
+    names = []
+    try:
+        for _ in range(count):
+            names.append(tempfile.NamedTemporaryFile(delete=False).name)
+        yield names
+    finally:
+        if delete:
+            for name in names:
+                os.unlink(name)
+
+
+def get_quantizer(model, args, optimizer=None):
+    quantizer = None
+    if args.diffq:
+        quantizer = DiffQuantizer(
+            model, min_size=args.q_min_size, group_size=8)
+        if optimizer is not None:
+            quantizer.setup_optimizer(optimizer)
+    elif args.qat:
+        quantizer = UniformQuantizer(
+                model, bits=args.qat, min_size=args.q_min_size)
+    return quantizer
+
+
+def load_model(path, strict=False):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        load_from = path
+        package = th.load(load_from, 'cpu')
+
+    klass = package["klass"]
+    args = package["args"]
+    kwargs = package["kwargs"]
+
+    if strict:
+        model = klass(*args, **kwargs)
+    else:
+        sig = inspect.signature(klass)
+        for key in list(kwargs):
+            if key not in sig.parameters:
+                warnings.warn("Dropping inexistant parameter " + key)
+                del kwargs[key]
+        model = klass(*args, **kwargs)
+
+    state = package["state"]
+    training_args = package["training_args"]
+    quantizer = get_quantizer(model, training_args)
+
+    set_state(model, quantizer, state)
+    return model
+
+
+def get_state(model, quantizer):
+    if quantizer is None:
+        state = {k: p.data.to('cpu') for k, p in model.state_dict().items()}
+    else:
+        state = quantizer.get_quantized_state()
+        buf = io.BytesIO()
+        th.save(state, buf)
+        state = {'compressed': zlib.compress(buf.getvalue())}
+    return state
+
+
+def set_state(model, quantizer, state):
+    if quantizer is None:
+        model.load_state_dict(state)
+    else:
+        buf = io.BytesIO(zlib.decompress(state["compressed"]))
+        state = th.load(buf, "cpu")
+        quantizer.restore_quantized_state(state)
+
+    return state
+
+
+def save_state(state, path):
+    buf = io.BytesIO()
+    th.save(state, buf)
+    sig = hashlib.sha256(buf.getvalue()).hexdigest()[:8]
+
+    path = path.parent / (path.stem + "-" + sig + path.suffix)
+    path.write_bytes(buf.getvalue())
+
+
+def save_model(model, quantizer, training_args, path):
+    args, kwargs = model._init_args_kwargs
+    klass = model.__class__
+
+    state = get_state(model, quantizer)
+
+    save_to = path
+    package = {
+        'klass': klass,
+        'args': args,
+        'kwargs': kwargs,
+        'state': state,
+        'training_args': training_args,
+    }
+    th.save(package, save_to)
+
+
+def capture_init(init):
+    @functools.wraps(init)
+    def __init__(self, *args, **kwargs):
+        self._init_args_kwargs = (args, kwargs)
+        init(self, *args, **kwargs)
+
+    return __init__
+
+class DummyPoolExecutor:
+    class DummyResult:
+        def __init__(self, func, *args, **kwargs):
+            self.func = func
+            self.args = args
+            self.kwargs = kwargs
+
+        def result(self):
+            return self.func(*self.args, **self.kwargs)
+
+    def __init__(self, workers=0):
+        pass
+
+    def submit(self, func, *args, **kwargs):
+        return DummyPoolExecutor.DummyResult(func, *args, **kwargs)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        return

lib/uvr5_pack/error_handling.py

@@ -0,0 +1,106 @@
+from datetime import datetime
+import traceback
+
+CUDA_MEMORY_ERROR = "CUDA out of memory"
+CUDA_RUNTIME_ERROR = "CUDNN error executing cudnnSetTensorNdDescriptor"
+DEMUCS_MODEL_MISSING_ERROR = "is neither a single pre-trained model or a bag of models."
+ENSEMBLE_MISSING_MODEL_ERROR = "local variable \'enseExport\' referenced before assignment"
+FFMPEG_MISSING_ERROR = """audioread\__init__.py", line 116, in audio_open"""
+FILE_MISSING_ERROR = "FileNotFoundError"
+MDX_MEMORY_ERROR = "onnxruntime::CudaCall CUDA failure 2: out of memory"
+MDX_MODEL_MISSING = "[ONNXRuntimeError] : 3 : NO_SUCHFILE"
+MDX_MODEL_SETTINGS_ERROR = "Got invalid dimensions for input"
+MDX_RUNTIME_ERROR = "onnxruntime::BFCArena::AllocateRawInternal"
+MODULE_ERROR = "ModuleNotFoundError"
+WINDOW_SIZE_ERROR = "h1_shape[3] must be greater than h2_shape[3]"
+SF_WRITE_ERROR = "sf.write"
+SYSTEM_MEMORY_ERROR = "DefaultCPUAllocator: not enough memory"
+MISSING_MODEL_ERROR = "'NoneType\' object has no attribute \'model_basename\'"
+ARRAY_SIZE_ERROR = "ValueError: \"array is too big; `arr.size * arr.dtype.itemsize` is larger than the maximum possible size.\""
+GPU_INCOMPATIBLE_ERROR = "no kernel image is available for execution on the device"
+
+CONTACT_DEV = 'If this error persists, please contact the developers with the error details.'
+
+ERROR_MAPPER = {
+    CUDA_MEMORY_ERROR:
+                        ('The application was unable to allocate enough GPU memory to use this model. ' + 
+                        'Please close any GPU intensive applications and try again.\n' + 
+                        'If the error persists, your GPU might not be supported.') ,
+    CUDA_RUNTIME_ERROR:
+                        (f'Your PC cannot process this audio file with the chunk size selected. Please lower the chunk size and try again.\n\n{CONTACT_DEV}'),
+    DEMUCS_MODEL_MISSING_ERROR:
+                        ('The selected Demucs model is missing. ' + 
+                        'Please download the model or make sure it is in the correct directory.'),
+    ENSEMBLE_MISSING_MODEL_ERROR:
+                        ('The application was unable to locate a model you selected for this ensemble.\n\n' + 
+                        'Please do the following to use all compatible models:\n\n1. Navigate to the \"Updates\" tab in the Help Guide.\n2. Download and install the model expansion pack.\n3. Then try again.\n\n' + 
+                        'If the error persists, please verify all models are present.'),
+    FFMPEG_MISSING_ERROR:
+                        ('The input file type is not supported or FFmpeg is missing. Please select a file type supported by FFmpeg and try again. ' + 
+                        'If FFmpeg is missing or not installed, you will only be able to process \".wav\" files until it is available on this system. ' + 
+                        f'See the \"More Info\" tab in the Help Guide.\n\n{CONTACT_DEV}'),
+    FILE_MISSING_ERROR:
+                        (f'Missing file error raised. Please address the error and try again.\n\n{CONTACT_DEV}'),
+    MDX_MEMORY_ERROR:
+                        ('The application was unable to allocate enough GPU memory to use this model.\n\n' + 
+                        'Please do the following:\n\n1. Close any GPU intensive applications.\n2. Lower the set chunk size.\n3. Then try again.\n\n' + 
+                        'If the error persists, your GPU might not be supported.'),
+    MDX_MODEL_MISSING:
+                        ('The application could not detect this MDX-Net model on your system. ' + 
+                        'Please make sure all the models are present in the correct directory.\n\n' + 
+                        'If the error persists, please reinstall application or contact the developers.'),
+    MDX_RUNTIME_ERROR:
+                        ('The application was unable to allocate enough GPU memory to use this model.\n\n' + 
+                        'Please do the following:\n\n1. Close any GPU intensive applications.\n2. Lower the set chunk size.\n3. Then try again.\n\n' + 
+                        'If the error persists, your GPU might not be supported.'),
+    WINDOW_SIZE_ERROR:
+                        ('Invalid window size.\n\n' + 
+                        'The chosen window size is likely not compatible with this model. Please select a different size and try again.'),
+    SF_WRITE_ERROR:
+                        ('Could not write audio file.\n\n' + 
+                        'This could be due to one of the following:\n\n1. Low storage on target device.\n2. The export directory no longer exists.\n3. A system permissions issue.'),
+    SYSTEM_MEMORY_ERROR:
+                        ('The application was unable to allocate enough system memory to use this model.\n\n' + 
+                        'Please do the following:\n\n1. Restart this application.\n2. Ensure any CPU intensive applications are closed.\n3. Then try again.\n\n' + 
+                        'Please Note: Intel Pentium and Intel Celeron processors do not work well with this application.\n\n' +
+                        'If the error persists, the system may not have enough RAM, or your CPU might not be supported.'),
+    MISSING_MODEL_ERROR:
+                        ('Model Missing: The application was unable to locate the chosen model.\n\n' + 
+                        'If the error persists, please verify any selected models are present.'),
+    GPU_INCOMPATIBLE_ERROR:
+                        ('This process is not compatible with your GPU.\n\n' + 
+                        'Please uncheck \"GPU Conversion\" and try again'),
+    ARRAY_SIZE_ERROR:
+                        ('The application was not able to process the given audiofile. Please convert the audiofile to another format and try again.'),
+}
+
+def error_text(process_method, exception):
+                 
+    traceback_text = ''.join(traceback.format_tb(exception.__traceback__))
+    message = f'{type(exception).__name__}: "{exception}"\nTraceback Error: "\n{traceback_text}"\n'
+    error_message = f'\n\nRaw Error Details:\n\n{message}\nError Time Stamp [{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}]\n'
+    process = f'Last Error Received:\n\nProcess: {process_method}\n\n'
+
+    for error_type, full_text in ERROR_MAPPER.items():
+        if error_type in message:
+            final_message = full_text
+            break
+    else:
+        final_message = (CONTACT_DEV) 
+        
+    return f"{process}{final_message}{error_message}"
+
+def error_dialouge(exception):
+    
+    error_name = f'{type(exception).__name__}'
+    traceback_text = ''.join(traceback.format_tb(exception.__traceback__))
+    message = f'{error_name}: "{exception}"\n{traceback_text}"'
+    
+    for error_type, full_text in ERROR_MAPPER.items():
+        if error_type in message:
+            final_message = full_text
+            break
+    else:
+        final_message = (f'{error_name}: {exception}\n\n{CONTACT_DEV}') 
+    
+    return final_message

lib/uvr5_pack/mdxnet.py

@@ -0,0 +1,140 @@
+from abc import ABCMeta
+
+import torch
+import torch.nn as nn
+from pytorch_lightning import LightningModule
+from .modules import TFC_TDF
+
+dim_s = 4
+
+class AbstractMDXNet(LightningModule):
+    __metaclass__ = ABCMeta
+
+    def __init__(self, target_name, lr, optimizer, dim_c, dim_f, dim_t, n_fft, hop_length, overlap):
+        super().__init__()
+        self.target_name = target_name
+        self.lr = lr
+        self.optimizer = optimizer
+        self.dim_c = dim_c
+        self.dim_f = dim_f
+        self.dim_t = dim_t
+        self.n_fft = n_fft
+        self.n_bins = n_fft // 2 + 1
+        self.hop_length = hop_length
+        self.window = nn.Parameter(torch.hann_window(window_length=self.n_fft, periodic=True), requires_grad=False)
+        self.freq_pad = nn.Parameter(torch.zeros([1, dim_c, self.n_bins - self.dim_f, self.dim_t]), requires_grad=False)
+
+    def configure_optimizers(self):
+        if self.optimizer == 'rmsprop':
+            return torch.optim.RMSprop(self.parameters(), self.lr)
+        
+        if self.optimizer == 'adamw':
+            return torch.optim.AdamW(self.parameters(), self.lr)
+
+class ConvTDFNet(AbstractMDXNet):
+    def __init__(self, target_name, lr, optimizer, dim_c, dim_f, dim_t, n_fft, hop_length,
+                 num_blocks, l, g, k, bn, bias, overlap):
+
+        super(ConvTDFNet, self).__init__(
+            target_name, lr, optimizer, dim_c, dim_f, dim_t, n_fft, hop_length, overlap)
+        self.save_hyperparameters()
+
+        self.num_blocks = num_blocks
+        self.l = l
+        self.g = g
+        self.k = k
+        self.bn = bn
+        self.bias = bias
+
+        if optimizer == 'rmsprop':
+            norm = nn.BatchNorm2d
+            
+        if optimizer == 'adamw':
+            norm = lambda input:nn.GroupNorm(2, input)
+            
+        self.n = num_blocks // 2
+        scale = (2, 2)
+
+        self.first_conv = nn.Sequential(
+            nn.Conv2d(in_channels=self.dim_c, out_channels=g, kernel_size=(1, 1)),
+            norm(g),
+            nn.ReLU(),
+        )
+
+        f = self.dim_f
+        c = g
+        self.encoding_blocks = nn.ModuleList()
+        self.ds = nn.ModuleList()
+        for i in range(self.n):
+            self.encoding_blocks.append(TFC_TDF(c, l, f, k, bn, bias=bias, norm=norm))
+            self.ds.append(
+                nn.Sequential(
+                    nn.Conv2d(in_channels=c, out_channels=c + g, kernel_size=scale, stride=scale),
+                    norm(c + g),
+                    nn.ReLU()
+                )
+            )
+            f = f // 2
+            c += g
+
+        self.bottleneck_block = TFC_TDF(c, l, f, k, bn, bias=bias, norm=norm)
+
+        self.decoding_blocks = nn.ModuleList()
+        self.us = nn.ModuleList()
+        for i in range(self.n):
+            self.us.append(
+                nn.Sequential(
+                    nn.ConvTranspose2d(in_channels=c, out_channels=c - g, kernel_size=scale, stride=scale),
+                    norm(c - g),
+                    nn.ReLU()
+                )
+            )
+            f = f * 2
+            c -= g
+
+            self.decoding_blocks.append(TFC_TDF(c, l, f, k, bn, bias=bias, norm=norm))
+
+        self.final_conv = nn.Sequential(
+            nn.Conv2d(in_channels=c, out_channels=self.dim_c, kernel_size=(1, 1)),
+        )
+
+    def forward(self, x):
+
+        x = self.first_conv(x)
+
+        x = x.transpose(-1, -2)
+
+        ds_outputs = []
+        for i in range(self.n):
+            x = self.encoding_blocks[i](x)
+            ds_outputs.append(x)
+            x = self.ds[i](x)
+
+        x = self.bottleneck_block(x)
+
+        for i in range(self.n):
+            x = self.us[i](x)
+            x *= ds_outputs[-i - 1]
+            x = self.decoding_blocks[i](x)
+
+        x = x.transpose(-1, -2)
+
+        x = self.final_conv(x)
+
+        return x
+    
+class Mixer(nn.Module):
+    def __init__(self, device, mixer_path):
+        
+        super(Mixer, self).__init__()
+        
+        self.linear = nn.Linear((dim_s+1)*2, dim_s*2, bias=False)
+        
+        self.load_state_dict(
+            torch.load(mixer_path, map_location=device)
+        )
+
+    def forward(self, x):
+        x = x.reshape(1,(dim_s+1)*2,-1).transpose(-1,-2)
+        x = self.linear(x)
+        return x.transpose(-1,-2).reshape(dim_s,2,-1)

lib/uvr5_pack/modules.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+
+
+class TFC(nn.Module):
+    def __init__(self, c, l, k, norm):
+        super(TFC, self).__init__()
+
+        self.H = nn.ModuleList()
+        for i in range(l):
+            self.H.append(
+                nn.Sequential(
+                    nn.Conv2d(in_channels=c, out_channels=c, kernel_size=k, stride=1, padding=k // 2),
+                    norm(c),
+                    nn.ReLU(),
+                )
+            )
+
+    def forward(self, x):
+        for h in self.H:
+            x = h(x)
+        return x
+
+
+class DenseTFC(nn.Module):
+    def __init__(self, c, l, k, norm):
+        super(DenseTFC, self).__init__()
+
+        self.conv = nn.ModuleList()
+        for i in range(l):
+            self.conv.append(
+                nn.Sequential(
+                    nn.Conv2d(in_channels=c, out_channels=c, kernel_size=k, stride=1, padding=k // 2),
+                    norm(c),
+                    nn.ReLU(),
+                )
+            )
+
+    def forward(self, x):
+        for layer in self.conv[:-1]:
+            x = torch.cat([layer(x), x], 1)
+        return self.conv[-1](x)
+
+
+class TFC_TDF(nn.Module):
+    def __init__(self, c, l, f, k, bn, dense=False, bias=True, norm=nn.BatchNorm2d):
+
+        super(TFC_TDF, self).__init__()
+
+        self.use_tdf = bn is not None
+
+        self.tfc = DenseTFC(c, l, k, norm) if dense else TFC(c, l, k, norm)
+
+        if self.use_tdf:
+            if bn == 0:
+                self.tdf = nn.Sequential(
+                    nn.Linear(f, f, bias=bias),
+                    norm(c),
+                    nn.ReLU()
+                )
+            else:
+                self.tdf = nn.Sequential(
+                    nn.Linear(f, f // bn, bias=bias),
+                    norm(c),
+                    nn.ReLU(),
+                    nn.Linear(f // bn, f, bias=bias),
+                    norm(c),
+                    nn.ReLU()
+                )
+
+    def forward(self, x):
+        x = self.tfc(x)
+        return x + self.tdf(x) if self.use_tdf else x
+

lib/uvr5_pack/name_params.json

@@ -0,0 +1,263 @@
+{
+    "equivalent" : [
+        {
+            "model_hash_name" : [
+                {
+                    "hash_name": "47939caf0cfe52a0e81442b85b971dfd",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100.json",
+                    "param_name": "4band_44100"
+                },
+                {
+                    "hash_name": "4e4ecb9764c50a8c414fee6e10395bbe",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_v2.json",
+                    "param_name": "4band_v2"
+                },
+                {
+                    "hash_name": "ca106edd563e034bde0bdec4bb7a4b36",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_v2.json",
+                    "param_name": "4band_v2"
+                },
+                {
+                    "hash_name": "e60a1e84803ce4efc0a6551206cc4b71",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100.json",
+                    "param_name": "4band_44100"
+                },
+                {
+                    "hash_name": "a82f14e75892e55e994376edbf0c8435",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100.json",
+                    "param_name": "4band_44100"
+                },
+                {
+                    "hash_name": "6dd9eaa6f0420af9f1d403aaafa4cc06",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_v2_sn.json",
+                    "param_name": "4band_v2_sn"
+                },
+                {
+                    "hash_name": "08611fb99bd59eaa79ad27c58d137727",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_v2_sn.json",
+                    "param_name": "4band_v2_sn"
+                },
+                {
+                    "hash_name": "5c7bbca45a187e81abbbd351606164e5",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/3band_44100_msb2.json",
+                    "param_name": "3band_44100_msb2"
+                },
+                {
+                    "hash_name": "d6b2cb685a058a091e5e7098192d3233",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/3band_44100_msb2.json",
+                    "param_name": "3band_44100_msb2"
+                },
+                {
+                    "hash_name": "c1b9f38170a7c90e96f027992eb7c62b",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100.json",
+                    "param_name": "4band_44100"
+                },
+                {
+                    "hash_name": "c3448ec923fa0edf3d03a19e633faa53",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100.json",
+                    "param_name": "4band_44100"
+                },
+                {
+                    "hash_name": "68aa2c8093d0080704b200d140f59e54",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/3band_44100.json",
+                    "param_name": "3band_44100"
+                },
+                {
+                    "hash_name": "fdc83be5b798e4bd29fe00fe6600e147",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/3band_44100_mid.json",
+                    "param_name": "3band_44100_mid.json"
+                },
+                {
+                    "hash_name": "2ce34bc92fd57f55db16b7a4def3d745",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/3band_44100_mid.json",
+                    "param_name": "3band_44100_mid.json"
+                },
+                {
+                    "hash_name": "52fdca89576f06cf4340b74a4730ee5f",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100.json",
+                    "param_name": "4band_44100.json"
+                },
+                {
+                    "hash_name": "41191165b05d38fc77f072fa9e8e8a30",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100.json",
+                    "param_name": "4band_44100.json"
+                },
+                {
+                    "hash_name": "89e83b511ad474592689e562d5b1f80e",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/2band_32000.json",
+                    "param_name": "2band_32000.json"
+                },
+                {
+                    "hash_name": "0b954da81d453b716b114d6d7c95177f",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/2band_32000.json",
+                    "param_name": "2band_32000.json"
+                }
+
+            ],
+            "v4 Models": [
+                {
+                    "hash_name": "6a00461c51c2920fd68937d4609ed6c8",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr16000_hl512.json",
+                    "param_name": "1band_sr16000_hl512"
+                },
+                {
+                    "hash_name": "0ab504864d20f1bd378fe9c81ef37140",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr32000_hl512.json",
+                    "param_name": "1band_sr32000_hl512"
+                },
+                {
+                    "hash_name": "7dd21065bf91c10f7fccb57d7d83b07f",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr32000_hl512.json",
+                    "param_name": "1band_sr32000_hl512"
+                },
+                {
+                    "hash_name": "80ab74d65e515caa3622728d2de07d23",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr32000_hl512.json",
+                    "param_name": "1band_sr32000_hl512"
+                },
+                {
+                    "hash_name": "edc115e7fc523245062200c00caa847f",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr33075_hl384.json",
+                    "param_name": "1band_sr33075_hl384"
+                },
+                {
+                    "hash_name": "28063e9f6ab5b341c5f6d3c67f2045b7",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr33075_hl384.json",
+                    "param_name": "1band_sr33075_hl384"
+                },
+                {
+                    "hash_name": "b58090534c52cbc3e9b5104bad666ef2",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl512.json",
+                    "param_name": "1band_sr44100_hl512"
+                },
+                {
+                    "hash_name": "0cdab9947f1b0928705f518f3c78ea8f",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl512.json",
+                    "param_name": "1band_sr44100_hl512"
+                },
+                {
+                    "hash_name": "ae702fed0238afb5346db8356fe25f13",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl1024.json",
+                    "param_name": "1band_sr44100_hl1024"
+                }
+            ]
+        }
+    ],
+    "User Models" : [
+        {
+            "1 Band": [
+                {
+                    "hash_name": "1band_sr16000_hl512",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr16000_hl512.json",
+                    "param_name": "1band_sr16000_hl512"
+                },
+                {
+                    "hash_name": "1band_sr32000_hl512",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr32000_hl512.json",
+                    "param_name": "1band_sr16000_hl512"
+                },
+                {
+                    "hash_name": "1band_sr33075_hl384",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr33075_hl384.json",
+                    "param_name": "1band_sr33075_hl384"
+                },
+                {
+                    "hash_name": "1band_sr44100_hl256",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl256.json",
+                    "param_name": "1band_sr44100_hl256"
+                },
+                {
+                    "hash_name": "1band_sr44100_hl512",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl512.json",
+                    "param_name": "1band_sr44100_hl512"
+                },
+                {
+                    "hash_name": "1band_sr44100_hl1024",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl1024.json",
+                    "param_name": "1band_sr44100_hl1024"
+                }
+            ],
+            "2 Band": [
+                {
+                    "hash_name": "2band_44100_lofi",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/2band_44100_lofi.json",
+                    "param_name": "2band_44100_lofi"
+                },
+                {
+                    "hash_name": "2band_32000",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/2band_32000.json",
+                    "param_name": "2band_32000"
+                },
+                {
+                    "hash_name": "2band_48000",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/2band_48000.json",
+                    "param_name": "2band_48000"
+                }
+            ],
+            "3 Band": [
+                {
+                    "hash_name": "3band_44100",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/3band_44100.json",
+                    "param_name": "3band_44100"
+                },
+                {
+                    "hash_name": "3band_44100_mid",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/3band_44100_mid.json",
+                    "param_name": "3band_44100_mid"
+                },
+                {
+                    "hash_name": "3band_44100_msb2",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/3band_44100_msb2.json",
+                    "param_name": "3band_44100_msb2"
+                }
+            ],
+            "4 Band": [
+                {
+                    "hash_name": "4band_44100",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100.json",
+                    "param_name": "4band_44100"
+                },
+                {
+                    "hash_name": "4band_44100_mid",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100_mid.json",
+                    "param_name": "4band_44100_mid"
+                },
+                {
+                    "hash_name": "4band_44100_msb",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100_msb.json",
+                    "param_name": "4band_44100_msb"
+                },
+                {
+                    "hash_name": "4band_44100_msb2",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100_msb2.json",
+                    "param_name": "4band_44100_msb2"
+                },
+                {
+                    "hash_name": "4band_44100_reverse",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100_reverse.json",
+                    "param_name": "4band_44100_reverse"
+                },
+                {
+                    "hash_name": "4band_44100_sw",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_44100_sw.json",
+                    "param_name": "4band_44100_sw"
+                },
+                {
+                    "hash_name": "4band_v2",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_v2.json",
+                    "param_name": "4band_v2"
+                },
+                {
+                    "hash_name": "4band_v2_sn",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/4band_v2_sn.json",
+                    "param_name": "4band_v2_sn"
+                },
+                {
+                    "hash_name": "tmodelparam",
+                    "model_params": "lib/uvr5_pack/vr_network/modelparams/tmodelparam.json",
+                    "param_name": "User Model Param Set"
+                }
+            ]
+        }
+    ]
+}

lib/uvr5_pack/pyrb.py

@@ -0,0 +1,92 @@
+import os
+import subprocess
+import tempfile
+import six
+import numpy as np
+import soundfile as sf
+import sys
+
+if getattr(sys, 'frozen', False):
+    BASE_PATH_RUB = sys._MEIPASS
+else:
+    BASE_PATH_RUB = os.path.dirname(os.path.abspath(__file__))
+
+__all__ = ['time_stretch', 'pitch_shift']
+
+__RUBBERBAND_UTIL = os.path.join(BASE_PATH_RUB, 'rubberband')
+
+if six.PY2:
+    DEVNULL = open(os.devnull, 'w')
+else:
+    DEVNULL = subprocess.DEVNULL
+
+def __rubberband(y, sr, **kwargs):
+
+    assert sr > 0
+
+    # Get the input and output tempfile
+    fd, infile = tempfile.mkstemp(suffix='.wav')
+    os.close(fd)
+    fd, outfile = tempfile.mkstemp(suffix='.wav')
+    os.close(fd)
+
+    # dump the audio
+    sf.write(infile, y, sr)
+
+    try:
+        # Execute rubberband
+        arguments = [__RUBBERBAND_UTIL, '-q']
+
+        for key, value in six.iteritems(kwargs):
+            arguments.append(str(key))
+            arguments.append(str(value))
+
+        arguments.extend([infile, outfile])
+
+        subprocess.check_call(arguments, stdout=DEVNULL, stderr=DEVNULL)
+
+        # Load the processed audio.
+        y_out, _ = sf.read(outfile, always_2d=True)
+
+        # make sure that output dimensions matches input
+        if y.ndim == 1:
+            y_out = np.squeeze(y_out)
+
+    except OSError as exc:
+        six.raise_from(RuntimeError('Failed to execute rubberband. '
+                                    'Please verify that rubberband-cli '
+                                    'is installed.'),
+                       exc)
+
+    finally:
+        # Remove temp files
+        os.unlink(infile)
+        os.unlink(outfile)
+
+    return y_out
+
+def time_stretch(y, sr, rate, rbargs=None):
+    if rate <= 0:
+        raise ValueError('rate must be strictly positive')
+
+    if rate == 1.0:
+        return y
+
+    if rbargs is None:
+        rbargs = dict()
+
+    rbargs.setdefault('--tempo', rate)
+
+    return __rubberband(y, sr, **rbargs)
+
+def pitch_shift(y, sr, n_steps, rbargs=None):
+
+    if n_steps == 0:
+        return y
+
+    if rbargs is None:
+        rbargs = dict()
+
+    rbargs.setdefault('--pitch', n_steps)
+
+    return __rubberband(y, sr, **rbargs)

lib/uvr5_pack/spec_utils.py

@@ -0,0 +1,817 @@
+import hashlib
+import json
+import os
+import librosa
+import numpy as np
+import soundfile as sf
+import math
+import random
+import math
+import platform
+import traceback
+from . import pyrb
+#cur
+OPERATING_SYSTEM = platform.system()
+SYSTEM_ARCH = platform.platform()
+SYSTEM_PROC = platform.processor()
+ARM = 'arm'
+
+# if OPERATING_SYSTEM == 'Windows':
+#     from pyrubberband import pyrb
+# else:
+from . import pyrb
+
+if OPERATING_SYSTEM == 'Darwin':
+    wav_resolution = "polyphase" if SYSTEM_PROC == ARM or ARM in SYSTEM_ARCH else "sinc_fastest" 
+else:
+    wav_resolution = "sinc_fastest"
+
+MAX_SPEC = 'Max Spec'
+MIN_SPEC = 'Min Spec'
+AVERAGE = 'Average'
+
+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_time = (h1_shape[3] - h2_shape[3]) // 2
+    e_time = s_time + h2_shape[3]
+    h1 = h1[:, :, :, s_time:e_time]
+
+    return h1
+
+def preprocess(X_spec):
+    X_mag = np.abs(X_spec)
+    X_phase = np.angle(X_spec)
+
+    return X_mag, X_phase
+
+def make_padding(width, cropsize, offset):
+    left = offset
+    roi_size = cropsize - offset * 2
+    if roi_size == 0:
+        roi_size = cropsize
+    right = roi_size - (width % roi_size) + left
+
+    return left, right, roi_size
+
+def wave_to_spectrogram(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False):
+    if reverse:
+        wave_left = np.flip(np.asfortranarray(wave[0]))
+        wave_right = np.flip(np.asfortranarray(wave[1]))
+    elif mid_side:
+        wave_left = np.asfortranarray(np.add(wave[0], wave[1]) / 2)
+        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1]))
+    elif mid_side_b2:
+        wave_left = np.asfortranarray(np.add(wave[1], wave[0] * .5))
+        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1] * .5))
+    else:
+        wave_left = np.asfortranarray(wave[0])
+        wave_right = np.asfortranarray(wave[1])
+
+    spec_left = librosa.stft(wave_left, n_fft=n_fft, hop_length=hop_length)
+    spec_right = librosa.stft(wave_right, n_fft=n_fft, hop_length=hop_length)
+    
+    spec = np.asfortranarray([spec_left, spec_right])
+
+    return spec
+   
+def wave_to_spectrogram_mt(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False):
+    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] * .5))
+        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1] * .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 normalize(wave, is_normalize=False):
+    """Save output music files"""
+    maxv = np.abs(wave).max()
+    if maxv > 1.0:
+        print(f"\nNormalization Set {is_normalize}: Input above threshold for clipping. Max:{maxv}")
+        if is_normalize:
+            print(f"The result was normalized.")
+            wave /= maxv
+        else:
+            print(f"The result was not normalized.")
+    else:
+        print(f"\nNormalization Set {is_normalize}: Input not above threshold for clipping. Max:{maxv}")
+    
+    return wave
+    
+def normalize_two_stem(wave, mix, is_normalize=False):
+    """Save output music files"""
+    
+    maxv = np.abs(wave).max()
+    max_mix = np.abs(mix).max()
+    
+    if maxv > 1.0:
+        print(f"\nNormalization Set {is_normalize}: Primary source above threshold for clipping. Max:{maxv}")
+        print(f"\nNormalization Set {is_normalize}: Mixture above threshold for clipping. Max:{max_mix}")
+        if is_normalize:
+            print(f"The result was normalized.")
+            wave /= maxv
+            mix /= maxv
+        else:
+            print(f"The result was not normalized.")
+    else:
+        print(f"\nNormalization Set {is_normalize}: Input not above threshold for clipping. Max:{maxv}")
+    
+    
+    print(f"\nNormalization Set {is_normalize}: Primary source - Max:{np.abs(wave).max()}")
+    print(f"\nNormalization Set {is_normalize}: Mixture - Max:{np.abs(mix).max()}")
+    
+    return wave, mix    
+
+def combine_spectrograms(specs, mp):
+    l = min([specs[i].shape[2] for i in specs])    
+    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
+                
+    return np.asfortranarray(spec_c)
+    
+def spectrogram_to_image(spec, mode='magnitude'):
+    if mode == 'magnitude':
+        if np.iscomplexobj(spec):
+            y = np.abs(spec)
+        else:
+            y = spec
+        y = np.log10(y ** 2 + 1e-8)
+    elif mode == 'phase':
+        if np.iscomplexobj(spec):
+            y = np.angle(spec)
+        else:
+            y = spec
+
+    y -= y.min()
+    y *= 255 / y.max()
+    img = np.uint8(y)
+
+    if y.ndim == 3:
+        img = img.transpose(1, 2, 0)
+        img = np.concatenate([
+            np.max(img, axis=2, keepdims=True), img
+        ], axis=2)
+
+    return img
+
+def reduce_vocal_aggressively(X, y, softmask):
+    v = X - y
+    y_mag_tmp = np.abs(y)
+    v_mag_tmp = np.abs(v)
+
+    v_mask = v_mag_tmp > y_mag_tmp
+    y_mag = np.clip(y_mag_tmp - v_mag_tmp * v_mask * softmask, 0, np.inf)
+
+    return y_mag * np.exp(1.j * np.angle(y))
+
+def merge_artifacts(y_mask, thres=0.01, min_range=64, fade_size=32):
+    mask = y_mask
+    
+    try:
+        if min_range < fade_size * 2:
+            raise ValueError('min_range must be >= fade_size * 2')
+
+        idx = np.where(y_mask.min(axis=(0, 1)) > thres)[0]
+        start_idx = np.insert(idx[np.where(np.diff(idx) != 1)[0] + 1], 0, idx[0])
+        end_idx = np.append(idx[np.where(np.diff(idx) != 1)[0]], idx[-1])
+        artifact_idx = np.where(end_idx - start_idx > min_range)[0]
+        weight = np.zeros_like(y_mask)
+        if len(artifact_idx) > 0:
+            start_idx = start_idx[artifact_idx]
+            end_idx = end_idx[artifact_idx]
+            old_e = None
+            for s, e in zip(start_idx, end_idx):
+                if old_e is not None and s - old_e < fade_size:
+                    s = old_e - fade_size * 2
+
+                if s != 0:
+                    weight[:, :, s:s + fade_size] = np.linspace(0, 1, fade_size)
+                else:
+                    s -= fade_size
+
+                if e != y_mask.shape[2]:
+                    weight[:, :, e - fade_size:e] = np.linspace(1, 0, fade_size)
+                else:
+                    e += fade_size
+
+                weight[:, :, s + fade_size:e - fade_size] = 1
+                old_e = e
+
+        v_mask = 1 - y_mask
+        y_mask += weight * v_mask
+        
+        mask = y_mask
+    except Exception as e:
+        error_name = f'{type(e).__name__}'
+        traceback_text = ''.join(traceback.format_tb(e.__traceback__))
+        message = f'{error_name}: "{e}"\n{traceback_text}"'
+        print('Post Process Failed: ', message)
+        
+
+    return mask
+
+def align_wave_head_and_tail(a, b):
+    l = min([a[0].size, b[0].size])  
+    
+    return a[:l,:l], b[:l,:l]
+    
+def spectrogram_to_wave(spec, hop_length, mid_side, mid_side_b2, reverse, clamp=False):
+    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, .4 * wave_left), np.subtract(wave_left / 1.25, .4 * wave_right)])
+    else:
+        return np.asfortranarray([wave_left, wave_right])
+    
+def spectrogram_to_wave_mt(spec, hop_length, mid_side, reverse, mid_side_b2):
+    import threading
+
+    spec_left = np.asfortranarray(spec[0])
+    spec_right = np.asfortranarray(spec[1])
+    
+    def run_thread(**kwargs):
+        global wave_left
+        wave_left = librosa.istft(**kwargs)
+        
+    thread = threading.Thread(target=run_thread, kwargs={'stft_matrix': spec_left, 'hop_length': hop_length})
+    thread.start()
+    wave_right = librosa.istft(spec_right, hop_length=hop_length)
+    thread.join()   
+    
+    if reverse:
+        return np.asfortranarray([np.flip(wave_left), np.flip(wave_right)])
+    elif mid_side:
+        return np.asfortranarray([np.add(wave_left, wave_right / 2), np.subtract(wave_left, wave_right / 2)])
+    elif mid_side_b2:
+        return np.asfortranarray([np.add(wave_right / 1.25, .4 * wave_left), np.subtract(wave_left / 1.25, .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):
+    bands_n = len(mp.param['band'])    
+    offset = 0
+
+    for d in range(1, bands_n + 1):
+        bp = mp.param['band'][d]
+        spec_s = np.ndarray(shape=(2, bp['n_fft'] // 2 + 1, spec_m.shape[2]), dtype=complex)
+        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 = np.add(wave, spectrogram_to_wave(spec_s, bp['hl'], mp.param['mid_side'], mp.param['mid_side_b2'], mp.param['reverse']))
+        else:
+            sr = mp.param['band'][d+1]['sr']
+            if d == 1: # lower
+                spec_s = fft_lp_filter(spec_s, bp['lpf_start'], bp['lpf_stop'])
+                wave = librosa.resample(spectrogram_to_wave(spec_s, bp['hl'], mp.param['mid_side'],mp.param['mid_side_b2'], mp.param['reverse']), orig_sr=bp['sr'], target_sr=sr, res_type=wav_resolution)
+            else: # mid
+                spec_s = fft_hp_filter(spec_s, bp['hpf_start'], bp['hpf_stop'] - 1)
+                spec_s = fft_lp_filter(spec_s, bp['lpf_start'], bp['lpf_stop'])
+                wave2 = np.add(wave, spectrogram_to_wave(spec_s, bp['hl'], mp.param['mid_side'], mp.param['mid_side_b2'], mp.param['reverse']))
+                wave = librosa.resample(wave2, orig_sr=bp['sr'], target_sr=sr, res_type=wav_resolution)
+        
+    return wave
+
+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 '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.j * 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)
+
+def adjust_aggr(mask, is_non_accom_stem, aggressiveness):
+    aggr = aggressiveness['value']
+
+    if aggr != 0:
+        if is_non_accom_stem:
+            aggr = 1 - aggr
+    
+        aggr = [aggr, aggr]
+    
+        if aggressiveness['aggr_correction'] is not None:
+            aggr[0] += aggressiveness['aggr_correction']['left']
+            aggr[1] += aggressiveness['aggr_correction']['right']
+
+        for ch in range(2):
+            mask[ch, :aggressiveness['split_bin']] = np.power(mask[ch, :aggressiveness['split_bin']], 1 + aggr[ch] / 3)
+            mask[ch, aggressiveness['split_bin']:] = np.power(mask[ch, aggressiveness['split_bin']:], 1 + aggr[ch])
+
+        # if is_non_accom_stem:
+        #     mask = (1.0 - mask)
+        
+    return mask
+
+def stft(wave, nfft, hl):
+    wave_left = np.asfortranarray(wave[0])
+    wave_right = np.asfortranarray(wave[1])
+    spec_left = librosa.stft(wave_left, nfft, hop_length=hl)
+    spec_right = librosa.stft(wave_right, nfft, hop_length=hl)
+    spec = np.asfortranarray([spec_left, spec_right])
+
+    return spec
+
+def istft(spec, hl):
+    spec_left = np.asfortranarray(spec[0])
+    spec_right = np.asfortranarray(spec[1])
+    wave_left = librosa.istft(spec_left, hop_length=hl)
+    wave_right = librosa.istft(spec_right, hop_length=hl)
+    wave = np.asfortranarray([wave_left, wave_right])
+
+    return wave
+
+def spec_effects(wave, algorithm='Default', value=None):
+    spec = [stft(wave[0],2048,1024), stft(wave[1],2048,1024)]
+    if algorithm == 'Min_Mag':
+        v_spec_m = np.where(np.abs(spec[1]) <= np.abs(spec[0]), spec[1], spec[0])
+        wave = istft(v_spec_m,1024)
+    elif algorithm == 'Max_Mag':
+        v_spec_m = np.where(np.abs(spec[1]) >= np.abs(spec[0]), spec[1], spec[0])
+        wave = istft(v_spec_m,1024)
+    elif algorithm == 'Default':
+        wave = (wave[1] * value) + (wave[0] * (1-value))
+    elif algorithm == 'Invert_p':
+        X_mag = np.abs(spec[0])
+        y_mag = np.abs(spec[1])            
+        max_mag = np.where(X_mag >= y_mag, X_mag, y_mag)  
+        v_spec = spec[1] - max_mag * np.exp(1.j * np.angle(spec[0]))
+        wave = istft(v_spec,1024)
+            
+    return wave      
+
+def spectrogram_to_wave_no_mp(spec, n_fft=2048, hop_length=1024):
+    wave = librosa.istft(spec, n_fft=n_fft, hop_length=hop_length)
+    
+    if wave.ndim == 1:
+        wave = np.asfortranarray([wave,wave])
+
+    return wave
+
+def wave_to_spectrogram_no_mp(wave):
+    
+    spec = librosa.stft(wave, n_fft=2048, hop_length=1024)
+    
+    if spec.ndim == 1:
+        spec = np.asfortranarray([spec,spec])
+
+    return spec
+
+def invert_audio(specs, invert_p=True):
+    
+    ln = min([specs[0].shape[2], specs[1].shape[2]])
+    specs[0] = specs[0][:,:,:ln]
+    specs[1] = specs[1][:,:,:ln]
+        
+    if invert_p:
+        X_mag = np.abs(specs[0])
+        y_mag = np.abs(specs[1])            
+        max_mag = np.where(X_mag >= y_mag, X_mag, y_mag)  
+        v_spec = specs[1] - max_mag * np.exp(1.j * np.angle(specs[0]))
+    else:
+        specs[1] = reduce_vocal_aggressively(specs[0], specs[1], 0.2)
+        v_spec = specs[0] - specs[1]
+
+    return v_spec
+
+def invert_stem(mixture, stem):
+    
+    mixture = wave_to_spectrogram_no_mp(mixture)
+    stem = wave_to_spectrogram_no_mp(stem)
+    output = spectrogram_to_wave_no_mp(invert_audio([mixture, stem]))
+
+    return -output.T
+
+def ensembling(a, specs):   
+    for i in range(1, len(specs)):
+        if i == 1:
+            spec = specs[0]
+
+        ln = min([spec.shape[2], specs[i].shape[2]])
+        spec = spec[:,:,:ln]
+        specs[i] = specs[i][:,:,:ln]
+        
+        if MIN_SPEC == a:
+            spec = np.where(np.abs(specs[i]) <= np.abs(spec), specs[i], spec)
+        if MAX_SPEC == a:
+            spec = np.where(np.abs(specs[i]) >= np.abs(spec), specs[i], spec)  
+        if AVERAGE == a:
+            spec = np.where(np.abs(specs[i]) == np.abs(spec), specs[i], spec)  
+
+    return spec
+
+def ensemble_inputs(audio_input, algorithm, is_normalization, wav_type_set, save_path):
+
+    wavs_ = []
+    
+    if algorithm == AVERAGE:
+        output = average_audio(audio_input)
+        samplerate = 44100
+    else:
+        specs = []
+        
+        for i in range(len(audio_input)):  
+            wave, samplerate = librosa.load(audio_input[i], mono=False, sr=44100)
+            wavs_.append(wave)
+            spec = wave_to_spectrogram_no_mp(wave)
+            specs.append(spec)
+        
+        wave_shapes = [w.shape[1] for w in wavs_]
+        target_shape = wavs_[wave_shapes.index(max(wave_shapes))]
+        
+        output = spectrogram_to_wave_no_mp(ensembling(algorithm, specs))
+        output = to_shape(output, target_shape.shape)
+
+    sf.write(save_path, normalize(output.T, is_normalization), samplerate, subtype=wav_type_set)
+
+def to_shape(x, target_shape):
+    padding_list = []
+    for x_dim, target_dim in zip(x.shape, target_shape):
+        pad_value = (target_dim - x_dim)
+        pad_tuple = ((0, pad_value))
+        padding_list.append(pad_tuple)
+    
+    return np.pad(x, tuple(padding_list), mode='constant')
+
+def to_shape_minimize(x: np.ndarray, target_shape):
+    
+    padding_list = []
+    for x_dim, target_dim in zip(x.shape, target_shape):
+        pad_value = (target_dim - x_dim)
+        pad_tuple = ((0, pad_value))
+        padding_list.append(pad_tuple)
+    
+    return np.pad(x, tuple(padding_list), mode='constant')
+
+def augment_audio(export_path, audio_file, rate, is_normalization, wav_type_set, save_format=None, is_pitch=False):
+
+    wav, sr = librosa.load(audio_file, sr=44100, mono=False)
+
+    if wav.ndim == 1:
+        wav = np.asfortranarray([wav,wav])
+
+    if is_pitch:
+        wav_1 = pyrb.pitch_shift(wav[0], sr, rate, rbargs=None)
+        wav_2 = pyrb.pitch_shift(wav[1], sr, rate, rbargs=None)
+    else:
+        wav_1 = pyrb.time_stretch(wav[0], sr, rate, rbargs=None)
+        wav_2 = pyrb.time_stretch(wav[1], sr, rate, rbargs=None)
+
+    if wav_1.shape > wav_2.shape:
+        wav_2 = to_shape(wav_2, wav_1.shape)
+    if wav_1.shape < wav_2.shape:
+        wav_1 = to_shape(wav_1, wav_2.shape)
+        
+    wav_mix = np.asfortranarray([wav_1, wav_2])
+    
+    sf.write(export_path, normalize(wav_mix.T, is_normalization), sr, subtype=wav_type_set)
+    save_format(export_path)
+    
+def average_audio(audio):
+    
+    waves = []
+    wave_shapes = []
+    final_waves = []
+
+    for i in range(len(audio)):
+        wave = librosa.load(audio[i], sr=44100, mono=False)
+        waves.append(wave[0])
+        wave_shapes.append(wave[0].shape[1])
+
+    wave_shapes_index = wave_shapes.index(max(wave_shapes))
+    target_shape = waves[wave_shapes_index]
+    waves.pop(wave_shapes_index)
+    final_waves.append(target_shape)
+
+    for n_array in waves:
+        wav_target = to_shape(n_array, target_shape.shape)
+        final_waves.append(wav_target)
+
+    waves = sum(final_waves)
+    waves = waves/len(audio)
+
+    return waves
+    
+def average_dual_sources(wav_1, wav_2, value):
+    
+    if wav_1.shape > wav_2.shape:
+        wav_2 = to_shape(wav_2, wav_1.shape)
+    if wav_1.shape < wav_2.shape:
+        wav_1 = to_shape(wav_1, wav_2.shape)
+
+    wave = (wav_1 * value) + (wav_2 * (1-value))
+
+    return wave
+    
+def reshape_sources(wav_1: np.ndarray, wav_2: np.ndarray):
+    
+    if wav_1.shape > wav_2.shape:
+        wav_2 = to_shape(wav_2, wav_1.shape)
+    if wav_1.shape < wav_2.shape:
+        ln = min([wav_1.shape[1], wav_2.shape[1]])
+        wav_2 = wav_2[:,:ln]
+
+    ln = min([wav_1.shape[1], wav_2.shape[1]])
+    wav_1 = wav_1[:,:ln]
+    wav_2 = wav_2[:,:ln]
+
+    return wav_2
+    
+def align_audio(file1, file2, file2_aligned, file_subtracted, wav_type_set, is_normalization, command_Text, progress_bar_main_var, save_format):
+    def get_diff(a, b):
+        corr = np.correlate(a, b, "full")
+        diff = corr.argmax() - (b.shape[0] - 1)
+        return diff
+  
+    progress_bar_main_var.set(10)
+    
+    # read tracks
+    wav1, sr1 = librosa.load(file1, sr=44100, mono=False)
+    wav2, sr2 = librosa.load(file2, sr=44100, mono=False)
+    wav1 = wav1.transpose()
+    wav2 = wav2.transpose()
+
+    command_Text(f"Audio file shapes: {wav1.shape} / {wav2.shape}\n")
+    
+    wav2_org = wav2.copy()
+    progress_bar_main_var.set(20)
+    
+    command_Text("Processing files... \n")
+    
+  # pick random position and get diff
+    
+    counts = {}       # counting up for each diff value
+    progress = 20
+    
+    check_range = 64
+
+    base = (64 / check_range)
+
+    for i in range(check_range):
+        index = int(random.uniform(44100 * 2, min(wav1.shape[0], wav2.shape[0]) - 44100 * 2))
+        shift = int(random.uniform(-22050,+22050))
+        samp1 = wav1[index      :index      +44100, 0]          # currently use left channel
+        samp2 = wav2[index+shift:index+shift+44100, 0]
+        progress += 1 * base
+        progress_bar_main_var.set(progress)
+        diff = get_diff(samp1, samp2)
+        diff -= shift
+        
+    if abs(diff) < 22050:
+        if not diff in counts:
+            counts[diff] = 0
+        counts[diff] += 1
+  
+  # use max counted diff value
+    max_count = 0
+    est_diff  = 0
+    for diff in counts.keys():
+        if counts[diff] > max_count:
+            max_count = counts[diff]
+            est_diff = diff
+    
+    command_Text(f"Estimated difference is {est_diff} (count: {max_count})\n")
+
+    progress_bar_main_var.set(90)
+    
+    audio_files = []
+
+    def save_aligned_audio(wav2_aligned):
+        command_Text(f"Aligned File 2 with File 1.\n")
+        command_Text(f"Saving files... ")
+        sf.write(file2_aligned, normalize(wav2_aligned, is_normalization), sr2, subtype=wav_type_set)
+        save_format(file2_aligned)
+        min_len = min(wav1.shape[0], wav2_aligned.shape[0])
+        wav_sub = wav1[:min_len] - wav2_aligned[:min_len]
+        audio_files.append(file2_aligned)
+        return min_len, wav_sub
+    
+  # make aligned track 2
+    if est_diff > 0:
+        wav2_aligned = np.append(np.zeros((est_diff, 2)), wav2_org, axis=0)
+        min_len, wav_sub = save_aligned_audio(wav2_aligned)
+    elif est_diff < 0:
+        wav2_aligned = wav2_org[-est_diff:]
+        min_len, wav_sub = save_aligned_audio(wav2_aligned)
+    else:
+        command_Text(f"Audio files already aligned.\n")
+        command_Text(f"Saving inverted track... ")
+        min_len = min(wav1.shape[0], wav2.shape[0])
+        wav_sub = wav1[:min_len] - wav2[:min_len]
+
+    wav_sub = np.clip(wav_sub, -1, +1)
+  
+    sf.write(file_subtracted, normalize(wav_sub, is_normalization), sr1, subtype=wav_type_set)
+    save_format(file_subtracted)
+  
+    progress_bar_main_var.set(95)
+
+# from RVC
+def cache_or_load(mix_path, inst_path, mp):
+    mix_basename = os.path.splitext(os.path.basename(mix_path))[0]
+    inst_basename = os.path.splitext(os.path.basename(inst_path))[0]
+
+    cache_dir = "mph{}".format(
+        hashlib.sha1(json.dumps(mp.param, sort_keys=True).encode("utf-8")).hexdigest()
+    )
+    mix_cache_dir = os.path.join("cache", cache_dir)
+    inst_cache_dir = os.path.join("cache", cache_dir)
+
+    os.makedirs(mix_cache_dir, exist_ok=True)
+    os.makedirs(inst_cache_dir, exist_ok=True)
+
+    mix_cache_path = os.path.join(mix_cache_dir, mix_basename + ".npy")
+    inst_cache_path = os.path.join(inst_cache_dir, inst_basename + ".npy")
+
+    if os.path.exists(mix_cache_path) and os.path.exists(inst_cache_path):
+        X_spec_m = np.load(mix_cache_path)
+        y_spec_m = np.load(inst_cache_path)
+    else:
+        X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}
+
+        for d in range(len(mp.param["band"]), 0, -1):
+            bp = mp.param["band"][d]
+
+            if d == len(mp.param["band"]):  # high-end band
+                X_wave[d], _ = librosa.load(
+                    mix_path, sr=bp["sr"], mono=False, dtype=np.float32, res_type=bp["res_type"]
+                )
+                y_wave[d], _ = librosa.load(
+                    inst_path,
+                    sr=bp["sr"],
+                    mono=False,
+                    dtype=np.float32,
+                    res_type=bp["res_type"],
+                )
+            else:  # lower bands
+                X_wave[d] = librosa.resample(
+                    X_wave[d + 1],
+                    orig_sr=mp.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=bp["res_type"],
+                )
+                y_wave[d] = librosa.resample(
+                    y_wave[d + 1],
+                    orig_sr=mp.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=bp["res_type"],
+                )
+
+            X_wave[d], y_wave[d] = align_wave_head_and_tail(X_wave[d], y_wave[d])
+
+            X_spec_s[d] = wave_to_spectrogram(
+                X_wave[d],
+                bp["hl"],
+                bp["n_fft"],
+                mp.param["mid_side"],
+                mp.param["mid_side_b2"],
+                mp.param["reverse"],
+            )
+            y_spec_s[d] = wave_to_spectrogram(
+                y_wave[d],
+                bp["hl"],
+                bp["n_fft"],
+                mp.param["mid_side"],
+                mp.param["mid_side_b2"],
+                mp.param["reverse"],
+            )
+
+        del X_wave, y_wave
+
+        X_spec_m = combine_spectrograms(X_spec_s, mp)
+        y_spec_m = combine_spectrograms(y_spec_s, mp)
+
+        if X_spec_m.shape != y_spec_m.shape:
+            raise ValueError("The combined spectrograms are different: " + mix_path)
+
+        _, ext = os.path.splitext(mix_path)
+
+        np.save(mix_cache_path, X_spec_m)
+        np.save(inst_cache_path, y_spec_m)
+
+    return X_spec_m, y_spec_m
+
+def 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")
+
+    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

lib/uvr5_pack/vr_network/dataset.py

@@ -0,0 +1,183 @@
+import os
+import random
+
+import numpy as np
+import torch
+import torch.utils.data
+from tqdm import tqdm
+
+from .. import spec_utils
+
+
+class VocalRemoverValidationSet(torch.utils.data.Dataset):
+    def __init__(self, patch_list):
+        self.patch_list = patch_list
+
+    def __len__(self):
+        return len(self.patch_list)
+
+    def __getitem__(self, idx):
+        path = self.patch_list[idx]
+        data = np.load(path)
+
+        X, y = data["X"], data["y"]
+
+        X_mag = np.abs(X)
+        y_mag = np.abs(y)
+
+        return X_mag, y_mag
+
+
+def make_pair(mix_dir, inst_dir):
+    input_exts = [".wav", ".m4a", ".mp3", ".mp4", ".flac"]
+
+    X_list = sorted(
+        [
+            os.path.join(mix_dir, fname)
+            for fname in os.listdir(mix_dir)
+            if os.path.splitext(fname)[1] in input_exts
+        ]
+    )
+    y_list = sorted(
+        [
+            os.path.join(inst_dir, fname)
+            for fname in os.listdir(inst_dir)
+            if os.path.splitext(fname)[1] in input_exts
+        ]
+    )
+
+    filelist = list(zip(X_list, y_list))
+
+    return filelist
+
+
+def train_val_split(dataset_dir, split_mode, val_rate, val_filelist):
+    if split_mode == "random":
+        filelist = make_pair(
+            os.path.join(dataset_dir, "mixtures"),
+            os.path.join(dataset_dir, "instruments"),
+        )
+
+        random.shuffle(filelist)
+
+        if len(val_filelist) == 0:
+            val_size = int(len(filelist) * val_rate)
+            train_filelist = filelist[:-val_size]
+            val_filelist = filelist[-val_size:]
+        else:
+            train_filelist = [
+                pair for pair in filelist if list(pair) not in val_filelist
+            ]
+    elif split_mode == "subdirs":
+        if len(val_filelist) != 0:
+            raise ValueError(
+                "The `val_filelist` option is not available in `subdirs` mode"
+            )
+
+        train_filelist = make_pair(
+            os.path.join(dataset_dir, "training/mixtures"),
+            os.path.join(dataset_dir, "training/instruments"),
+        )
+
+        val_filelist = make_pair(
+            os.path.join(dataset_dir, "validation/mixtures"),
+            os.path.join(dataset_dir, "validation/instruments"),
+        )
+
+    return train_filelist, val_filelist
+
+
+def augment(X, y, reduction_rate, reduction_mask, mixup_rate, mixup_alpha):
+    perm = np.random.permutation(len(X))
+    for i, idx in enumerate(tqdm(perm)):
+        if np.random.uniform() < reduction_rate:
+            y[idx] = spec_utils.reduce_vocal_aggressively(
+                X[idx], y[idx], reduction_mask
+            )
+
+        if np.random.uniform() < 0.5:
+            # swap channel
+            X[idx] = X[idx, ::-1]
+            y[idx] = y[idx, ::-1]
+        if np.random.uniform() < 0.02:
+            # mono
+            X[idx] = X[idx].mean(axis=0, keepdims=True)
+            y[idx] = y[idx].mean(axis=0, keepdims=True)
+        if np.random.uniform() < 0.02:
+            # inst
+            X[idx] = y[idx]
+
+        if np.random.uniform() < mixup_rate and i < len(perm) - 1:
+            lam = np.random.beta(mixup_alpha, mixup_alpha)
+            X[idx] = lam * X[idx] + (1 - lam) * X[perm[i + 1]]
+            y[idx] = lam * y[idx] + (1 - lam) * y[perm[i + 1]]
+
+    return X, y
+
+
+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 make_training_set(filelist, cropsize, patches, sr, hop_length, n_fft, offset):
+    len_dataset = patches * len(filelist)
+
+    X_dataset = np.zeros((len_dataset, 2, n_fft // 2 + 1, cropsize), dtype=np.complex64)
+    y_dataset = np.zeros((len_dataset, 2, n_fft // 2 + 1, cropsize), dtype=np.complex64)
+
+    for i, (X_path, y_path) in enumerate(tqdm(filelist)):
+        X, y = spec_utils.cache_or_load(X_path, y_path, sr, hop_length, n_fft)
+        coef = np.max([np.abs(X).max(), np.abs(y).max()])
+        X, y = X / coef, y / coef
+
+        l, r, roi_size = make_padding(X.shape[2], cropsize, offset)
+        X_pad = np.pad(X, ((0, 0), (0, 0), (l, r)), mode="constant")
+        y_pad = np.pad(y, ((0, 0), (0, 0), (l, r)), mode="constant")
+
+        starts = np.random.randint(0, X_pad.shape[2] - cropsize, patches)
+        ends = starts + cropsize
+        for j in range(patches):
+            idx = i * patches + j
+            X_dataset[idx] = X_pad[:, :, starts[j] : ends[j]]
+            y_dataset[idx] = y_pad[:, :, starts[j] : ends[j]]
+
+    return X_dataset, y_dataset
+
+
+def make_validation_set(filelist, cropsize, sr, hop_length, n_fft, offset):
+    patch_list = []
+    patch_dir = "cs{}_sr{}_hl{}_nf{}_of{}".format(
+        cropsize, sr, hop_length, n_fft, offset
+    )
+    os.makedirs(patch_dir, exist_ok=True)
+
+    for i, (X_path, y_path) in enumerate(tqdm(filelist)):
+        basename = os.path.splitext(os.path.basename(X_path))[0]
+
+        X, y = spec_utils.cache_or_load(X_path, y_path, sr, hop_length, n_fft)
+        coef = np.max([np.abs(X).max(), np.abs(y).max()])
+        X, y = X / coef, y / coef
+
+        l, r, roi_size = make_padding(X.shape[2], cropsize, offset)
+        X_pad = np.pad(X, ((0, 0), (0, 0), (l, r)), mode="constant")
+        y_pad = np.pad(y, ((0, 0), (0, 0), (l, r)), mode="constant")
+
+        len_dataset = int(np.ceil(X.shape[2] / roi_size))
+        for j in range(len_dataset):
+            outpath = os.path.join(patch_dir, "{}_p{}.npz".format(basename, j))
+            start = j * roi_size
+            if not os.path.exists(outpath):
+                np.savez(
+                    outpath,
+                    X=X_pad[:, :, start : start + cropsize],
+                    y=y_pad[:, :, start : start + cropsize],
+                )
+            patch_list.append(outpath)
+
+    return VocalRemoverValidationSet(patch_list)

lib/uvr5_pack/vr_network/layers.py

@@ -0,0 +1,118 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+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

lib/uvr5_pack/vr_network/layers_new.py

@@ -0,0 +1,125 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+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

lib/uvr5_pack/vr_network/model_param_init.py

@@ -0,0 +1,59 @@
+import json
+import pathlib
+
+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:
+            with open(config_path, 'r') as f:
+                self.param = json.loads(f.read(), 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 not k in self.param:
+                self.param[k] = False

lib/uvr5_pack/vr_network/modelparams/1band_sr16000_hl512.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 16000,
+			"hl": 512,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 16000,
+	"pre_filter_start": 1023,
+	"pre_filter_stop": 1024
+}

lib/uvr5_pack/vr_network/modelparams/1band_sr32000_hl512.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 32000,
+			"hl": 512,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 32000,
+	"pre_filter_start": 1000,
+	"pre_filter_stop": 1021
+}

lib/uvr5_pack/vr_network/modelparams/1band_sr33075_hl384.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 33075,
+			"hl": 384,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 33075,
+	"pre_filter_start": 1000,
+	"pre_filter_stop": 1021
+}

lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl1024.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 44100,
+			"hl": 1024,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 1023,
+	"pre_filter_stop": 1024
+}

lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl256.json

@@ -0,0 +1,19 @@
+{
+	"bins": 256,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 44100,
+			"hl": 256,
+			"n_fft": 512,
+			"crop_start": 0,
+			"crop_stop": 256,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 256,
+	"pre_filter_stop": 256
+}

lib/uvr5_pack/vr_network/modelparams/1band_sr44100_hl512.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 1023,
+	"pre_filter_stop": 1024
+}